Scaling Transformers for Low-Bitrate High-Quality Speech Coding

Authors: Julian D Parker, Anton Smirnov, Jordi Pons, CJ Carr, Zack Zukowski, Zach Evans, Xubo Liu
Primary area: generative models
Keywords: Audio coding, neural audio codecs, transformers

Abstract:
The tokenization of audio with neural audio codec models is a vital part of modern AI pipelines for the generation or understanding of speech, alone or in a multimodal context. Traditionally such tokenization models have concentrated on low parameter-count architectures using only components with strong inductive biases. In this work we show that by applying a transformer architecture with large parameter count to this problem, and applying a flexible Finite Scalar Quantization (FSQ) based bottleneck, it is possible to reach state-of-the-art speech quality at extremely low bit-rates of $400$ or $700$ bits-per-second. The trained models strongly out-perform existing baselines in both objective and subjective tests.

Decision: accept (poster)
Reviewer scores: 6, 8, 8, 6

Reviews and discussion:
<thread>
creation date:
1730664874968

thread_type:
official_review

<thread_object>
creator: reviewer_baq6

<content>
summary:
The authors present an encoder-decoder transformer-based architecture for encoding speech at very low bitrate, below 1kbps. The Transformer Audio AutoEncoder (TAAE) uses a polyphase time-frequency representation to perform downsampling/upsampling before the enocder and after the decoder of TAAE. Finite Scalar Quantization (FSQ) is employed within the encoder/decoder bottleneck to mitigate codebook underutilization typically seen with vector quantized (VQ) and residual vector quantized (RVQ) approaches. The authors combine an L1 discriminator feature loss with decaying L1 waveform loss and late perceptual reconstruction loss for training. The TAAE is trained on 105k hours of English speech sampled at 16kHz. The reconstruction capability of TAAE is compared to the Descript Audio Codec (DAC), Encodec, SpeechTokenizer, SematiCodec, and Mimi. A mean opinion score (MOS) is also produced from a perceptual evaluation comprised of 23 participants comparing TAAE to Mimi and SemantiCodec. The authors demonstrate that TAAE obtains better reconstruction performance according to both objective measures and MOS. The authors also demonstrate that one variant of the TAAE codebook attains 98% utilization.

strengths:
Originality:
The results presented in Appendix B are enlightening regarding the use of non-power-scaled FFT sizes. The FSQ-based bottleneck also seems to overcome common issues in the training of RVQ systems. 

Quality:
The authors provide a wide variety of objective assessments for their architecture's performance. The authors also do a good job of citing current literature.

Clarity: 
The authors very clearly describe their architecture and the motivations for their architectural decisions. The appendices were well organized and helpful.

Significance:
Appendix B and the FSQ bottleneck are worthwhile contributions.

weaknesses:
I do not think it is novel to scale the parameter count of a neural network autoencoder and demonstrate better compression ratios/reconstruction compared to smaller architectures.  This is a well known result.

I do not think it is novel to restrict the domain of a neural network autoencoder in comparison to another architecture trained on a more general domain and demonstrate better compression ratios/reconstruction. This is a well known result.

By restricting the domain of their speech audio corpus to English speech, the authors have produced an English speech audio codec. In order to claim that this is a "speech codec," the authors should evaluate on non-English speech to demonstrate generalization capabilities.  

I do not think DAC, Encodec, and SemantiCodec are reasonable baselines to compare to, as none claim to be English speech codecs. 

Mimi focuses on streaming and causality with 1/10 the parameter count of TAAE, which makes no claims regarding streaming capability. This also leads to an odd comparison as the goals of Mimi and TAAE are not aligned.  

It is unclear why SpechTokenizer was left out of the perceptual evaluation, as it is the most comparable to TAAE in terms of architecture and training domain. Comparison to SpeechTokenizer could also boost claims that the FSQ significantly outperforms RVQ schema.  

I think the presented MOS scores are confusing, as MOSNet has estimated the MOS closer to 3 than to 5. A MUSHRA evaluation should have been used instead for pairwise comparison between codecs and is standard in the literature cited in this paper. Furthermore, the authors should include demographic breakdowns of the perceptual evaluation, as well as a description of the listening setup, as is standard for speech codecs.

questions:
Why was SpeechTokenizer not included in the perceptual evaluation?

How do the design goals of Mimi align with that of TAAE? Why is Mimi a good baseline comparison to TAAE?

Why was a MOS evaluation chosen instead of MUSHRA?

How do you explain the gap in your perceptual evaluation MOS score and the estimate provided by MOSNet? 

