Paper: Zipformer: A faster and better encoder for automatic speech recognition (ICLR-2024)

Reviewer comment (Experiments):
There is uncertainty about whether the proposed optimizer, which has many adhoc modifications compared with the much more widely used Adam optimizer, can be more widely used for other deep learning optimization tasks, or whether it only works for speech tasks, or only works for zipformer. Given that zipformer-L is a relatively small model (only 60M parameters) trained on limited data (<1000hrs) according to today's standard, will the proposed change still work better than standard Adam when the model and the training data scale up?

Author response:
> 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.

Reviewer comment (Experiments):
The authors propose changes to many widely used components, which have been well tested over time and over different tasks beyond just speech recognition. For example, the authors claim that the mean extraction in layer norm is not necessary and also remove the variance normalization. While this may seem to work for the authors’ use case, it is unclear how generalizable this claim can be.

Author response:
> 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.

Reviewer comment (Experiments):
Some of the proposed modifications seem arbitrary. For example, in Eq (4), for the swooshR and swooshL functions, it is not explained how the -0.08 is selected. The authors mentioned that the coefficient 0.035 in Eq (4) is slightly tuned. The generality of this parameter is also not addressed.

Author response:
> 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).

Reviewer comment (Evaluation):
In the ablation study, quite a few modifications result in less than 0.2% absolute WER change. It is unclear how sensitive the WERs on librispeech are relative to other factors like checkpoint choosing and different training runs with different seeds. In the reviewer’s own experiments, these factors can result in up to 0.1% WER changes.

Author response:
> 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.

Reviewer comment (Novelty):
The authors miss a series of work in the literature that directly apply standard transformer (without much modification) for ASR. For example, Y. Wang, A. Mohamed, D. Le, C. Liu et al., “Transformer-based Acoustic Modeling for Hybrid Speech Recognition,” in Proc ICASSP, 2020, reports a WER of 2.26% / 4.85% on test-clean/test-other of librispeech, achieved by CTC + external NNLM. This is followed by Zhang, Frank, et al. "Faster, simpler and more accurate hybrid asr systems using wordpieces." arXiv preprint arXiv:2005.09150 (2020), which achieves a WER of 2.10% and 4.20% using a similar architecture.

Author response:
> Missing references

We have added these references in Section 2.

Reviewer comment (Experiments):
In Table 5 and section 4.0.3, it is unclear why "no temporal downsampling" results in an increase in the number of parameters (from 65.6M to 94.2M). All the major components in zipformers: attention, feedforward, and convolution use the same parameters for different temporal resolutions.

Author response:
> 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).

Reviewer comment (Experiments):
There are maybe too many new things being introduced here, which are all interesting by themselves, but each of them would maybe require more investigation and analysis on their own. E.g. introducing a new optimizer (ScaledAdam) is interesting, but it should be tested on a couple of different models and benchmarks, and this would basically be a work on its own. Now we have way too little analysis for each of the introduced methods to really tell how good they are. The ablation study is basically only a single experiment, showing ScaledAdam vs Adam, or LayerNorm vs BiasNorm. A single experiment is not really enough to really see how ScaledAdam vs Adam performs, or BiasNorm vs LayerNorm, etc. E.g. replacing LayerNorm by BiasNorm in a standard Transformer, maybe for language modeling, how would this perform? Many of these introduced methods seem to be quite orthogonal to each other.

Author response:
>  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.

Reviewer comment (Experiments):
For comparisons, it would have been better to use a more standard transducer, not a pruned transducer. Or maybe even CTC, or CTC/AED. This would have made it easier to compare the results to other results from the literature (Table 2).

Author response:
> 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.

Reviewer comment (Evaluation):
Measuring just FLOPs can sometimes be misleading as an indicator for speed, because certain operations might be faster than others, and certain operations can be executed in parallelized, while others can not. It would be interesting to also see the actual measured speed on some given hardware.

Author response:
> 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.

Reviewer comment (Experiments):
No systematic comparison to other Conformer variants, e.g. E-Branchformer, etc. (I don't just mean to have the number from literature as a comparison, as done in Table 2. I mean that the differences are compared, i.e. discussed and maybe also experimentally compared.)

Author response:
> 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.

Reviewer comment (Writing):
The capitalization is a bit inconsistent. Transformer is always lower case (transformer), while Conformer, ZipFormer etc are all capitalized.

Author response:
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.

Reviewer comment (Reproducibility):
It would make sense to state that (or whether) the code is published (even if you want to put the link only in the final version).

Author response:
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.

Reviewer comment (Reproducibility):
How long does the training take? What batch sizes are used?

Author response:
> How long does the training take? What batch sizes are used?

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

Reviewer comment (Experiments):
The downsample module averages every 2 frames with 2 learnable scalar weights (after softmax normalization) — I don't exactly understand how this is done. For the whole module, there are two learnable scalar weights only? So it could learn to always take really the average, but also to always take the first and ignore the second, or vice versa? This does not make sense to me. Why does it make sense to learn that? What happens if you just take the average? Or maybe use max-pooling instead? What is the influence of this?

Author response:
> "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

Reviewer comment (Experiments):
If you downsample at the very end to 25 Hz, why not do that earlier, and downsample right in the beginning to 25 Hz, and then never to 50 Hz again? That would make it quite a bit faster. So how much worse does this get?

Author response:
> 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.

Reviewer comment (Theory):
LayerNorm/BiasNorm: When there is a large value, as even explicitly modeled by the bias in BiasNorm, isn't this a problem, that there is a large number in the denominator, so then all the values become very small? The exp(γ) can compensate that then again, however, I wonder if there are numerical problems.

Author response:
> 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.

Reviewer comment (Experiments):
What is the influence of exp(γ) vs just γ (without exp) in BiasNorm?

Author response:
> 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.

Reviewer comment (Experiments):
How does ScaledAdam compare to methods like WeightNorm, where the weights are also normalized?

Author response:
> 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.

Reviewer comment (Experiments):
Since Adam is nearly invariant to changes in the gradient scale, for simplicity we replace this with ht = gt ·(rt−1 ⊙θt′−1) = gt ·θt−1 — 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?

Author response:
> 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.

Reviewer comment (Reproducibility):
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?

Author response:
> 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.

Reviewer comment (Writing):
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.

Author response:
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.

Reviewer comment (Evaluation):
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?

Author response:
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.

Reviewer comment (Experiments):
Curious about whether the authors tried the new activation functions, biasnorm and scaledadam on other tasks? Do they still show superiority?

Author response:
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

Reviewer comment (Experiments):
How could you make the model work in streaming fasion?

Author response:
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.