Scaling Transformers for Low-Bitrate High-Quality Speech Coding

Dear reviewer Baq6,

Thank you for your efforts in reviewing this paper, and for raising your concerns with the work. We have grouped your questions and comments and provided detailed responses to each point below.

Q1: 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.

We agree that scaling is a well known result in other fields (e.g. LLMs), which is one of the primary motivations for us pursuing this direction (as articulated in Sec.1). However, we’re not aware of existing work showing that audio codec performance scales with parameter count. If we’re missing some important work that establishes this, it’d be much appreciated if you could share it. Our opinion (explained in Sec.1) is that this was not a goal for previous codec models as they were more directly targeting the transmission/storage use-cases, which benefit much more from small parameter-efficient models. This is also likely to be related to the overall difficulty of scaling CNNs to the billion parameter level (see e.g. ConvNeXTv2). In our work we establish that:

1. Leveraging transformer blocks as the primary component of a model is viable for an end-to-end audio waveform codec, once a number of architectural challenges are overcome.
2. Scaling parameter count does indeed provide improved performance in audio waveform coding (which we also now demonstrate more directly in the new ablation in Appendix A.2).

We strongly believe that both are novel contributions to the field of audio coding, which will be helpful in informing future work.

Q2: 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. I do not think DAC, Encodec, and SemantiCodec are reasonable baselines to compare to, as none claim to be English speech codecs.

We are certainly not claiming this as a point of novelty, nor is it something we are trying to demonstrate. We acknowledge the relatively limited domain (English speech) as a limitation of this initial iteration of our model, and we intend to address this by scaling to more diverse datasets in the future. This has been made more explicit in Sec.3.7. We would argue that the presented results, along with the new experiment in Appendix A.5 which shows good generalization to unseen languages, strongly suggest that model performance will increase even further when extended to a larger and more diverse dataset. Our opinion is that this enhances the attractiveness of this architecture as the basis for future work.

Regarding the fairness of baselines choices, we agree that all the chosen baselines were not exactly designed with the same goals as our proposed model, nor necessarily trained on the same data. This is not an intentional choice on our part, as we’re restricted by which models are released to the public. The chosen baselines represent generally the most widely used and positively considered neural codecs for speech, showing strong performance. Hopefully there will be fairer baselines available in the future for this comparison, as larger neural audio codecs specifically targeting downstream generative use cases become more prevalent.

Hi reviewer Baq6,

We'd really appreciate if you could take a look at the changes we made to address you're concerns. We think the results on multi-lingual generalisation are especially interesting! We also made some significant changes to address your other concerns. We're available here to discuss.

Dear reviewer bLHu,

Thank you for your efforts in reviewing this paper, and for your many insightful comments. We’ve tried to address your concerns with the following changes:

Q1: 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 ϵ is provided.

In Appendix B.1 we have added some further discussion about the ϵ value of the layer norms, as well as the exact value used in the experiments. Thank you for spotting this missing piece of information.

Q2: 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.

We have expanded Sec. 2.4. to give much more detail about the training objective, in both the pretraining and finetuning stages.

Q3: 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?

• We’ve added Appendix B.2 discussing the relative receptive fields of our proposed model and the baselines. Interestingly, our proposed model does not have a wildly different receptive field to existing models - partly as very few are purely convolutional. Several baselines use RNNs with effectively unlimited receptive fields, and use chunked inference to counteract any downside from this.

• We’ve added an experiment addressing length generalization by testing inference of the TAAE and baselines with a variety of utterances lengths. This can be found in Appendix A.6. The results show mild degradation at longer utterances, which can be mitigated if necessary by chunked inference.

• We’ve added an experiment which examines empirical inference time. This can be found in Appendix A.9. This experiment shows that the TAAE architecture is competitive with the baselines in terms of empirical inference time, despite utilizing a much larger number of parameters. The architecture benefits greatly from the patched representation taking care of much of the down/upsampling, as well as the large amount of effort invested by other researchers and engineers in optimizing the components of the transformer architecture.

• We have interpreted ‘time-complexity’ to mean Big O analysis with respect to varying sequence length. This is now discussed briefly in Appendix B.2. Our model (and all the baselines) scale as O(n) with sequence length.

Dear reviewer VKr9 ,

Thank you for reading our paper and for providing such insightful comments. We made the following changes to address your review:

Q1: “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).”