Who was included in the perceptual evaluation, and what was their listening setup?

How does TAAE perform on non-English speech? And how does that compare to the more generalist NAC?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730667468983

thread_type:
official_review

<thread_object>
creator: reviewer_blhu

<content>
summary:
This paper proposed  TAAE, an audio codec model that uses Transformer encoders as the main building block to replace conventional convolution-based modules. To accommodate the choice, TAAE performs downsampling mainly by patchifying the time domain signal and training transformer encoder stacks on top of the downsampled sequence. For discretizing audio, TAAE relied on FSQ-based bottleneck that approximates continuous low-dimensional latent numerically. Experiment results show TAAE achieved outstanding speech quality on autoencoding at a significantly lower bit rate comparing to existing models.

strengths:
- The idea of using Transformer and the main architecture for the neural audio codec learning is novel and well executed.
- Judging from the audio samples on the demo page and MOS study, TAAE is clearly state-of-the-art in low bit rate speech compression.
- This paper provided a lot of detailed knowledge, empirical findings, and engineering improvements that can truly benefit the audio codec research community. I personally learned a lot in the details such as the discussion on systematic bias of discriminator, choice of filterbank, observation on the latent representation of silence frames with self-attention, etc.

---

### Justification for rating
Overall, I believe this work is novel enough and provides solid contributions to the field.
However, some improvement might be necessary (see weaknesses below).
If the authors can properly address these concerns and update the submission accordingly, I would be more than happy to raise my rating on this paper.

weaknesses:
Given the main contribution of this work is in exploring an alternative architecture for codec models, completeness in terms of design details and reproducibility are expected. In contrast, I found a lot of details missing or vague. (Although the authors state the code will be released later, the paper itself should still be comprehensive alone.) Here are some examples:

---

> ($\S$2.1) ... Instead we raise the $\epsilon$ constant used in the calculation of normalization factors in the layer norm blocks ... allows convergence of the architecture.

This appears to be an interesting observation and a critical hyper-parameter for training the model as the authors spent a paragraph discussing it, but neither the exact value nor the study/experiment on $\epsilon$  is provided.

---

> ($\S$2.4)...For training the codec itself, we primarily use a normalized feature-matching L1 loss on the per-layer features of the discriminator network ... In addition we found it beneficial to include a traditional L1 reconstruction loss to boost convergence at the beginning of the training process ...

The overall objective of the model is not explicitly given but described in a more hand-wavy style instead, which could easily lead to misunderstanding. The full objective should be listed explicitly together with the weight/importance for each term/component in the main paper or appendix.

---

> ($\S$2.1) ... The self-attention uses a sliding window of size 128, to restrict receptive field and aid generalization of the architecture to arbitrary length sequences. 

This choice appears as one simple sentence, but self-attention is the key difference between TAAE and prior works, which changes the properties of the model dramatically.
If my understanding is correct, this means the receptive field of the first layer is already 2.56 seconds (128 frames $\times$ 20 ms-per-frame), and the number doubles for every layer. It is obvious that TAAE has a much larger receptive field size comparing to convolution-based models. While this is an advantage, it could also lead to some problems that are not discussed in the paper.
 
- What is the trade-off between length generalization and sliding window size for TAAE? How do time complexity and empirical inference time change accordingly? How do these numbers compare to those of CNN-based models?
- Beyond length generalization, can TAAE perform streaming encoding/decoding (as most of the existing works compared in this paper)?
  - If so, what is the size receptive field? how does it affect the latency of the codec? how does it compares to conventional CNN-based codec models?
  - If not, this should still be explicitly discussed as a limitation of the proposed framework in the paper.

These are just some examples. In short, I believe the fundamental differences between TAAE and CNN-based codec models should be discussed in the paper more throughout and carefully. Both advantages and disadvantages should be clearly stated and summarized in the main body of the paper.

---

I believe these concerns can all be addressed without additional training, thus should be easy enough to complete within the rebuttal period.

questions:
(please see weaknesses above)
</content>
</thread_object>

</thread>
<thread>
creation date:
1730697589488

thread_type:
official_review

<thread_object>
creator: reviewer_vkr9

<content>
summary:
This work describes an approach to leverage scaled transformer model architecture to achieve low-bitrate and high-quality speech coding. Different from conventional neural audio codec approaches, this work leverages the finite scalar quantization (FSQ) based bottleneck to achieve high codebook utilization rate and avoid the difficulty in training a VQ-VAE auto-encoder. Experimental results show this works outperform existing baselines in both objective and subject tests.

