Zipformer: A faster and better encoder for automatic speech recognition

I just want to leave a short comment on this:

these improvements set a new state of the art for supervised Librispeech

This is only partly true.

First of all, the standard Librispeech task is usually together with a language model. Together with a language model, as far as I know, the best result is 3.65% WER on test-other with the E-Branchformer (https://arxiv.org/abs/2210.00077).

Ok, but it's also legitimate to specifically look at results without external language model, as the authors do here. But in that case, they also don't quite reach state-of-the-art on test-other, which is still by the original Conformer-L with 4.3%, better than 4.38% which they reach here.

On test-clean, without language model, yes, I think they reached state-of-the-art. But this is maybe the not-so-interesting-case.

Nonetheless, the results they present are still very good, and specifically their relative improvements.

But more critically speaking, you could also say, their Conformer baseline with 5.55% WER on test-other is a bit weak baseline compared to other standard Conformer results, e.g. the 4.3% by Gulati et al., 2020, or 4.74% in https://arxiv.org/abs/2210.00077.

Thank you for your detailed review, below we address the key concerns.

In my option, this paper has the same problem that it identifies in the original Conformer write-up: It's closed-source, complicated, and it will likely not be possible for the results to be reproduced. This is not by itself disqualifying; it is the nature of the field that SOTA systems are often published without source code. However, I expect that future work iterating on this design will struggle with direct comparison, much in the same way that this paper compares to a weaker reproduction of the original Conformer.

Our code is open-sourced. The link will be put in the final version.

The comparison to the original conformer is not quite apples-to-apples, in that Zipformer-L is substantially larger than Conformer-L in both the original publication and in this paper's reproduction. I think the authors should consider including a variant of Zipformer with ~120M parameters, as originally used by Gulati et al. This would prevent the suspicion that the stronger WER simply reflects the fact that the Zipformer is larger. Still, I think Zipformer's substantially reduced computational needs demonstrate that superiority over Conformer does not only reflect model size. So, I'll recommend this paper for acceptance despite the flaw in this comparison.

We thank the reviewer for the observation. We would like to suggest that the majority of the experiments are all conducted using M-scaled Zipformer model, hence the difference of # params on L-scaled model between Zipformer and Conformer does not invalidate the conclusions we claimed in the manuscript. When it comes to Zipformer-L, we would like to slightly extend the boundary of where it could reach to given its efficiency in resource usage and computational time.

Thank you for your detailed review, below we address the key concerns.

However, it is unclear to me whether this is because of the model structure or because of the modified optimizer? .

As in Table 5, the new model structure, BiasNorm, Swoosh functions, and ScaledAdam all contribute to the performance. The model structure is mainly for higher efficiency. Swoosh functions and ScaledAdam have greater impact on performance improvement.

Will the other variants of transformer get better WER with the author proposed optimizer ?

We validated ScaledAdam on LibriSpeech using a 12-layer Conformer (attention dim=512, feed-forward dim=2048). Results indicate that ScaledAdam leads to faster convergence and better performance on Conformer.

• Adam (epoch index, WER(%) on test-clean/test-other):

epoch-10, 3.13/7.62; epoch-20, 2.83/6.88; epoch-30, 2.76/6.67;

epoch-40, 2.79/6.62; epoch-50, 2.77/6.48; epoch-60, 2.74/6.43.

• ScaledAdam:

epoch-10, 3.15/7.35; epoch-20, 2.68/6.32; epoch-30, 2.55/6.12;

epoch-40, 2.53/5.96; epoch-50, 2.55/6.02; epoch-60, 2.55/6.0.

I am wondering whether the proposed optimizer ... does it only work for speech tasks or does it only work for zipformer?

We validated ScaledAdam and Swoosh on LM task. Results of LMs with 18 transformer layers using GeLU (#param. 98M) trained on the Libriheavy transcripts (https://arxiv.org/pdf/2309.08105.pdf) suggest they can improve the performance:

Exp1: Adam w. CosineAnnealingLR and lr=4e-4;

Exp2: ScaledAdam w. Eden and $\alpha_{\text{base}}=0.025$;

Exp3: ScaledAdam w. Eden and $\alpha_{\text{base}}=0.025$, replace GeLU with SwooshL.

Trained for 12000 iter with 1B tokens, the validation loss values (negative log-likelihood):

Exp1: 0.947; Exp2: 0.9369; Exp3: 0.9229

Trained for 24000 iter with 2B tokens, the validation loss values (negative log-likelihood):

Exp1: 0.9068; Exp2: 0.885; Exp3: 0.8725

Given the zipformer-L is a pretty small model ... when the model and the training data scales up?

As in Table 4, experiments on the 10000-hr WenetSpeech manifest that Zipformer-M/L outperform all other models on Test Net and Test Meeting sets.

A 288M Zipformer on the 10000-hr GigaSpeech gets WERs of 10.07/10.19 on dev/test sets, while the 66M Zipformer-M gets 10.25/10.38. To the best of our knowledge, this is so far the lowest WER reported on GigaSpeech as on https://github.com/SpeechColab/GigaSpeech.

The authors propose changes to many widely used components, ..., I am intrigued to learn how generalizable this claim can be.

We would like to emphasize the components are focused on ASR. Our experiments on LM task suggest that these can be generalized to other tasks to a certain extent, e.g., ScaledAdam and Swoosh. Please refer to the response of question 3.

Empirically, the mean subtraction is used by systems with LayerNorm as a way to "defeat LayerNorm", by making certain elements very large and nearly constant. This is not necessary with BiasNorm because we can "defeat normalization" by having a large learnable bias. We normalize the vector by RMS instead of Var. We found an extra sqrt in the denominator in Eq.2 and removed that in the revised version.

Some of the proposed modification seem arbitrary. It would be great if the authors can explain a bit. For example, in Eq (4), for the swooshR and swooshL functions, how the -0.08 is selected ? The authors mentioned that the coefficient 0.035 in Eq (4) is slightly tuned. How general is this parameter?

The -0.08 is chosen to make the negative slope to be 10% of the positive slope. The coefficient 0.035 was tuned on LibriSpeech and maintained the same on all other experiments including the LM ones, suggesting its capability of generalization (refer to the response of question 3).

How sensitive of the WERs on librispeech relative to other factors like checkpoint choosing, different training run with different seed?

In our experiments, we search for different checkpoints for decoding and report the best result. In training, we set a fixed random seed for the libraries, such as random, numpy, and torch.

Missing references

We have added these references in Section 2.

About the referred paper Accelerating rnn-t training and inference using ctc guidance

We cited it in Section 2 as a related work. It leverages the guidance from the auxiliary CTC head to skip frames on the encoder output and even in the middle of the encoder. However, our work focus on improving the sequence encoder structure itself.

In Table 5 and section 4.0.3, it is unclear to me ... for different temporal resolution.

As described in Section 3.1 and Section 4.0.1 (in Table 1), stacks are of various embedding dims and feed-forward dims. For the ablation study ``no temporal downsampling", we use Conv-Embed with downsampling rate of 4, and 12 Zipformer blocks with the same dimensions as in the middle stack (embedding dim=512, feed-forward dim=1536).

Thank you for your detailed review, below we address the key concerns. Spelling issues are fixed in the revised version, link to the code will be attached in the final version.

About the proposed BiasNorm, Swoosh, ScaledAdam on language modeling task

We validated these components on LM task. Results indicate that using ScaleAdam and Swoosh function achieves better performance. A 18-layer transformer LM was built w. GeLU (# param. is 98M). The model was trained on the Libriheavy transcripts (https://arxiv.org/pdf/2309.08105.pdf).

Exp1: Adam w. CosineAnnealingLR and lr=4e-4;

Exp2: ScaledAdam w. Eden and $\alpha_{\text{base}}=0.025$;

Exp3: ScaledAdam w. Eden and $\alpha_{\text{base}}=0.025$, replace GeLU with SwooshL.

Trained for 12000 iter with 1B tokens, the validation loss values (negative log-likelihood):

Exp1: 0.947; Exp2: 0.9369; Exp3: 0.9229

Trained for 24000 iter with 2B tokens, the validation loss values (negative log-likelihood):

Exp1: 0.9068; Exp2: 0.885; Exp3: 0.8725

However, BiasNorm causes slightly worse result. It is possible that the BiasNorm might not work for LM task, since our modification are done based on the observation on ASR experiments. The ablation study in Tab. 5 manifests that it is effective on Zipformer models.

For comparisons, it would have been better to use a more standard transducer, ... would have made it easier to compare the results to other results from the literature (Table 2).

We updated results of Zipformer models with CTC and CTC/AED systems in Appx. Sec. A.4.2 in the revised version. Overall, Zipformer models still outperform other ASR encoders.

The actual measured speed on some given hardware.

We compared the speed and memory usage between different ASR encoders in inference mode on a 32G V100 in Figure 3. In overall, Zipformer achieves a better trade-off between performance and efficiency than other encoders.

No systematic comparison to other Conformer variants

Besides the performance comparison in Table 2, Figure 3 compares their efficiency in terms of speed and memory usage, and Sec. 2 describes their architectures and main difference between Zipformer and these variants.

How long does the training take? What batch sizes are used?

Please refer to Appx. Sec. A.4.1 of the revised version.

"Downsample module ... What happens if you just take the average? Or maybe use max-pooling instead?

With a factor of 2, the Downsample just performs the weighted average on every 2 frames with 2 learnable scalars that sum to one. This slightly outperforms average-pooling and max-pooling, on LibriSpeech test-clean/test-other:

• ours, 2.21/4.79

• average, 2.18/4.96

• max-pooling, 2.24/4.87

If you downsample at the very end to 25 Hz, why not do that earlier ... So how much worse does this get?

This degrades the performance of Zipformer-M from 2.21/4.79 to 2.45/5.43 on LibriSpeech test-clean/test-other. This level of downsampling at the middle stack might be too aggressive to cause information loss.

LayerNorm/BiasNorm: When there is a large value, ..., I wonder if there are numerical problems.

It wouldn't be so large as to cause numerical problems. It just needs large enough to dominate the denominator so that the some length information could be retained, e.g., 50 in LayerNorm.

What is the influence of exp(γ) vs just γ (without exp) in BiasNorm?

Learning the scale in log space is safer in theory, since it avoids the gradient oscillation problem that we describe in Sec. 3.3.

How does ScaledAdam compare to methods like WeightNorm, where the weights are also normalized? For ScaledAdam, as I understand, adding the learnable parameter scale would actually change the model, as this scale now is also part of the model? Is this added for every parameter, or only some selected ones, e.g. the linear transformations?

WeightNorm decouples the magnitude of a parameter from its direction, resulting in two parameters. ScaledAdam does not learn an extra parameter for the scale. It learns the scale by adding an update term (in Eq. 7), which amounts to adding an extra gradient term in the direction of rescaling each parameter. It eliminates the need to specify the modules to apply the weight normalization.

We mentioned this in Sec. 3.5 of the revised version.

I'm not sure that this would make the code simpler, to have this special case? But I also wonder, what is the influence of this?

This makes the code simpler as we don't need to compute $\boldsymbol\theta_{t-1}'$, which makes almost no impact since Adam's update magnitudes are invariant to gradient rescaling.

What is the influence of the Eden learning rate scheduling?

Eden is specially designed for ScaledAdam. Since ScaledAdam is more stable than Adam, we don't need a long warm-up period like in Conformer's learning rate schedule. We also use larger learning rate values, because we scale the update by parameter RMS (in Eq. 6), which is normally much smaller than one.

Thank you for your detailed review, below we address the key concerns.

While the motivation for biasnorm and scaledadam are well explained, the motivation for zipformer blcok, especially those downsampling and upsampling modules are not well presented.

Zipformer uses a downsampled encoder structure which processes the sequence at various frame rates. It allows to efficiently learn representation at different temporal resolutions. According to the comment, we have made this motivation more clear in the revised version.

As described in Section 3.2, Zipformer block has about twice the depth of the Conformer block, reusing the attention weights in different module to save computation and memory.

The results on aishell1 is not quite convincing compared to other conformer variants. The author could elaborate more on the performance. Could be that this is a small dataset?

Yes, Aishell-1 is a small dataset containing only 170 hours of speech, which is described in Section 4.0.1. Note that as presented in Table 3, Zipformer-S achieves better performance than Conformer implemented in ESPnet toolkit with fewer parameters, and scaling up Zipformer achieves better performance.

Curious about whether the authors tried the new activation functions, biasnorm and scaledadam on other tasks? Do they still show superiority?

We tried Swoosh function, BiasNorm, and ScaledAdam on language modeling task. Experimental results indicate that ScaledAdam and Swoosh function are able to improve the performance, while BiasNorm can't. Specifically, we build an 98M language model with 18 transformer layers with GeLU activation. The model is trained on the transcripts of Libriheavy dataset (https://arxiv.org/pdf/2309.08105.pdf). The transcripts are converted into 500-class BPE tokens. We conduct following experiments:

Exp1: use Adam optimizer with CosineAnnealingLR schedule and lr=4e-4;

Exp2: use ScaledAdam optimizer with Eden schedule and $\alpha_{\text{base}}=0.025$;

Exp3: use ScaledAdam optimizer with Eden schedule and $\alpha_{\text{base}}=0.025$, replace GeLU with SwooshL.

When trained for 12000 iterations with 1B tokens, the validation loss values (negative log-likelihood) are:

Exp1: 0.947; Exp2: 0.9369; Exp3: 0.9229

When trained for 24000 iterations with 2B tokens, the validation loss values (negative log-likelihood) are:

Exp1: 0.9068; Exp2: 0.885; Exp3: 0.8725

How could you make the model work in streaming fasion?

The way to make conformer streaming also applies to Zipformer. Specifically, for attention modules, in training we could limit the future contexts by applying the block-triangular masks on attention weights. We should also replace regular convolutions with causal convolutions. In inference, we could cache the left contexts for attention modules and convolution modules.