• We trained a causal version of the codec as a finetune from the main presented model. This is presented in Appendix A.4. We achieved this by switching convolutions to causal versions and changing the sliding attention window to be purely causal. The model performs very close to the main model presented in the paper, so the trade-off for applying the model in a streaming situation appears to be minimal, assuming sufficient computation resources are available to execute the model. We did not want to alter the main intention of the paper (we were not concentrated on streaming use-cases), so this is kept as an extension rather than used to replace the non-causal version of the model in the main comparisons. New audio examples created by the causal TAAE model can be found on the anonymous website.

• We have added some discussion in Appendix B.2 which compares the receptive fields, causality and latency of our proposed model and the baselines.

Thank you again for your reviewing efforts. We’re happy to address any more comments or concerns you may have.

Dear reviewer 1mKc,

Thank you for your efforts in reviewing this paper, and for the very insightful comments. We’ve made a number of modifications in order to address your concerns.

Q1: "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."

Our initial draft did not contain discussions on the Huffman coded bitrates of the proposed model or the baselines, as our use-case for this model was not focused on transmission or storage of speech. However, we agree that this information is important and valuable to the wider speech coding community. To address this, we have added a new section in Appendix A.8 that examines and computes the codebook utilisation and Huffman-coded bitrates for the proposed model and the baselines. Our analysis shows that Huffman coding achieves a moderate reduction in bits-per-second for earlier RVQ-based codecs such as Encodec. In contrast, more recent RVQ-based codecs (e.g., DAC, SpeechTokenizer, Mimi), which leverage advanced methods like EMA updates, factorized codes, and L2-normalized embeddings to boost RVQ codebook utilisation significantly, derive only marginal benefits from Huffman coding. Similarly, FSQ-based tokens produced by TAAE achieve comparable levels of gains in coding efficiency with Huffman coding to these modern RVQ models.

Q2: "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."

We conducted scaling experiments with TAAE architectures containing approximately 250M, 500M, and 1B parameters, which is shown in Appendix A.2. The results demonstrate that scaling up the parameter count yields clear improvements in objective metrics, although the smaller models still perform respectably compared to the baselines.

We would also like to highlight that scaling transformer-based codec architectures to 1B parameters is a key contribution of this work. Unlike traditional CNN-based codecs, TAAE uses a transformer-based architecture that offers enhanced scalability. However, in our early experiments we found it challenging to make Transformers work effectively for audio coding tasks. The success of the TAAE model is due to our specific architecture design, empirical findings, and optimizations, as described in Section 2 and further analyzed in Appendix B.

Hi reviewer 1mKc, Thanks for updating your score! If you have any remaining concerns with the work, please feel free to share and we can discuss.

Hi Zhen,

Thank you for your kind words. We're happy that you enjoyed the paper!

I'm sorry—I really wanted to try your Patched Transform scheme, but I'm not familiar with DSP, so implementing the code is hard for me.

Here's an implementation for you:

class PatchedPretransform(nn.Module):
    def __init__(self, channels, patch_size):
        super().__init__()
        self.channels = channels
        self.patch_size = patch_size
    def encode(self, x):
        x = rearrange(x, "b c (l h) -> b (c h) l", h=self.patch_size)
        return x
    def decode(self, z):
        z = rearrange(z, "b (c h) l -> b c (l h)", h=self.patch_size)
        return z

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

Yes, we've tried this - although we haven't fully trained a model at this temporal resolution as we were targeting a lower bitrate. The reconstruction performance should be better than the 25Hz model we presented (assuming the bottleneck is the same) as the compression ratio is less.

I noticed that the number of channels in your discriminator implementation is 256. I'm using the MRD from DAC (link). 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?

We're using the Encodec discriminator which has less params in general than the DAC discriminator. Total param count is 35.5M.

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:

Your code looks correct, but you'll need to set normalized = True in the torch.stft call - otherwise it'll produce very large values. This is turned on in Encodec (I'd assume also in DAC) as otherwise you'll get huge activations inside your discriminator.

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?

Our feeling is that a powerful encoder actually helps the model effectively communicate information through the very restrictive FSQ bottleneck. We did experiment with smaller encoders, but we found that this makes the decoder act more like a generative model - we could still get plausible speech output but the alignment with the input was less precise. Investigating this relationship in more detail (and if it holds with VQ instead of FSQ) could be a very interesting line of future research.

Thanks again for your comments! Happy to discuss more if you have further questions.