strengths:
This paper is well written and very clear to follow. In the introduction part, it clearly presents the motivations and has an excellent survey of the existing methods. 

Though using transformers to scale and leverage FSQ for high codebook utilization is not something new, this paper presents the motivations of these changes, the associated challenges and their mitigations. This paper also introduces a new method so that FSQ can be used in a similar way as RVQ where a varying bits-per-second rate can be achieved. 

This paper presents strong experiment results, significant improving over the existing baselines (but at the cost of increased computation and latency).

weaknesses:
If I understand the proposed model correctly, it is based on transformer layer with a local attention of 128 (both left and right), which means different from DAC/Encodec/Mimi etc which use causal encoders, the encoder in the proposed method is not causal, and it will introduce a latency up to the patch length (which is 320/16k ~ 20ms?). It would be great if the author can present the results with causal encoder so that it can be compared with DAC/Encodec/Mimi in a relative fair comparison (apart from the model size difference).

questions:
N/A
</content>
</thread_object>

</thread>
<thread>
creation date:
1730764285836

thread_type:
official_review

<thread_object>
creator: reviewer_1mkc

<content>
summary:
The paper presents a new Transformer architecture for speech coding. It is characterized by the new scalar quantization method performed in a dimension-reduced space, which showed improved coding gain compared to other methods that are based on residual vector quantization. The paper also provides various training strategies that appear to be useful for neural codec training.

strengths:
- The proposed system appears to work well according to the objective metrics and subjective tests.
- The proposed FSQ idea seems to be a solid quantization option, improving the codebook utilization.
- The authors put a lot of effort in making it more scalable by adding multiple levels of quantization.

weaknesses:
- The proposed method relies on the dimension reduction part for its dimension-specific scalar quantization to work. And that's why they could achieve higher codebook utilization. Meanwhile, there is also a trend that higher codebook utilization leads to lower coding gain if entropy coding is applied after tokenization. Indeed, the paper does not mention anything about Huffman coding results, which the proposed method might not be able to take advantage of due to the low dimensionality and high codebook utilization. At the same time, the RVQ-based ones might have a better chance of compressing more via Huffman coding. I wish the paper provided an in-depth discussion about it. In my opinion, all the coding gain and performance-related arguments must be based on the entropy-coded bitrates of all codecs mentioned. 

- The other main criticism is that the proposed model is just a lot bigger than the other models. I don't mean to argue that a bigger codec necessarily results in a better coding gain, but in general, it is also true that there is a relation. I wish the paper had provided an ablation test that investigated the impact of the different sizes of the proposed model. 

- The paper provides various tips and useful information about their model training, but they are scattered in different places without a clear organization.

questions:
- How does it compare to the HuBERT-based codec? HuBERT can be as large as the proposed model (its X-Large version), and can turn into a codec with a vocoder attached to it as shown in [a].

[a] A. Polyak et al. “Speech Resynthesis from Discrete Disentangled Self-Supervised Representations,” Interspeech 2021
</content>
</thread_object>

</thread>
<thread>
creation date:
1732300680403

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General comments to all reviewers

comment:
Dear reviewers,

We have now uploaded a new version of the paper with a significant amount of new content that we believe should address the majority of your comments. We've replied to each of you individually to give specific replies to your concerns, but we also wanted to highlight here the new content suggested by reviewers.

* We performed an ablation showing that the performance of the TAAE scales clearly with parameter count. This is described in **Appendix A.2**.

* We trained a causal finetune of the presented model. This demonstrates minor performance degradation compared to the main presented model but is competitive with the strongest streaming-focused baseline. This is described in **Appendix A.4**.

* We re-ran the subjective perceptual test with a more conventional MUSHRA format, and the addition of SpeechTokenizer as an extra baseline. The results are consistent with those presented in the previous version of the paper.

* We tested the generalization of the presented model and baselines to unseen languages, by evaluating objective metrics on the Multilingual LibriSpeech dataset. The model shows strong generalization performance, reaching parity or outperforming models trained on multi-lingual or general audio data. This is presented in **Appendix A.5**.

* We tested the generalization of the model and baselines to utterances of a variety of lengths. This is presented in **Appendix A.6**

* We performed a comparison with an alternative style of codec model, constructed using a pretrained HuBERT semantic model a HifiGAN-derived vocoder model. This is presented in **Appendix A.7**.

* We performed entropy coding experiments on the model and baselines. This revealed that our model and most baselines exhibited high codebook utilization and, and did not benefit strongly from entropy coding. This is presented in **Appendix A.8**.

