arXiv · 2024 · Preprint
Parker et al. (Stability AI) · → Paper · Demo: ✓ · Code: ✓
A predominantly transformer-based neural audio codec (TAAE) scaled to ~950M parameters with a Finite Scalar Quantization bottleneck achieves state-of-the-art speech coding quality at 400 and 700 bits per second, outperforming all convolutional baselines in both objective and subjective evaluation while operating at lower bitrates than any of them.
Problem
Neural audio codecs have been almost exclusively built on convolutional or recurrent architectures, which offer strong inductive bias and parameter efficiency but are notoriously difficult to scale beyond roughly 1B parameters. Meanwhile, modern generative speech pipelines (autoregressive language models, latent diffusion systems) collectively contain billions of parameters, yet the codec component — responsible for tokenizing speech for downstream generation — has remained comparatively tiny and under-scaled. A second limitation is the prevalence of Residual Vector Quantization (RVQ) in codec bottlenecks: RVQ produces hierarchical multi-stream token sequences whose causal cross-stream dependencies complicate training and inference in generative models, and the codebooks can suffer from poor utilization.
Method
The paper introduces the Transformer Audio AutoEncoder (TAAE), a codec architecture in which both the encoder and decoder are built primarily from transformer blocks rather than convolutions. Each encoder/decoder stage consists of a strided 1D convolution for temporal down- or upsampling, followed by a chain of transformer blocks using pre-norm layer normalization, QK-norm, sliding-window self-attention with RoPE positional embeddings, and gated MLP feedforward blocks with LayerScale for training stability. The key design insight is to front-load temporal downsampling: rather than performing it progressively inside the transformer stack (which would produce very short embeddings in early layers), most downsampling is handled at the input and output boundaries using a polyphase filterbank (patched) transform, keeping sequence lengths tractable for self-attention. The main model has two encoder blocks (8 and 20 transformer blocks respectively) and a symmetric decoder, with an embedding dimension of 1024, giving approximately 950M parameters.
The bottleneck uses a modified Finite Scalar Quantization (FSQ) scheme in place of RVQ. FSQ projects the latent representation to a low-dimensional space (6 dimensions by default) and scalar-quantizes each dimension independently into a fixed number of levels, mapping every combination to a unique integer token. Near-perfect codebook utilization follows from this formulation without any explicit utilization regularization. Training mixes uniform noise quantization (50% of steps) with straight-through gradient estimation, and randomly samples quantization level configurations during training to allow post-training flexibility in token resolution. Crucially, the FSQ bottleneck supports post-hoc residual decomposition: a model trained with a single-token-per-frame bottleneck can be reconfigured after training into a hierarchical residual tokenization (Residual FSQ) without retraining, enabling the bitrate to be adjusted at inference time. The 400 bps variant uses a single token per frame at 25 Hz, and the 700 bps variant decomposes into two residual tokens per frame.
Training proceeds in two stages. Pretraining uses an adversarial feature-matching loss from a multi-resolution complex STFT discriminator, augmented with exponentially decaying L1 waveform and STFT reconstruction losses. The discriminator uses FFT sizes chosen to be maximally inharmonic (spaced by the golden ratio) to avoid systematic spectral biases observed in standard power-of-two configurations. Finetuning replaces the reconstruction losses with a WavLM-Large perceptual feature-matching loss, which the paper shows is essential for achieving intelligible, high-quality reconstruction (Table 3, Appendix A.1). Training uses 105k hours of 16 kHz English speech (Librilight and MLS-English), 16 H100 GPUs, and FlashAttention for efficiency.
Key Results
On LibriSpeech test-clean, the 700 bps TAAE variant achieves PESQ 3.09 and STOI 0.92, outperforming all baselines by a substantial margin despite operating at a lower bitrate than most of them (Mimi at 1100 bps, SpeechTokenizer at 1000 bps, DAC at 1000-2000 bps). The 400 bps variant reaches PESQ 2.96 and STOI 0.90. On the MOSNet automatic quality predictor, both TAAE variants score 3.36, compared to Mimi’s 3.24 at 1100 bps (Table 2).
The MUSHRA subjective test with 24 participants confirms the objective pattern: TAAE significantly outperforms all baselines and approaches ground-truth quality, a result that none of the baselines achieves at any tested bitrate (Figure 2).
Scaling ablations (Table 4) confirm consistent quality improvement as parameter count grows from 240M to 540M to 950M, with all metrics improving monotonically. A causal streaming variant (Table 6), obtained by finetuning with causal sliding attention and causal convolutions for 200k steps, performs competitively with Mimi at 1100 bps despite the TAAE running at 700 bps, with a streaming latency of 40 ms versus Mimi’s 80 ms.
On multilingual LibriSpeech (Table 7), TAAE trained on English only outperforms all baselines across all seven evaluated languages on most metrics, including Mimi which was trained on 7M hours (roughly 70 times more data), though Mimi leads slightly on PESQ in several languages.
Real-time factor (Table 10) on an H100 GPU is 0.014 for 5s clips, slower than lightweight baselines (DAC: 0.005, Mimi: 0.005) but competitive with longer utterances, and over 80 times faster than real-time for 30s+ audio, making it viable in batch processing contexts.
Novelty Assessment
The core contribution is architectural: replacing the convolutional backbone of neural audio codecs with a predominantly transformer-based structure and demonstrating that this enables meaningful scaling (250M to 950M) with consistent quality gains. The observation that transformers’ input-dependent attention is better suited to the uneven information density of speech (dense during phonation, sparse during silence) provides a principled motivation beyond analogical transfer from vision. The FSQ bottleneck itself is not new (Mentzer et al., 2023), but the post-hoc residual decomposition (Residual FSQ) is a genuine contribution: it allows a single trained model to emit tokens at different bitrates without retraining. The discriminator design (inharmonic FFT sizes motivated by a per-sample sensitivity analysis, magnitude scaling to suppress noise-floor bias) is a careful engineering contribution that addresses systematic failure modes in adversarial training.
The main limitation is the restricted training scope: 16 kHz English speech only, 105k hours, audiobook-dominated. This is notably narrower than Mimi (7M hours) and constrains the generalization claims.
Field Significance
Tip
High — TAAE demonstrates that the transformer scaling paradigm transfers to neural audio coding, shifting the design space of codecs used by downstream speech generation systems. Its 400–700 bps operating points and single-token-per-frame tokenization simplify generative modeling of speech, and the post-hoc residual FSQ decomposition provides a practical tool for adapting a single codec to varied inference budgets. With 40 corpus citations in under a year, it provides empirical evidence that the field views large-parameter-count transformer codecs as a viable, high-value direction rather than an engineering curiosity.
Claims
- Scaling transformer architecture parameter count in neural audio codecs produces consistent quality improvements across objective and subjective metrics. (§4.6, Table 4)
- Finite scalar quantization achieves near-perfect codebook utilization without explicit utilization regularization, simplifying downstream generative modeling compared to RVQ. (§3.2, §A.8, Table 9)
- A neural codec trained exclusively on English speech can generalize effectively to unseen languages, outperforming multilingual-trained baselines of similar scale on most objective metrics. (§A.5, Table 7)
- Perceptual losses derived from self-supervised speech models (WavLM-Large features) are critical for achieving intelligible reconstruction at low bitrates, beyond what adversarial and spectral reconstruction losses alone provide. (§3.4, §A.1, Table 3)
- Systematic spectral bias in multi-resolution STFT discriminators, arising from power-of-two FFT configurations, causes periodic reconstruction artifacts that disproportionately affect large-capacity codec architectures. (§3.3, §B.5)
Limitations and Open Questions
Warning
Training data is 16 kHz English audiobook speech only (105k hours). Multilingual generalization results are promising but the model was not trained or optimized for non-English data; claims about multilingual capability should be interpreted cautiously relative to models with dedicated multilingual training at scale.
The model has not been evaluated on noisy speech, overlapping speakers, or environmental audio, which are common real-world conditions. The large parameter count (950M) requires substantially more compute than lighter baselines (DAC at 76M, Mimi at 80M); while the RTF is acceptable on H100 GPUs for longer utterances, latency for short clips is roughly 3x that of smaller models, which matters for streaming applications. The post-hoc Residual FSQ decomposition is restricted to specific level configurations (L = 2^n + 1); arbitrary bitrate targets are not directly achievable without retraining. The systematic bias analysis in the discriminator (§B.5) raises open questions about whether similar biases appear in other convolutional discriminator architectures and how to address them in the general case.
Wiki Connections
Core architecture: neural-codec | autoregressive-codec-tts | spoken-language-model
In-corpus papers cited: 2305.09636 (SoundStorm) | 2410.00037 (Moshi) | 2210.13438 (EnCodec) | 2301.02111 (VALL-E) | 2104.00355 (speech resynthesis from disentangled SSL representations)
Evaluation context: evaluation-metrics | self-supervised-speech (WavLM used as perceptual loss)