* We measured RTF for the proposed model and baselines, showing that despite a much larger parameter count, inference time for the TAAE is competitive with baselines. This is presented in **Appendix A.9**

* We added further clarification about the training objective in **Sec.2.4**, and clarified many hyperparameters in **Sec.3.2**

* We added discussion of the receptive field, causality and latency of the presented model and baselines as well as additional discussion about the merits of convolution vs attention. This is presented in **Appendix B2--B3**

* We added new audio examples demonstrating the generalization of the model to unseen languages and also demonstrating the causal variant of the model. These can be found at the following link: https://taae-iclr-2025.github.io/taae_anonymised/

Thank you for working with us on making this paper the best it can be! We very much appreciate your input.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732776637226

thread_type:
public_comment

<thread_object>
creator: ~zhen_ye2

<content>
title:
Excellent paper

comment:
Hello,

I am very, very impressed with your work, as I am currently researching codecs with single VQ and Transformer decoders. Your work has given me a lot of inspiration. If it's convenient, I would love to discuss it with you.

First, I tried your proposed new discriminator  FFT parameters {78, 126, 206, 334, 542, 876, 1418, 2296} to replace old one { 128, 256, 512, 1024, 2048} and found it works exceptionally well, improving the speaker similarity before and after from 0.80 to 0.82.

My current experiments are trained on 16kHz audio, using 50 tokens per second (50x1), with an STOI of about 0.91 and a PESQ of 3.01 and sim of 0.82. My decoder is a Transformer (hidden size 1024, 12 layers) with vocos. I'm sorry—I really wanted to try your Patched Transform scheme, but I'm not familier with DSP, so implementing the code is hard for me. I'm also using FSQ with a codebook size of 65536 (4^8).

I have a few questions I'd like to ask you:

1. Have you tried the 50Hz x1 setting? Or, what performance do you expect your method to achieve under this setting?

2. I noticed that the number of channels in your discriminator implementation is 256. I'm using the MRD from DAC ([link](https://github.com/descriptinc/descript-audio-codec/blob/c7cfc5d2647e26471dc394f95846a0830e7bec34/dac/model/discriminator.py#L136)). I tried setting the channels to 256, and the discriminator's parameters are about 100M. Is this correct? How does it compare to the number of parameters in your discriminator?

3. Regarding SYSTEMATIC BIAS, I tried it but suspect there may be an issue with my implementation. Here's my code; I'm not sure if I understood your paper correctly:

```python
x_stft = torch.stft(x, fft_size, hop_size, win_length, window.to(x.device), return_complex=True)
gamma = 1/2
X_magnitude = torch.abs(x_stft)
x_stft = x_stft * torch.pow(X_magnitude + 1e-8, gamma)
```

4. About the encoder, I tried using mel + Transformer or STFT + Transformer, but the former resulted in worse performance, and the latter's performance was similar to DAC's encoder (about 40M parameters). I'm not sure if your method has experimented with different encoders. If you still use the CNN-based DAC encoder, is there a significant difference in your method? My thinking about the encoder is that since the encoder's features have to go through VQ, it means that no matter how powerful your encoder is, after VQ, it can only retain key information. Does this mean that for the encoder, maybe we don't need a very powerful model—perhaps DAC's encoder is sufficient?

If you are busy, please don't feel pressured to reply immediately, but I am really looking forward to discussing this with you.

Thank you very much for your paper. I have discussed your paper with many collaborators working on codec, and we all think it's an excellent paper. Your discussions on various technical details are really great. I'm very certain that your work will have a significant impact on the development of audio codec!

Looking forward to your reply!
</content>
</thread_object>

</thread>


Metareview:
metareview:
**Paper Summary:**

This paper describes low-bitrate neural discrete audio codec. This is achieved by scaling up an encoder-decoder Transformer architecture using finite scalar quantization (Mentzer et al., ICLR 2024) in contrast to the more traditional a VQ-VAE discretization. Experimental results are strong for both quantitative metrics and human evaluation.

**Strengths:**

There is general consensus that the empirical analysis is strong, demonstrating a new state-of-the-art in low bit rate speech compression.

All reviewers praised the paper's insights. There is also mostly positive praise for the paper's clarity and organization.

The authors are committed to releasing the artifacts of this work, which will be of great value to the community.

**Weaknesses:**

As Reviewers 1mKc and Baq6 point out, achieving better results by scaling up neural networks is no longer a surprising result for the community. On the other hand, this work identifies and documents many valuable insights into the architectural decisions required to effectively scale neural audio codecs.