Continuous First, Discrete Later: VQ-VAEs Without Dimensional Collapse

# Continuous First, Discrete Later: VQ-VAEs Without Dimensional Collapse
Source: [https://arxiv.org/html/2605.06870](https://arxiv.org/html/2605.06870)
Xinyu Zhao MIT lxz@mit\.edu &Nikita Karagodin11footnotemark:1 MIT nikitus@mit\.edu Hamed Hassani University of Pennsylvania &Sinan Hersek Google &Paul Pu Liang MIT &Yury Polyanskiy MIT

###### Abstract

While many approaches to improve VQ\-VAE performance focus on codebook size and utilization, the effect of dimensional collapse, where trained VQ\-VAE representations live in an extremely low\-dimensional subspace \(1−2%1\-2\\%of full rank\), remains unaddressed\. We show theoretically and empirically that dimension collapse causes a hard loss lower bound that various codebook improvement techniques fail to surpass\. Our analytic framework extends the sequential learning effect ofSaxeet al\.\([2014](https://arxiv.org/html/2605.06870#bib.bib27)\)by introducing ideas from rate\-distortion theory and explains how the latent collapse is caused by the VQ suppressing lower\-variance directions\. Our theory justifies a simple solution: a “warm\-up phase” that trains the model as an \(unquantized\) autoencoder before introducing VQ\. On both synthetic experiments and large\-scale image \(VQGAN\) and audio \(WavTokenizer\) VQ\-VAEs, we show that AE Warm\-Up successfully restores representation dimension, leading to lower reconstruction and perceptual loss at the same training budget\. Across codebook sizesK∈\{210,214,216\}K\\in\\\{2^\{10\},2^\{14\},2^\{16\}\\\}, AE warm\-up raises VQGAN codebook effective dimension from 3–5 to 17–19 and reduces rFID by 17–35%; on WavTokenizer atK∈\{213,214\}K\\in\\\{2^\{13\},2^\{14\}\\\}, it raises codebook dimension from 4 to 17–19 and improves PESQ by 11–14%\. We empirically characterize how warm\-up duration governs the achievable final loss\. In agreement with experiment, our theoretical analysis predicts downstream performance as a function of warm\-up length, enabling an adaptive criterion for switching from AE Warm\-up to VQ\-VAE training\.111Code for our experiments can be found at[https://github\.com/xyz\-zy/vqvae\_latent\_span\_collapse](https://github.com/xyz-zy/vqvae_latent_span_collapse)\.

## 1Introduction

Vector\-quantized variational autoencoders \(VQ\-VAEs\) are the backbone of modern discrete tokenization for audio and image representation learning and generative models\(van den Oordet al\.,[2017](https://arxiv.org/html/2605.06870#bib.bib3); Razaviet al\.,[2019](https://arxiv.org/html/2605.06870#bib.bib4); Esseret al\.,[2021](https://arxiv.org/html/2605.06870#bib.bib5); Zeghidouret al\.,[2021](https://arxiv.org/html/2605.06870#bib.bib13); Jiet al\.,[2025](https://arxiv.org/html/2605.06870#bib.bib36)\)\. Scaling the codebook is the standard lever for improving VQ\-VAE reconstruction, and a sustained line of work has improved codebook utilization through techniques such as EMA updates, dead\-code respawn, and codebook reparameterization\(Łańcuckiet al\.,[2020](https://arxiv.org/html/2605.06870#bib.bib12); Zheng and Vedaldi,[2023](https://arxiv.org/html/2605.06870#bib.bib15); Zhuet al\.,[2024](https://arxiv.org/html/2605.06870#bib.bib16),[2025](https://arxiv.org/html/2605.06870#bib.bib18)\)\. It is thus quite surprising that large\-scale VQ\-VAEs use only a small fraction of allocated latent dimensions: we find that WavTokenizer on LibriTTS, with a 512\-dimensional pre\-quantization latent, uses only 4 effective dimensions at codebook sizeK=212K=2^\{12\}despite 100% codebook utilization, or asZhanget al\.\([2025](https://arxiv.org/html/2605.06870#bib.bib20)\)shows, VQGAN on ImageNet collapses to 2–5 effective dimensions across codebook sizes fromK=210K=2^\{10\}toK=214K=2^\{14\}\. This dimensional collapse creates an irreducible loss floor that does not improve with codebook scaling\. The bottleneck is not the codebook size, but the effective dimension of the latent subspace\.

To theoretically analyze the mechanism behind collapse, we build on the seminal work ofSaxeet al\.\([2014](https://arxiv.org/html/2605.06870#bib.bib27)\); Refinetti and Goldt \([2022](https://arxiv.org/html/2605.06870#bib.bib28)\), which demonstrated how plain autoencoders learn the principal modes of their data sequentially, in decreasing order of variance\. Unfortunately, extending their idea to VQ\-VAE is difficult, since it requires analyzing coupled dynamics of theKKcodebook vectors and encoder\-decoder matrices, complicated by the non\-differentiable nearest\-neighbor quantizer\. Our contribution is the introduction of a tractable model,RD\-AE flow, which replaces VQ with a stochastic noising operation, chosen to minimize mean\-squared error subject to mutual information upper bound oflog2⁡K\\log\_\{2\}K\. We show that under the optimal coupling, latent directions corresponding to low\-variance modes may fall below a so\-calledwater leveland stop evolving due to interaction between the commitment loss term and the STE, see Fig\.[2](https://arxiv.org/html/2605.06870#S2.F2), thereby collapsing the latent dimension\.

This model shows that the cause of dimensional collapse is the timing of quantization: latent directions that were not yet properly learned when the quantizer turns on are deactivated by the rate constraint and never recover\. Our proposed fix trains the encoder and decoder as a plain autoencoder until enough directions are represented by the encoder before introducing VQ; we call this*AE warm\-up*\. Our theory predicts dependence of the latent effective dimension on the length of the warm\-up, and is corroborated by experiments on VQGAN, see Figure[5](https://arxiv.org/html/2605.06870#S3.F5)\. AE Warm\-up raises the VQGAN codebook dimension from 3 to 16 atK=214K=2^\{14\}and the WavTokenizer codebook dimension from 4 to 21 atK=216K=2^\{16\}, breaking the loss floor that cold\-start training hits regardless of codebook size and restoring expected scaling withKK\. Increasing warm\-up steps up to the predicted stopping point both increases final codebook dimension and reduces final loss, with gains scaling with target codebook size\.

![Refer to caption](https://arxiv.org/html/2605.06870v1/figures/headline_figure5.png)Figure 1:Dimensional collapse in VQ\-VAEs and the warm\-up fix\.During plain autoencoder training \(gray\), the latent effective dimensiondeffd\_\{\\text\{eff\}\}grows as directions of data variance \(“modes"\) are learned sequentially\. Turning on the quantizer too early collapsesdeffd\_\{\\text\{eff\}\}; reconstruction loss hits a floor that is independent of codebook size\|C\|\|C\|\(blue\)\. Training the encoder–decoder as a plain autoencoder before introducing the quantizer \(AE Warm\-up\) allows the encoder to learn lower\-variance directions first, preserving higherdeffd\_\{\\text\{eff\}\}after introducing VQ and restoring the expected scaling of reconstruction quality with\|C\|\|C\|\(yellow\-red\)\. We give an information\-theoretic account of this mechanism \([Section˜3](https://arxiv.org/html/2605.06870#S3)\) and validate it in both image \(VQGAN trained with ImageNet\-100,[Section˜4\.1](https://arxiv.org/html/2605.06870#S4.SS1)\) and audio \(WavTokenizer trained with LibriTTS,[Section˜4\.2](https://arxiv.org/html/2605.06870#S4.SS2)\) VQ\-VAEs\.Contributions\.

- •\(practice\) We demonstrate dimensionality collapse in popular VQ\-VAEs for audio and image tokenization \(WavTokenizer, VQGAN\) and show that the resulting performance degradation cannot be fixed by increasing codebook sizeKK\(“loss floor”\)\.
- •\(practice\) We propose a simple fix,*AE Warm\-up*, in which the bottleneck between encoder and decoder remains unquantized during the initial training phase\.
- •\(theory\) We introduce*RD\-AE*, a toy model for studying VQ\-VAE training dynamics that extends the well\-known sequential learning analysis ofSaxeet al\.\([2014](https://arxiv.org/html/2605.06870#bib.bib27)\)to the VQ\-VAE setting\.
- •\(theory\+practice\) We derive a tight bound on the latent dimensionality and reconstruction loss as a function of warm\-up duration\. This allows us to adaptively detect the optimal warm\-up duration\.
- •\(practice\) Our approach improves VQGAN \(ImageNet\-100\) and WavTokenizer \(LibriTTS\), where AE Warm\-up raises codebook effective dimension by5×5\\timesand restores loss scaling withKK\.

Prior work\.The warm\-up approach has been previously proposed in various ways, although without similar justification or connection to latent dimension\.Łańcuckiet al\.\([2020](https://arxiv.org/html/2605.06870#bib.bib12)\)delay quantization as a stability measure within their codebook\-utilization method,Zhaoet al\.\([2024](https://arxiv.org/html/2605.06870#bib.bib17)\)report empirical gains from AE pretraining followed by VQ fine\-tuning,Wanget al\.\([2025](https://arxiv.org/html/2605.06870#bib.bib39)\)take this to its limit by applying post\-hoc quantization to a frozen pretrained continuous VAE without any VQ training, and \(concurrently with this work\)Zhaoet al\.\([2026](https://arxiv.org/html/2605.06870#bib.bib21)\)propose the warm\-up under the name*deferred quantization*, though not connecting it to the latent dimension collapse\.

## 2Background

![Refer to caption](https://arxiv.org/html/2605.06870v1/figures/fig_bottom.png)Figure 2:Sequential learning in AE and VQ\-VAE\.A 2\-layer linear AE learns latent modes sequentiallySaxeet al\.\([2014](https://arxiv.org/html/2605.06870#bib.bib27)\)\(left\)\. In VQ\-VAEs, however, quantizing the bottleneck freezes lower modes \(those under the water levelD∗D^\{\*\}\), thus severely constraining effective dimension \(right\)\. See Section[3](https://arxiv.org/html/2605.06870#S3)for details\.Autoencoders and vector quantization\.An autoencoder pairs an encoderfθ:𝒳→ℝkf\_\{\\theta\}:\\mathcal\{X\}\\to\\mathbb\{R\}^\{k\}with a decodergϕ:ℝk→𝒳g\_\{\\phi\}:\\mathbb\{R\}^\{k\}\\to\\mathcal\{X\}and minimizes the reconstruction loss‖x−gϕ​\(fθ​\(x\)\)‖2\\\|x\-g\_\{\\phi\}\(f\_\{\\theta\}\(x\)\)\\\|^\{2\}\. A vector\-quantized autoencoder\(van den Oordet al\.,[2017](https://arxiv.org/html/2605.06870#bib.bib3)\)inserts a discretization step between the two: a finite codebook𝒞=\{c1,…,cK\}⊂ℝk\\mathcal\{C\}=\\\{c\_\{1\},\\dots,c\_\{K\}\\\}\\subset\\mathbb\{R\}^\{k\}and a nearest\-neighbor quantizerQ​\(z\)=arg⁡minc∈𝒞⁡‖z−c‖Q\(z\)=\\arg\\min\_\{c\\in\\mathcal\{C\}\}\\\|z\-c\\\|\. The encoder produces a continuous pre\-quantization latentz=fθ​\(x\)z=f\_\{\\theta\}\(x\); the quantizer replaces it with the nearest codewordzq=Q​\(z\)z\_\{q\}=Q\(z\); the decoder reconstructsx^=gϕ​\(zq\)\\hat\{x\}=g\_\{\\phi\}\(z\_\{q\}\)\. Training minimizes

ℒ​\(ϕ,θ\)=𝔼​\[‖x−x^‖2\+β​‖z−sg​\(zq\)‖2\],\\mathcal\{L\}\(\\phi,\\theta\)\\;=\\;\\mathbb\{E\}\[\\\|x\-\\hat\{x\}\\\|^\{2\}\\;\+\\;\\beta\\,\\\|z\-\\mathrm\{sg\}\(z\_\{q\}\)\\\|^\{2\}\],\(1\)whereβ\\betais the commitment weight andsg\\mathrm\{sg\}is the stop\-gradient and expectation is over the batch \(practice\) or population \(theory\)\. The encoder receives its gradient through the straight\-through estimator\(van den Oordet al\.,[2017](https://arxiv.org/html/2605.06870#bib.bib3)\),∂zq/∂z:=I\\partial z\_\{q\}/\\partial z:=I\. Codebook entries are typically updated by exponential moving averages of assigned latents\.

Evaluation metrics\.VQ\-VAEs are evaluated on reconstruction quality and on the health of the learned codebook\. Reconstruction is measured directly via pixel\- or sample\-level losses \(L1, L2, mel\-spectrogram\), and perceptually via task\-specific metrics — LPIPS\(Zhanget al\.,[2018](https://arxiv.org/html/2605.06870#bib.bib31)\)and rFID for images, PESQ\(Rixet al\.,[2001](https://arxiv.org/html/2605.06870#bib.bib32)\)and STOI\(Taalet al\.,[2010](https://arxiv.org/html/2605.06870#bib.bib33)\)for audio\. Modern VQ\-VAEs additionally include an adversarial discriminator loss\(Esseret al\.,[2021](https://arxiv.org/html/2605.06870#bib.bib5); Jiet al\.,[2025](https://arxiv.org/html/2605.06870#bib.bib36)\)\. On the codebook side,*utilization*measures how many allocated codewords are actually in use\.

Even at full utilization, standard VQ\-VAE representations live in a low\-dimensional subspace \(typically 4–10 dimensions\), as measured bydeffd\_\{\\mathrm\{eff\}\}, the number of principal components needed to explain99%99\\%of the pre\-quantization latent variance:

deff=min⁡\{m:∑i≤mλi/∑iλi≥0\.99\},d\_\{\\mathrm\{eff\}\}=\\min\\Big\\\{m:\\textstyle\\sum\_\{i\\leq m\}\\lambda\_\{i\}\\big/\\sum\_\{i\}\\lambda\_\{i\}\\geq 0\.99\\Big\\\},\(2\)with\{λi\}\\\{\\lambda\_\{i\}\\\}as the PCA eigenvalues ofzzin decreasing order\. The same definition applies to the codebook entries themselves; we report bothdeff​\(z\)d\_\{\\mathrm\{eff\}\}\(z\)anddeff​\(𝒞\)d\_\{\\mathrm\{eff\}\}\(\\mathcal\{C\}\)in Section[4](https://arxiv.org/html/2605.06870#S4)\.

Autoencoder gradient flow\.Saxeet al\.\([2014](https://arxiv.org/html/2605.06870#bib.bib27)\)analyzed the gradient flow of the plain\-AE objectiveLAE​\(W1,W2\)=𝔼​‖x−W2​W1​x‖2L\_\{\\mathrm\{AE\}\}\(W\_\{1\},W\_\{2\}\)=\\mathbb\{E\}\\\|x\-W\_\{2\}W\_\{1\}x\\\|^\{2\}for the two\-layer linear networkX^=W2​W1​X\\hat\{X\}=W\_\{2\}W\_\{1\}X\. On the diagonal manifoldW1=diag⁡\(uj\)W\_\{1\}=\\operatorname\{diag\}\(u\_\{j\}\),W2=diag⁡\(vj\)W\_\{2\}=\\operatorname\{diag\}\(v\_\{j\}\), the loss decouples asLAE​\(u,v\)=∑jσj2​\(1−uj​vj\)2L\_\{\\mathrm\{AE\}\}\(u,v\)=\\sum\_\{j\}\\sigma\_\{j\}^\{2\}\(1\-u\_\{j\}v\_\{j\}\)^\{2\}, and the flow acts independently on each mode:

u˙j=2​σj2​vj​\(1−uj​vj\),v˙j=2​σj2​uj​\(1−uj​vj\)\.\\dot\{u\}\_\{j\}=2\\sigma\_\{j\}^\{2\}\\,v\_\{j\}\\,\(1\-u\_\{j\}v\_\{j\}\),\\qquad\\dot\{v\}\_\{j\}=2\\sigma\_\{j\}^\{2\}\\,u\_\{j\}\\,\(1\-u\_\{j\}v\_\{j\}\)\.\(3\)With balanced initializationuj​\(0\)=vj​\(0\)=su\_\{j\}\(0\)=v\_\{j\}\(0\)=\\sqrt\{s\}at small scales∈\(0,1\)s\\in\(0,1\), the activation for each moderj​\(t\):=uj​\(t\)​vj​\(t\)r\_\{j\}\(t\):=u\_\{j\}\(t\)v\_\{j\}\(t\)has the closed form

rj​\(t\)=\(1\+1−ss​e−4​σj2​t\)−1,r\_\{j\}\(t\)=\\Bigl\(1\+\\tfrac\{1\-s\}\{s\}\\,e^\{\-4\\sigma\_\{j\}^\{2\}t\}\\Bigr\)^\{\-1\},\(4\)so modes activate sequentially in decreasing order of variance \(Figure[2](https://arxiv.org/html/2605.06870#S2.F2), left\)\. Consequentlydeffd\_\{\\text\{eff\}\}dips as top modes saturate first, then recovers as the rest catch up \(Figure[4](https://arxiv.org/html/2605.06870#S3.F4), black\) — a pattern also visible in large\-scale VQGAN and WavTokenizer AEs trained with AdamW \(Figure[5](https://arxiv.org/html/2605.06870#S3.F5)and Figure[7](https://arxiv.org/html/2605.06870#A2.F7)\)\.

## 3Theoretical model of collapse

We propose to model VQ\-VAE training dynamics by replacing the nearest\-neighbor quantizer with the rate\-distortion\-optimal continuous channel at rateR=log2⁡KR=\\log\_\{2\}K, resulting in a model that we call*RD\-AE*\. This relaxation makes the joint dynamics of the encoder, decoder, and channel analytically tractable\. We study the framework in the classical linear\-Gaussian setting where AE training dynamics admit closed\-form analysis\(Saxeet al\.,[2014](https://arxiv.org/html/2605.06870#bib.bib27)\)\. Despite the simplifications, the resulting toy model exhibits dimensional collapse \(Figure[4](https://arxiv.org/html/2605.06870#S3.F4)\)\. Rigorous analysis predicts that the collapse results from the interplay of two mechanisms: sequential learning in autoencoders\(Saxeet al\.,[2014](https://arxiv.org/html/2605.06870#bib.bib27); Refinetti and Goldt,[2022](https://arxiv.org/html/2605.06870#bib.bib28)\)and a fundamental property of MSE\-trained vector\-quantizers to collapse low\-variance dimensions \(which we model here information\-theoretically via rate\-distortion bottleneck\)\.

This section is organized as follows\. RD\-AE is introduced in Section[3\.1](https://arxiv.org/html/2605.06870#S3.SS1), which also simulates the dense case of the model\. The dense case is still difficult to analyze, and thus in Section[3\.2](https://arxiv.org/html/2605.06870#S3.SS2)we further simplify it to case of diagonal covariance matrices and weight initialization, and present the collapse mechanism in detail\. Section[3\.3](https://arxiv.org/html/2605.06870#S3.SS3)explains how to predict effective dimension of the resulting VQ\-VAE from the state of the AE Warm\-up checkpoint used\. Subsequent Section[4](https://arxiv.org/html/2605.06870#S4)validates that the behavior of non\-linear VQ\-VAEs trained on real data is in qualitative agreement with the behavior of our toy model\.

### 3\.1Setup

xxW1​xW\_\{1\}xzzQ​\(z\)=arg⁡minI​\(Q​\(z\),z\)≤R⁡𝔼​‖Q​\(z\)−z‖2Q\(z\)=\\arg\\min\_\{I\(Q\(z\),z\)\\leq R\}\\mathbb\{E\}\\\|Q\(z\)\-z\\\|^\{2\}zqz\_\{q\}W2​zqW\_\{2\}z\_\{q\}x^\\hat\{x\}STE backward:∂zq/∂z:=I\\partial z\_\{q\}/\\partial z:=IFigure 3:The linear RD\-AE\.The linear encoderW1W\_\{1\}maps data to the pre\-quantization latentzz; the \(potentially stochastic\) quantizerQQassignszzto codezqz\_\{q\}; the linear decoderW2W\_\{2\}reconstructsx^\\hat\{x\}\. The backward pass through the quantizer uses the straight\-through estimator\.Architecture\.The data isx∼𝒩​\(0,Σ\)x\\sim\\mathcal\{N\}\(0,\\Sigma\)onℝd\\mathbb\{R\}^\{d\}withΣ=diag⁡\(σ12,…,σd2\)\\Sigma=\\operatorname\{diag\}\(\\sigma\_\{1\}^\{2\},\\ldots,\\sigma\_\{d\}^\{2\}\)andσ12≥⋯≥σd2\>0\\sigma\_\{1\}^\{2\}\\geq\\cdots\\geq\\sigma\_\{d\}^\{2\}\>0\. A linear encoderW1∈ℝd×dW\_\{1\}\\in\\mathbb\{R\}^\{d\\times d\}produces the pre\-quantization latentz=W1​xz=W\_\{1\}x; a \(potentially stochastic\) quantizerQQproduceszq:=Q​\(z\)z\_\{q\}:=Q\(z\); a linear decoderW2∈ℝd×dW\_\{2\}\\in\\mathbb\{R\}^\{d\\times d\}returnsx^=W2​zq\\hat\{x\}=W\_\{2\}z\_\{q\}\(Figure[3](https://arxiv.org/html/2605.06870#S3.F3)\)\. Training objective is the standard VQ\-VAE loss \([1](https://arxiv.org/html/2605.06870#S2.E1)\)\.

The learning dynamics follow standard VQ\-VAE conventions: the encoder reconstruction “gradient” uses the straight\-through estimator∂zq/∂z:=I\\partial z\_\{q\}/\\partial z:=I, and the commitment term treatszqz\_\{q\}as constant via stop\-gradient\. As a result, the update equations forW1W\_\{1\}andW2W\_\{2\}are not the gradient of \([1](https://arxiv.org/html/2605.06870#S2.E1)\)\.

Information\-theoretic relaxation\.Analyzing VQ’s effect on the training dynamics requires jointly trackingKKhigh\-dimensional codewords alongside the encoder and decoder\. The codebook also breaks the diagonal structure exploited bySaxeet al\.\([2014](https://arxiv.org/html/2605.06870#bib.bib27)\), preventing a per\-mode decomposition\. Our key insight is that we can replace the dynamically evolving VQ with its information\-theoretic optimal counterpart\. Specifically, we assume that the mapz↦zqz\\mapsto z\_\{q\}at every instant of training is the one that minimizes MSE subject to the information bottleneck constraint:

minPzq\|z⁡\{𝔼​\[‖z−zq‖2\]:I​\(z;zq\)≤log2⁡K\}≜D​\(R\),\\min\_\{P\_\{z\_\{q\}\|z\}\}\\left\\\{\\mathbb\{E\}\[\\\|z\-z\_\{q\}\\\|^\{2\}\]:I\(z;z\_\{q\}\)\\leq\\log\_\{2\}K\\right\\\}\\,\\triangleq D\(R\),\(5\)where we denote the minimum MSE as the optimal rate\-RRdistortion in*latent space*,D​\(R\)D\(R\)\. This can be understood as relaxing the hard constraint\|supp​zq\|≤K\|\\mathrm\{supp\}z\_\{q\}\|\\leq Kto a softer mutual information constraint\. Although we optimistically assume that quantizer is trained fast enough to remain optimal for the evolving statistics ofzz, RD\-AE nevertheless exhibits dimensional collapse\. Due to the Gaussianity ofzz, the solution of \([5](https://arxiv.org/html/2605.06870#S3.E5)\) is given by a classical waterfilling formula\(Polyanskiy and Wu,[2025](https://arxiv.org/html/2605.06870#bib.bib2), Ex\. I\.9\)and conveniently diagonalizes when the initialization andΣ\\Sigmadoes\. Thus we obtain a similarly simple and insightful dynamics as inSaxeet al\.\([2014](https://arxiv.org/html/2605.06870#bib.bib27)\)except now including the VQ part\.

RD\-AE flow \(dense case\)\.Training is governed by the system of ODEs onW1,W2W\_\{1\},W\_\{2\}\(derived in[Section˜A\.1](https://arxiv.org/html/2605.06870#A1.SS1)\):

W˙1\\displaystyle\\dot\{W\}\_\{1\}=−2​\(W2⊤​W2​M​W1−W2⊤\)​Σ−2​β​\(W1−M​W1\)​Σ,\\displaystyle=\-2\\big\(W\_\{2\}^\{\\top\}W\_\{2\}\\,M\\,W\_\{1\}\-W\_\{2\}^\{\\top\}\\big\)\\Sigma\-2\\beta\\big\(W\_\{1\}\-MW\_\{1\}\\big\)\\Sigma,\(6\)W˙2\\displaystyle\\dot\{W\}\_\{2\}=2​Σ​W1⊤​M−2​W2​U​diag⁡\(cj​λj\)​U⊤,\\displaystyle=2\\Sigma W\_\{1\}^\{\\top\}M\-2W\_\{2\}\\,U\\operatorname\{diag\}\(c\_\{j\}\\lambda\_\{j\}\)U^\{\\top\}\\,,\(7\)whereMMis a covariance of the noise in thez→zqz\\to z\_\{q\}channel, andcjc\_\{j\}andλj\\lambda\_\{j\}are the spectra ofMMandCov​\(z\)\\mathrm\{Cov\}\(z\), respectively \(see[Section˜A\.1](https://arxiv.org/html/2605.06870#A1.SS1)\)\. Importantly, whenK→∞K\\to\\inftythe channel matrixM→IM\\to I, and the commitment term vanishes\. In this limit the RD\-AE flow \([6](https://arxiv.org/html/2605.06870#S3.E6)\)–\([7](https://arxiv.org/html/2605.06870#S3.E7)\) reduces to the*gradient flow*of the AE objective in[Figure˜2](https://arxiv.org/html/2605.06870#S2.F2), though theKKrequired for this to happen is impractically large\.

Numerical simulation \(dense case\)\.Before deriving the dynamics formally, we preview the qualitative behavior\. Figure[4](https://arxiv.org/html/2605.06870#S3.F4)simulates theRD\-AE flowwithd=64d=64,σj2=j−1\\sigma\_\{j\}^\{2\}=j^\{\-1\},β=1\\beta=1, with i\.i\.d\. Gaussian initialization, compared with the plain\-AE limit\. The plain AE drivesLrec→0L\_\{\\mathrm\{rec\}\}\\to 0whiledeffd\_\{\\textrm\{eff\}\}dips and recovers\. The RD\-AE, by contrast, exhibits dimensional collapse at every rateR∈\{10,14,18\}R\\in\\\{10,14,18\\\}bits, todeff=3,4,5d\_\{\\textrm\{eff\}\}=3,4,5respectively, mirroring the behavior we see in practice\. The reconstruction loss plateaus strictly above the Shannon distortion floorD​\(R\)D\(R\)\.222AlthoughD​\(R\)D\(R\)is the minimum*latent*\-space MSE, it also lower\-bounds the data\-space MSE: by the data processing inequality,I​\(x;x^\)≤I​\(z;zq\)≤RI\(x;\\hat\{x\}\)\\leq I\(z;z\_\{q\}\)\\leq R, sox^\\hat\{x\}is a rate\-RRdescription ofxxand must satisfy𝔼​‖x−x^‖2≥D​\(R\)\\mathbb\{E\}\\\|x\-\\hat\{x\}\\\|^\{2\}\\geq D\(R\)\.This gap is the cost of dimensionality collapse\.

![Refer to caption](https://arxiv.org/html/2605.06870v1/figures/fig_theory_main.png)Figure 4:Plain AE vs\. RD\-AE flow\.d=64d=64,Σ=diag⁡\(j−1\)\\Sigma=\\operatorname\{diag\}\(j^\{\-1\}\),β=1\\beta=1\.W1,W2W\_\{1\},W\_\{2\}ared×dd\\times dmatrices with i\.i\.d\. entries\(Wi\)j​k∼𝒩​\(0,s2/\(4​d\)\)\(W\_\{i\}\)\_\{jk\}\\sim\\mathcal\{N\}\(0,s^\{2\}/\(4d\)\),s=0\.01s=0\.01, evolved under the RD\-AE flow \([6](https://arxiv.org/html/2605.06870#S3.E6)\)–\([7](https://arxiv.org/html/2605.06870#S3.E7)\)\. Each curve is the median over 128 seeds; cross\-seed variance is below10−310^\{\-3\}\.\(Left\)Effective dimensiondeffd\_\{\\mathrm\{eff\}\}\. Plain AE \(black\) dips and recovers; RD\-AE \(colored by rateRR\) dips and plateaus\.\(Right\)Reconstruction loss\. Dashed lines are Shannon floorsD​\(R\)D\(R\), achieved by optimal enc\-dec pair\.
### 3\.2Analysis of the RD\-AE dynamics in the diagonal case

To get more insight into the dynamics, we restrict attention to diagonalW1,W2W\_\{1\},W\_\{2\}:

W1\(t\)=:diag\(uj\(t\)\),W2\(t\)=:diag\(vj\(t\)\),Σ=:diag\(λj\)\.W\_\{1\}\(t\)=:\\operatorname\{diag\}\(u\_\{j\}\(t\)\),\\quad W\_\{2\}\(t\)=:\\operatorname\{diag\}\(v\_\{j\}\(t\)\),\\quad\\Sigma=:\\operatorname\{diag\}\(\\lambda\_\{j\}\)\\,\.It is clear that in this case equations \([6](https://arxiv.org/html/2605.06870#S3.E6)\)\-\([7](https://arxiv.org/html/2605.06870#S3.E7)\) decouple and diagonality at initialization is preserved throughout training\. Next, we derive equations foruj,vju\_\{j\},v\_\{j\}and characterize the dimensional collapse mechanism via: \(1\) a per\-mode flow that reduces to aSaxeet al\.\([2014](https://arxiv.org/html/2605.06870#bib.bib27)\)logistic with a channel\-dependent timescale, \(2\) a dynamical water level whose rise pushes active modes below threshold and freezes them, and \(3\) a closed\-form expression for the resulting loss plateau\. The restriction to the diagonal case is only done for analytic convenience; this simplification still expresses all the major phenomena in the dense dynamics ofW1,W2W\_\{1\},W\_\{2\}\(cf\. \([6](https://arxiv.org/html/2605.06870#S3.E6)\)\-\([7](https://arxiv.org/html/2605.06870#S3.E7)\)\) that is simulated[Figure˜2](https://arxiv.org/html/2605.06870#S2.F2)and[Figure˜4](https://arxiv.org/html/2605.06870#S3.F4)\.

The rate\-distortion channel\.SinceW1W\_\{1\}is diagonal,zj=uj​xjz\_\{j\}=u\_\{j\}x\_\{j\}has independent coordinateszj∼𝒩​\(0,λj\)z\_\{j\}\\sim\\mathcal\{N\}\(0,\\lambda\_\{j\}\)withλj=uj2​σj2\\lambda\_\{j\}=u\_\{j\}^\{2\}\\sigma\_\{j\}^\{2\}\. The rate\-RRchannelPzq∣zP\_\{z\_\{q\}\\mid z\}minimizing𝔼​‖z−zq‖2\\mathbb\{E\}\\\|z\-z\_\{q\}\\\|^\{2\}under the information constraint factorizes across coordinates\(Polyanskiy and Wu,[2025](https://arxiv.org/html/2605.06870#bib.bib2), Sec\. 20\.4, Ch\. 26\)and is given by reverse water\-filling: a unique*water level*D⋆=D⋆​\(λ1,…,λd;R\)D^\{\\star\}=D^\{\\star\}\(\\lambda\_\{1\},\\ldots,\\lambda\_\{d\};R\), representing the variance threshold below which modes are dropped from the channel, solves

∑j12​log2\+⁡\(λj/D⋆\)=R,\\textstyle\\sum\_\{j\}\\tfrac\{1\}\{2\}\\log\_\{2\}^\{\+\}\(\\lambda\_\{j\}/D^\{\\star\}\)\\;=\\;R,\(8\)and the channel acts coordinatewise aszq,j∣zj∼𝒩​\(cj​zj,τj2\)z\_\{q,j\}\\mid z\_\{j\}\\sim\\mathcal\{N\}\(c\_\{j\}z\_\{j\},\\tau\_\{j\}^\{2\}\)with per\-mode distortionDjD\_\{j\}, channel gaincjc\_\{j\}, and noise varianceτj2\\tau\_\{j\}^\{2\}given by

Dj=min⁡\(λj,D⋆\),cj=1−Dj/λj,τj2=cj​D⋆\.D\_\{j\}\\;=\\;\\min\(\\lambda\_\{j\},D^\{\\star\}\),\\qquad c\_\{j\}\\;=\\;1\-D\_\{j\}/\\lambda\_\{j\},\\qquad\\tau\_\{j\}^\{2\}\\;=\\;c\_\{j\}D^\{\\star\}\.\(9\)Coordinates withλj≤D⋆\\lambda\_\{j\}\\leq D^\{\\star\}satisfycj=τj2=0c\_\{j\}=\\tau\_\{j\}^\{2\}=0and are passed through as zero; we call these*inactive*and the rest*active*, and write the active set𝒜​\(t\)=\{j:λj​\(t\)\>D⋆​\(t\)\}\\mathcal\{A\}\(t\)=\\\{j:\\lambda\_\{j\}\(t\)\>D^\{\\star\}\(t\)\\\}withk​\(t\)=\|𝒜​\(t\)\|k\(t\)=\|\\mathcal\{A\}\(t\)\|\. The channel decomposes aszj=zq,j\+njz\_\{j\}=z\_\{q,j\}\+n\_\{j\}jointly, withnj∼𝒩​\(0,Dj\)n\_\{j\}\\sim\\mathcal\{N\}\(0,D\_\{j\}\)independent ofzq,jz\_\{q,j\}, so that the total latent distortion𝔼​‖z−zq‖2=∑jDj\\mathbb\{E\}\\\|z\-z\_\{q\}\\\|^\{2\}=\\sum\_\{j\}D\_\{j\}equalsD​\(R\)D\(R\)from \([5](https://arxiv.org/html/2605.06870#S3.E5)\)\.

Diagonal RD\-AE flow\.Differentiating the full training objective \([1](https://arxiv.org/html/2605.06870#S2.E1)\) requires care\. We begin with the reconstruction lossLrec=𝔼​‖x−W2​zq‖2L\_\{\\mathrm\{rec\}\}=\\mathbb\{E\}\\\|x\-W\_\{2\}z\_\{q\}\\\|^\{2\}, which decomposes per mode using the joint decompositionz=zq\+nz=z\_\{q\}\+n, coupled across modes throughD⋆D^\{\\star\}:

Lrec​\(u,v\)=∑j=1d\[σj2​\(1−cj​uj​vj\)2\+vj2​τj2\],L\_\{\\mathrm\{rec\}\}\(u,v\)\\;=\\;\\sum\_\{j=1\}^\{d\}\\Bigl\[\\,\\sigma\_\{j\}^\{2\}\(1\-c\_\{j\}u\_\{j\}v\_\{j\}\)^\{2\}\\;\+\\;v\_\{j\}^\{2\}\\tau\_\{j\}^\{2\}\\,\\Bigr\],\(10\)Under these conditions, the decoder gradient is directly computed from \([10](https://arxiv.org/html/2605.06870#S3.E10)\)\. For the encoder, we apply the chain rule with the STE and SG conventions \(∂zq/∂z=:I\\partial z\_\{q\}/\\partial z=:Iand∂sg⁡\(x\)/∂x:=0\\partial\\operatorname\{sg\}\(x\)/\\partial x:=0\) per\-sample and take expectation\. We get \(full derivation in[Section˜A\.2](https://arxiv.org/html/2605.06870#A1.SS2)\):

∂Lrec∂vj=−2​cj​uj​σj2​\(1−uj​vj\),∂Lrec∂uj\|STE=−2​σj2​vj​\(1−cj​uj​vj\)\\frac\{\\partial L\_\{\\mathrm\{rec\}\}\}\{\\partial v\_\{j\}\}\\;=\\;\-2c\_\{j\}u\_\{j\}\\sigma\_\{j\}^\{2\}\(1\-u\_\{j\}v\_\{j\}\),\\qquad\\left\.\\frac\{\\partial L\_\{\\mathrm\{rec\}\}\}\{\\partial u\_\{j\}\}\\right\|\_\{\\mathrm\{STE\}\}\\;=\\;\-2\\sigma\_\{j\}^\{2\}v\_\{j\}\(1\-c\_\{j\}u\_\{j\}v\_\{j\}\)\(11\)The encoder commitment gradient gives∂Lcom/∂uj=2​Dj/uj\\partial L\_\{\\mathrm\{com\}\}/\\partial u\_\{j\}=2D\_\{j\}/u\_\{j\}\. Assembling, we get the diagonal RD\-AE flow

u˙j=2​σj2​vj​\(1−cj​uj​vj\)−2​β​Djuj,v˙j=2​σj2​cj​uj​\(1−uj​vj\)\.\\dot\{u\}\_\{j\}\\;=\\;2\\sigma\_\{j\}^\{2\}v\_\{j\}\(1\-c\_\{j\}u\_\{j\}v\_\{j\}\)\\;\-\\;\\frac\{2\\beta D\_\{j\}\}\{u\_\{j\}\},\\qquad\\dot\{v\}\_\{j\}\\;=\\;2\\sigma\_\{j\}^\{2\}c\_\{j\}u\_\{j\}\(1\-u\_\{j\}v\_\{j\}\)\.\(12\)We note that due to the straight\-through estimator, the flow \([12](https://arxiv.org/html/2605.06870#S3.E12)\) isnot a gradientflow, since the symmetry condition∂u˙j/∂vj=∂v˙j/∂uj\\partial\\dot\{u\}\_\{j\}/\\partial v\_\{j\}=\\partial\\dot\{v\}\_\{j\}/\\partial u\_\{j\}fails whenevercj≠1c\_\{j\}\\neq 1\.

##### Mode dynamics at balanced initialization\.

We specialize toβ=1\\beta=1and balanced initializationuj​\(0\)=vj​\(0\)=su\_\{j\}\(0\)=v\_\{j\}\(0\)=\\sqrt\{s\}\. Balance is preserved throughout training: forj∈𝒜j\\in\\mathcal\{A\},Dj=D⋆=λj​\(1−cj\)=uj2​σj2​\(1−cj\)D\_\{j\}=D^\{\\star\}=\\lambda\_\{j\}\(1\-c\_\{j\}\)=u\_\{j\}^\{2\}\\sigma\_\{j\}^\{2\}\(1\-c\_\{j\}\), and atuj=vj=:aju\_\{j\}=v\_\{j\}=:a\_\{j\}, we haveu˙j−v˙j=0\\dot\{u\}\_\{j\}\-\\dot\{v\}\_\{j\}=0\.

Henceuj​\(t\)=vj​\(t\)u\_\{j\}\(t\)=v\_\{j\}\(t\)for allttwhere modejjis active\. Then,rj:=uj​vj=uj2r\_\{j\}:=u\_\{j\}v\_\{j\}=u\_\{j\}^\{2\}obeys

r˙j=4​σj2​cj​\(t\)​rj​\(1−rj\)for​j∈𝒜,r˙j=0for​j∉𝒜,\\dot\{r\}\_\{j\}\\;=\\;4\\sigma\_\{j\}^\{2\}\\,c\_\{j\}\(t\)\\,r\_\{j\}\(1\-r\_\{j\}\)\\quad\\text\{for \}j\\in\\mathcal\{A\},\\qquad\\dot\{r\}\_\{j\}\\;=\\;0\\quad\\text\{for \}j\\notin\\mathcal\{A\},\(13\)derived in Appendix[A\.4](https://arxiv.org/html/2605.06870#A1.SS4)\. This is structurally similar to logistic \([4](https://arxiv.org/html/2605.06870#S2.E4)\) withcj​\(t\)c\_\{j\}\(t\)acting as a time\-dependent scale factor: atcj​\(t\)≡1c\_\{j\}\(t\)\\equiv 1\(theK→∞K\\to\\inftylimit\), \([13](https://arxiv.org/html/2605.06870#S3.E13)\) reduces to \([3](https://arxiv.org/html/2605.06870#S2.E3)\)\. Forj∉𝒜j\\notin\\mathcal\{A\},cj=0c\_\{j\}=0andDj=λjD\_\{j\}=\\lambda\_\{j\}giveu˙j=v˙j=0\\dot\{u\}\_\{j\}=\\dot\{v\}\_\{j\}=0on the balanced manifold\.

Collapse mechanism\.Implicit differentiation of \([8](https://arxiv.org/html/2605.06870#S3.E8)\) yields the water level dynamics:

D˙⋆D⋆=2k​∑i∈𝒜u˙iui,\\frac\{\\dot\{D\}^\{\\star\}\}\{D^\{\\star\}\}\\;=\\;\\frac\{2\}\{k\}\\sum\_\{i\\in\\mathcal\{A\}\}\\frac\{\\dot\{u\}\_\{i\}\}\{u\_\{i\}\},\(14\)derived in Appendix[A\.3](https://arxiv.org/html/2605.06870#A1.SS3)\. As active modes grow,D⋆D^\{\\star\}rises through \([14](https://arxiv.org/html/2605.06870#S3.E14)\), and modes whoseλj\\lambda\_\{j\}falls below the water line drop out of𝒜\\mathcal\{A\}and freeze permanently\. The active set is therefore non\-increasing, unlike in the plain AE, where every mode eventually activates\.*This is the dimensional collapse phenomenon*: weak modes that have not yet activated whenD⋆D^\{\\star\}overtakes them will remain unrepresented in the latent space forever\.

Loss plateau\.At convergence, withk∞=\|𝒜∞\|k\_\{\\infty\}=\|\\mathcal\{A\}\_\{\\infty\}\|active modes surviving, the reconstruction loss is

Lrec∞=k∞​D∞⋆\+∑j∉𝒜∞σj2,L\_\{\\mathrm\{rec\}\}^\{\\infty\}\\;=\\;k\_\{\\infty\}D^\{\\star\}\_\{\\infty\}\\;\+\\;\\sum\_\{j\\notin\\mathcal\{A\}\_\{\\infty\}\}\\sigma\_\{j\}^\{2\},\(15\)where the rate budgetRRis spent only on𝒜∞\\mathcal\{A\}\_\{\\infty\}, leaving frozen modes unrepresented \(paying their full varianceσj2\\sigma\_\{j\}^\{2\}\)\. The plateau therefore exceeds the Shannon optimumD​\(R\)D\(R\)whenever𝒜∞\\mathcal\{A\}\_\{\\infty\}is a strict subset of the optimal active set \(Figure[4](https://arxiv.org/html/2605.06870#S3.F4)\)\.

Special case ofβ=0\\beta=0\.Removing the commitment term by settingβ=0\\beta=0leads to very different dynamics: \([12](https://arxiv.org/html/2605.06870#S3.E12)\) givesu˙j=2​σj2​vj\>0\\dot\{u\}\_\{j\}=2\\sigma\_\{j\}^\{2\}v\_\{j\}\>0on inactive modes, so RD\-AE always eventually recovers from dimensional collapse atβ=0\\beta=0\. This does not, however, justify dropping the commitment loss in practice: in VQGAN, we find thatβ=0\\beta=0allows dimension recovery when trained with dead code reset, but achieves worse final loss \(Appendix[C\.2](https://arxiv.org/html/2605.06870#A3.SS2)\)\. We attribute this discrepancy to two aspects that RD\-AE elides: the finite codebook adaptation speed \(causing staleness between codewords and statistics ofzz\) and weight\-decay that enables continued evolution \(shrinkage\) of inactive modes\.

### 3\.3Effect of AE Warm\-up

Section[3\.2](https://arxiv.org/html/2605.06870#S3.SS2)showed that dimensional collapse happens because the rate constraint imposes a rising water level that freezes modes falling below it\. A natural fix is*AE Warm\-up*: delay quantization until enough modes have activated\. This runs the plain\-AE flow \([3](https://arxiv.org/html/2605.06870#S2.E3)\) from balanced diagonal initializationW1​\(0\)=W2​\(0\)=ε​IW\_\{1\}\(0\)=W\_\{2\}\(0\)=\\varepsilon Ifor warm\-up durationTwuT\_\{\\mathrm\{wu\}\}, then switch to the RD\-AE flow \([12](https://arxiv.org/html/2605.06870#S3.E12)\) with rateRRandβ=1\\beta=1\. The following theorem bounds how many modes survive the switch as a function ofTwuT\_\{\\mathrm\{wu\}\}\.

###### Theorem 1\.

Letx∼𝒩​\(0,Σ\)x\\sim\\mathcal\{N\}\(0,\\Sigma\)withΣ=diag⁡\(σ12,…,σd2\)\\Sigma=\\operatorname\{diag\}\(\\sigma\_\{1\}^\{2\},\\ldots,\\sigma\_\{d\}^\{2\}\)andσ12\>⋯\>σd2\>0\\sigma\_\{1\}^\{2\}\>\\cdots\>\\sigma\_\{d\}^\{2\}\>0\. Letgj​\(T\)=\(1\+1−ε2ε2​e−4​σj2​T\)−1g\_\{j\}\(T\)=\\big\(1\+\\tfrac\{1\-\\varepsilon^\{2\}\}\{\\varepsilon^\{2\}\}e^\{\-4\\sigma\_\{j\}^\{2\}T\}\\big\)^\{\-1\}denote the activation of modejjat timeTTin plain\-AE training \([4](https://arxiv.org/html/2605.06870#S2.E4)\)\. Define

mwu​\(Twu,R\):=max⁡\{m∈\{1,…,d\}:R\>12​∑i=1m−1log2⁡σi2​gi​\(Twu\)σm2​gm​\(Twu\)\}\.m^\{\\mathrm\{wu\}\}\(T\_\{\\mathrm\{wu\}\},R\)\\;:=\\;\\max\\\!\\Bigl\\\{m\\in\\\{1,\\ldots,d\\\}\\;:\\;R\>\\tfrac\{1\}\{2\}\\sum\_\{i=1\}^\{m\-1\}\\log\_\{2\}\\tfrac\{\\sigma\_\{i\}^\{2\}\\,g\_\{i\}\(T\_\{\\mathrm\{wu\}\}\)\}\{\\sigma\_\{m\}^\{2\}\\,g\_\{m\}\(T\_\{\\mathrm\{wu\}\}\)\}\\Bigr\\\}\.\(16\)After switching to RD\-AE flow atTwuT\_\{\\mathrm\{wu\}\}, at mostmwum^\{\\mathrm\{wu\}\}modes remain active at convergence, and

Lrec∞≥mwu​δmwu​\(R\)\+∑j\>mwuσj2,δm​\(R\):=\(∏i=1mσi2\)1/m​2−2​R/m\.L\_\{\\mathrm\{rec\}\}^\{\\infty\}\\;\\geq\\;m^\{\\mathrm\{wu\}\}\\,\\delta\_\{m^\{\\mathrm\{wu\}\}\}\(R\)\\;\+\\;\\sum\_\{j\>m^\{\\mathrm\{wu\}\}\}\\sigma\_\{j\}^\{2\},\\qquad\\delta\_\{m\}\(R\):=\\Bigl\(\\prod\_\{i=1\}^\{m\}\\sigma\_\{i\}^\{2\}\\Bigr\)^\{\\\!1/m\}2^\{\-2R/m\}\.\(17\)

Proof sketch\.At the switch time,λj​\(Twu\)=gj​\(Twu\)​σj2\\lambda\_\{j\}\(T\_\{\\mathrm\{wu\}\}\)=g\_\{j\}\(T\_\{\\mathrm\{wu\}\}\)\\sigma\_\{j\}^\{2\}, and the condition in \([16](https://arxiv.org/html/2605.06870#S3.E16)\) is the water\-filling requirement for the topmwum^\{\\mathrm\{wu\}\}modes to lie aboveD⋆​\(Twu\)D^\{\\star\}\(T\_\{\\mathrm\{wu\}\}\)\. Modes inactive atTwuT\_\{\\mathrm\{wu\}\}stay atrj=gj​\(Twu\)r\_\{j\}=g\_\{j\}\(T\_\{\\mathrm\{wu\}\}\), contributing varianceσj2\\sigma\_\{j\}^\{2\}to the reconstruction loss\. The bound onLrec∞L\_\{\\mathrm\{rec\}\}^\{\\infty\}follows from optimal water\-filling of rateRRover the topmwum^\{\\mathrm\{wu\}\}modes; full proof in Appendix[A\.6](https://arxiv.org/html/2605.06870#A1.SS6)\.

Empirical validation of the warm\-up theory\.Although Theorem[1](https://arxiv.org/html/2605.06870#Thmtheorem1)is stated for the diagonal RD\-AE model \(where it is nearly tight; see Appendix[A\.6\.1](https://arxiv.org/html/2605.06870#A1.SS6.SSS1)\), its core construction transfers directly to non\-linear VQ\-VAEs\. Indeed, one can snapshot a state of AE during the warm\-up stage, compute PCA of the covariance matrix of the latents, evaluate water\-filling rate allocation, and compute the number of surviving dimensions\. The resulting number turns out to give a rather accurate upper bound on the final latent dimension of the trained VQ\-VAE if the warm\-up is switched off at the time of the snapshot: see Figure[5](https://arxiv.org/html/2605.06870#S3.F5)\(right\) and Appendix[C](https://arxiv.org/html/2605.06870#A3)\)\. This allows to make on\-the\-fly automatic decision about when to switch from the AE \(warm\-up\) phase and start training the actual VQ\-VAE\.

![Refer to caption](https://arxiv.org/html/2605.06870v1/figures/cb16k_3panel_v3_brokenleft.png)Figure 5:VQGAN on ImageNet\-100 \(20k\),\|𝒞\|=214\|\\mathcal\{C\}\|=2^\{14\}\. Codebookdeffd\_\{\\text\{eff\}\}\(left\) and validation reconstruction loss \(center\) during training for AE warm\-up durationsTwu∈\{0,1​k,…,40​k\}T\_\{\\text\{wu\}\}\\in\\\{0,1\\text\{k\},\\dots,40\\text\{k\}\\\}; longer warm\-up raises codebook dimension and lowers the loss floor, with returns diminishing pastTwu≈20​kT\_\{\\text\{wu\}\}\\approx 20\\text\{k\}\.Right:reverse water\-filling on the AE checkpoint’s PCA spectrum at rateR=14R=14is an empirical upper bound on the final trained codebookdeffd\_\{\\text\{eff\}\}\(circles\)\. See Appendix[C](https://arxiv.org/html/2605.06870#A3)for\|𝒞\|∈\{210,216\}\|\\mathcal\{C\}\|\\in\\\{2^\{10\},2^\{16\}\\\}\.

## 4Escaping Dimensional Collapse with Autoencoder Warm\-up

[Section˜3](https://arxiv.org/html/2605.06870#S3)argues that VQ\-VAE reconstruction is bottlenecked by the latent effective dimension, not the codebook size: modes inactive when the quantizer turns on are frozen by the rate constraint and never recover\. The suggested fix,*AE Warm\-up*, trains the encoder–decoder as a plain autoencoder forTwuT\_\{\\text\{wu\}\}steps so that lower\-variance directions activate first, and then introduces VQ\. A single warm\-up checkpoint can be reused across codebook sizes, amortizing its cost over a sweep\.

ChoosingTwuT\_\{\\text\{wu\}\}\.The theory tells us that reverse water\-filling at rateR=log2⁡KR=\\log\_\{2\}Kon the PCA spectrum of the encoder’s latents upper\-bounds the codebookdeffd\_\{\\text\{eff\}\}that the downstream VQ\-VAE will achieve \([Figure˜5](https://arxiv.org/html/2605.06870#S3.F5), right\)\. In practice,TwuT\_\{\\text\{wu\}\}should be chosen by monitoringdeffd\_\{\\text\{eff\}\}during AE training and stopping once it plateaus — further AE training does not lift additional modes above the water level at the target rate\. For VQGAN on ImageNet\-100 this occurs around 20,000 steps \([Figure˜7\(a\)](https://arxiv.org/html/2605.06870#A2.F7.sf1)\)\.

Empirical validation\.We test the prediction at scale on both image and audio modalities, training VQGAN on ImageNet\-100 \([Section˜4\.1](https://arxiv.org/html/2605.06870#S4.SS1)\) and WavTokenizer on LibriTTS \([Section˜4\.2](https://arxiv.org/html/2605.06870#S4.SS2)\)\. Each setting addresses three questions:\(Q1\)does cold\-start VQ training collapse the latent to a low\-dimensional subspace in large\-scale architectures;\(Q2\)does the resulting loss floor persist asKKgrows; and\(Q3\)does AE warm\-up recover latent dimensionality and restore the expected scaling of reconstruction quality withKK\.

### 4\.1VQGAN on ImageNet\-100

Table 1:VQGAN on ImageNet\-100 \(20k\) warm\-up ablation\.L1, LPIPS, and rFID: lower is better\. D: codebookdeffd\_\{\\text\{eff\}\}\(higher is better\)\. Bold indicates the best value within each codebook size\. Vanilla VQGAN uses uniform codebook init and no respawn; VQGAN w/ Respawn adds k\-means init and respawn;Twu=T\_\{\\text\{wu\}\}=\{1k, …, 40k\} additionally pre\-trains the encoder–decoder as a plain autoencoder for the indicated number of steps before introducing the codebook\. All runs share the same 100k\-step budget, so warm\-up trades VQ steps for AE steps\.##### Setup\.

We train VQGAN\(Esseret al\.,[2021](https://arxiv.org/html/2605.06870#bib.bib5)\)on 20,000 images subsampled from ImageNet\-100 at128×128128\\times 128resolution, with L1 and LPIPS losses for 100k steps \(the discriminator loss inEsseret al\.\([2021](https://arxiv.org/html/2605.06870#bib.bib5)\)only activates after 250k steps\)\. We compare three methods across codebook sizesK∈\{210,214,216\}K\\in\\\{2^\{10\},2^\{14\},2^\{16\}\\\}:Vanilla VQGAN\(random codebook init, straight\-through updates\);VQGAN w/ Respawn, the standard modern VQ\-VAE training method with k\-means init, EMA codebook updates \(decay 0\.99\), and dead\-code respawn at usage threshold0\.010\.01\(Takidaet al\.,[2022](https://arxiv.org/html/2605.06870#bib.bib6); Mentzeret al\.,[2024](https://arxiv.org/html/2605.06870#bib.bib7); Dhariwalet al\.,[2020](https://arxiv.org/html/2605.06870#bib.bib8)\); andAE Warm\-up, which trains as a plain autoencoder forTwu∈\{1​k,5​k,10​k,15​k,20​k,30​k,40​k\}T\_\{\\text\{wu\}\}\\in\\\{1\\text\{k\},5\\text\{k\},10\\text\{k\},15\\text\{k\},20\\text\{k\},30\\text\{k\},40\\text\{k\}\\\}steps before resuming the Respawn protocol\. All variants share a 100k\-step budget\. Full training and evaluation details can be found in[Section˜B\.1](https://arxiv.org/html/2605.06870#A2.SS1)\.

##### Key findings\.

Table[1](https://arxiv.org/html/2605.06870#S4.T1)shows the final metrics\. Without warm\-up, neither Vanilla VQGAN nor the Respawn\-only baseline improves substantially asKKgrows\. AE Warm\-up restores scaling with codebook size and improves all metrics over both baselines: 8\.1–18\.6% over Vanilla VQGAN and 5\.3–11\.0% over Respawn on L1, 18\.0–31\.4% and 14\.5–23\.1% on LPIPS, and 30\.6–47\.5% and 17\.9–35\.4% on rFID\. Returns plateau beyondTwu=20​kT\_\{\\text\{wu\}\}=20\\text\{k\}\.

Vanilla VQGAN does not benefit from a larger codebook\.The three vanilla rows sit at the same L1 loss \(0\.153\) and the same codebook dimension \(deff=2d\_\{\\text\{eff\}\}=2\) despiteKKgrowing by64×64\\times, with utilization collapsing from 82\.7% atK=210K=2^\{10\}to 2\.2% atK=216K=2^\{16\}\.

Respawn partially rescues the floor but cannot recover zeroed\-out modes\.VQGAN w/ Respawn improves on Vanilla at everyKKbut falls short of AE warm\-up in all cases\. Its codebook dimension reaches only 3 atK∈\{210,214\}K\\in\\\{2^\{10\},2^\{14\}\\\}and 5 atK=216K=2^\{16\}— far below warm\-up values — consistent with respawn redistributing dead codes within the active subspace rather than reviving collapsed modes\.

Warm\-up lowers the floor\.IncreasingTwuT\_\{\\text\{wu\}\}from 1k to 40k steps monotonically raises codebook effective dimension \(from 4–5 to 17–19\) and improves LPIPS \(from 0\.31–0\.36 to 0\.28–0\.33\), with noisier but consistent gains on rFID\. As predicted by Theorem[1](https://arxiv.org/html/2605.06870#Thmtheorem1), each additional warm\-up step allows more modes to be represented by the encoder before the quantizer is introduced\.

### 4\.2WavTokenizer on LibriTTS

Table 2:WavTokenizer audio reconstruction performance on LibriTTS test\-clean/test\-other\.Each row pair compares cold\-start WavTokenizer training against AE Warm\-up at the same codebook size\.Boldindicates the winner within each codebook size;underlineindicates the overall best\. Without warm\-up, codebook utilization collapses asKKgrows and codebook effective dimension is pinned at44\. AE Warm\-up maintains near\-full utilization at everyKKand raises codebook effective dimension to1111–2121\. Additionally,DD\(codebookdeffd\_\{\\text\{eff\}\}\), Mel, PESQ, and STOI improve monotonically withKK\. Cold\-start atK=216K=2^\{16\}is omitted given the severity of collapse atK=214K=2^\{14\}\.Method\|𝒞\|\|\\mathcal\{C\}\|Util↑\\uparrowD↑D\\uparrowMel↓\\downarrowUTMOS↑\\uparrowPESQ↑\\uparrowSTOI↑\\uparrowV/UV F1↑\\uparrowWavTokenizer4096100\.0%40\.480/0\.5484\.051/3\.5632\.336/2\.0540\.912/0\.8780\.940/0\.916AE Warm\-up4096100\.0%110\.469/0\.5263\.980/3\.4842\.275/2\.0320\.907/0\.8730\.931/0\.905WavTokenizer819216\.9%40\.510/0\.5844\.009/3\.5472\.179/1\.9370\.902/0\.8670\.934/0\.908AE Warm\-up8192100\.0%170\.448/0\.5043\.991/3\.4802\.425/2\.1470\.915/0\.8820\.933/0\.907WavTokenizer163848\.4%40\.497/0\.5664\.007/3\.5392\.209/1\.9770\.904/0\.8700\.936/0\.912AE Warm\-up16384100\.0%190\.432/0\.4864\.041/3\.5362\.514/2\.2290\.920/0\.8890\.942/0\.917AE Warm\-up6553699\.7%210\.428/0\.4804\.050/3\.5212\.626/2\.3280\.925/0\.8950\.940/0\.915##### Setup\.

To validate our theory in the audio domain and on a larger\-scale setting, we train WavTokenizer\(Jiet al\.,[2025](https://arxiv.org/html/2605.06870#bib.bib36)\)on LibriTTS\(Zenet al\.,[2019](https://arxiv.org/html/2605.06870#bib.bib38)\)with full reconstruction and adversarial losses\. We use four codebook sizes \(K∈\{4096,8192,16384,65536\}K\\in\\\{4096,8192,16384,65536\\\}\) and a fixed training budget of 100 epochs in two regimes: \(i\) cold\-start WavTokenizer training \(the baseline\) and \(ii\) AE Warm\-up for 43 epochs before initializing the codebook with k\-means and training as a VQ\-VAE with dead code respawn\. All other hyperparameters match the original WavTokenizer; details are in[Section˜B\.2](https://arxiv.org/html/2605.06870#A2.SS2)\.

##### Key findings\.

Table[2](https://arxiv.org/html/2605.06870#S4.T2)confirms the same pattern as VQGAN\. Cold\-start training pins codebookdeffd\_\{\\text\{eff\}\}at44regardless ofKK—fewer than1%1\\%of the 512 available directions—and utilization collapses from100%100\\%atK=212K\{=\}2^\{12\}to8\.4%8\.4\\%atK=214K\{=\}2^\{14\}\. Performance does not merely plateau but*degrades*withKK: test\-clean PESQ drops from2\.342\.34atK=212K\{=\}2^\{12\}to2\.182\.18atK=213K\{=\}2^\{13\}\. The exception isK=212K\{=\}2^\{12\}, where a 4\-dimensional latent already tiles the codebook and warm\-up’s extra dimensions \(deff=11d\_\{\\text\{eff\}\}\{=\}11\) are surplus capacity—consistent with the theory\.

AE Warm\-up restores monotone scaling: codebookdeffd\_\{\\text\{eff\}\}rises from1111to2121asKKgrows from2122^\{12\}to2162^\{16\}, all at near\-full utilization\. The reconstruction metrics also improve: Mel loss drops from0\.4690\.469to0\.4280\.428on test\-clean, and PESQ improves from2\.2752\.275to2\.6262\.626\. We did not attempt cold\-start atK=216K\{=\}2^\{16\}given the severity of collapse already atK=214K\{=\}2^\{14\}\.

## 5Discussion and Limitations

The bottleneck in VQ\-VAE training is the span of the latent subspace the encoder delivers to the quantizer, not the number of codes\. AE warm\-up addresses this directly, and because it acts before quantization begins, it is agnostic to the downstream quantizer: residual VQ\(Zeghidouret al\.,[2021](https://arxiv.org/html/2605.06870#bib.bib13); Kumaret al\.,[2023](https://arxiv.org/html/2605.06870#bib.bib9); Zhanget al\.,[2024](https://arxiv.org/html/2605.06870#bib.bib10); Défossezet al\.,[2024](https://arxiv.org/html/2605.06870#bib.bib11)\), product quantization\(Zhanget al\.,[2025](https://arxiv.org/html/2605.06870#bib.bib20)\), and structural reparameterizations such as SimVQ\(Zhuet al\.,[2025](https://arxiv.org/html/2605.06870#bib.bib18)\)are all compatible\. Indeed, the mode\-suppression mechanism operates within each residual stage or product factor, so these methods stand to benefit from warm\-up as well\.

Scope and future work\.Our theory assumes linear encoders and decoders with independent Gaussian inputs — enough to expose the mechanism, but restrictive\. Extending to nonlinear encoders and predicting the loss floor quantitatively in realistic models are natural next steps\. The method also raises practical questions: how far can codebook size be pushed once warm\-up is in place, and which existing codebook\-geometry techniques continue to work in the higher\-deffd\_\{\\text\{eff\}\}regime warm\-up unlocks?

## References

- A\. Défossez, L\. Mazaré, M\. Orsini, A\. Royer, P\. Pérez, H\. Jégou, E\. Grave, and N\. Zeghidour \(2024\)Moshi: a speech\-text foundation model for real\-time dialogue\.arXiv preprint arXiv:2410\.00037\.Cited by:[§5](https://arxiv.org/html/2605.06870#S5.p1.1)\.
- P\. Dhariwal, H\. Jun, C\. Payne, J\. W\. Kim, A\. Radford, and I\. Sutskever \(2020\)Jukebox: a generative model for music\.arXiv preprint arXiv:2005\.00341\.Cited by:[§4\.1](https://arxiv.org/html/2605.06870#S4.SS1.SSS0.Px1.p1.4)\.
- P\. Esser, R\. Rombach, and B\. Ommer \(2021\)Taming transformers for high\-resolution image synthesis\.InProceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition \(CVPR\),pp\. 12873–12883\.Cited by:[§B\.1](https://arxiv.org/html/2605.06870#A2.SS1.SSS0.Px1.p1.7),[§1](https://arxiv.org/html/2605.06870#S1.p1.3),[§2](https://arxiv.org/html/2605.06870#S2.p2.1),[§4\.1](https://arxiv.org/html/2605.06870#S4.SS1.SSS0.Px1.p1.4)\.
- S\. Ji, Z\. Jiang, W\. Wang, Y\. Chen, M\. Fang, J\. Zuo, Q\. Yang, X\. Cheng, Z\. Wang, R\. Li, Z\. Zhang, X\. Yang, R\. Huang, Y\. Jiang, Q\. Chen, S\. Zheng, and Z\. Zhao \(2025\)WavTokenizer: an efficient acoustic discrete codec tokenizer for audio language modeling\.InInternational Conference on Learning Representations \(ICLR\),Note:arXiv:2408\.16532Cited by:[§B\.2](https://arxiv.org/html/2605.06870#A2.SS2.SSS0.Px1.p1.1),[§B\.2](https://arxiv.org/html/2605.06870#A2.SS2.SSS0.Px6.p1.1),[§1](https://arxiv.org/html/2605.06870#S1.p1.3),[§2](https://arxiv.org/html/2605.06870#S2.p2.1),[§4\.2](https://arxiv.org/html/2605.06870#S4.SS2.SSS0.Px1.p1.1)\.
- J\. W\. Kim, J\. Salamon, P\. Li, and J\. P\. Bello \(2018\)CREPE: a convolutional representation for pitch estimation\.InProceedings of the IEEE International Conference on Acoustics, Speech and Signal Processing \(ICASSP\),pp\. 161–165\.Cited by:[§B\.2](https://arxiv.org/html/2605.06870#A2.SS2.SSS0.Px6.p1.1)\.
- R\. Kumar, P\. Seetharaman, A\. Luebs, I\. Kumar, and K\. Kumar \(2023\)High\-fidelity audio compression with improved RVQGAN\.InAdvances in Neural Information Processing Systems \(NeurIPS\),Cited by:[§5](https://arxiv.org/html/2605.06870#S5.p1.1)\.
- A\. Łańcucki, J\. Chorowski, G\. Sanchez, R\. Marxer, N\. Chen, H\. J\. G\. A\. Dolfing, S\. Khurana, T\. Alumae, and A\. Laurent \(2020\)Robust training of vector quantized bottleneck models\.InInternational Joint Conference on Neural Networks \(IJCNN\),Cited by:[§1](https://arxiv.org/html/2605.06870#S1.p1.3),[§1](https://arxiv.org/html/2605.06870#S1.p5.1)\.
- F\. Mentzer, D\. Minnen, E\. Agustsson, and M\. Tschannen \(2024\)Finite scalar quantization: vq\-vae made simple\.InInternational Conference on Learning Representations \(ICLR\),Cited by:[§4\.1](https://arxiv.org/html/2605.06870#S4.SS1.SSS0.Px1.p1.4)\.
- M\. Morrison, R\. Kumar, K\. Kumar, P\. Seetharaman, A\. Courville, and Y\. Bengio \(2022\)Chunked autoregressive GAN for conditional waveform synthesis\.InInternational Conference on Learning Representations \(ICLR\),Cited by:[§B\.2](https://arxiv.org/html/2605.06870#A2.SS2.SSS0.Px6.p1.1)\.
- Y\. Polyanskiy and Y\. Wu \(2025\)Information theory: from coding to learning\.Cambridge university press\.Cited by:[§A\.3](https://arxiv.org/html/2605.06870#A1.SS3.p2.3),[§3\.1](https://arxiv.org/html/2605.06870#S3.SS1.p3.8),[§3\.2](https://arxiv.org/html/2605.06870#S3.SS2.p2.8)\.
- A\. Razavi, A\. van den Oord, and O\. Vinyals \(2019\)Generating diverse high\-fidelity images with VQ\-VAE\-2\.InAdvances in Neural Information Processing Systems \(NeurIPS\),Vol\.32\.Cited by:[§1](https://arxiv.org/html/2605.06870#S1.p1.3)\.
- M\. Refinetti and S\. Goldt \(2022\)The dynamics of representation learning in shallow, non\-linear autoencoders\.InProceedings of the 39th International Conference on Machine Learning,K\. Chaudhuri, S\. Jegelka, L\. Song, C\. Szepesvari, G\. Niu, and S\. Sabato \(Eds\.\),Proceedings of Machine Learning Research, Vol\.162,pp\. 18499–18519\.External Links:[Link](https://proceedings.mlr.press/v162/refinetti22a.html)Cited by:[§1](https://arxiv.org/html/2605.06870#S1.p2.2),[§3](https://arxiv.org/html/2605.06870#S3.p1.1)\.
- A\. W\. Rix, J\. G\. Beerends, M\. P\. Hollier, and A\. P\. Hekstra \(2001\)Perceptual evaluation of speech quality \(PESQ\)—a new method for speech quality assessment of telephone networks and codecs\.InProceedings of the IEEE International Conference on Acoustics, Speech, and Signal Processing \(ICASSP\),Vol\.2,pp\. 749–752\.Cited by:[§B\.2](https://arxiv.org/html/2605.06870#A2.SS2.SSS0.Px6.p1.1),[§2](https://arxiv.org/html/2605.06870#S2.p2.1)\.
- T\. Saeki, D\. Xin, W\. Nakata, T\. Koriyama, S\. Takamichi, and H\. Saruwatari \(2022\)UTMOS: UTokyo\-SaruLab system for VoiceMOS challenge 2022\.InProc\. Interspeech,pp\. 4521–4525\.Cited by:[§B\.2](https://arxiv.org/html/2605.06870#A2.SS2.SSS0.Px6.p1.1)\.
- A\. M\. Saxe, J\. L\. McClelland, and S\. Ganguli \(2014\)Exact solutions to the nonlinear dynamics of learning in deep linear neural networks\.InInternational Conference on Learning Representations \(ICLR\),Cited by:[§A\.5](https://arxiv.org/html/2605.06870#A1.SS5.p1.8),[3rd item](https://arxiv.org/html/2605.06870#S1.I1.i3.p1.1),[§1](https://arxiv.org/html/2605.06870#S1.p2.2),[Figure 2](https://arxiv.org/html/2605.06870#S2.F2),[Figure 2](https://arxiv.org/html/2605.06870#S2.F2.2.1.1),[§2](https://arxiv.org/html/2605.06870#S2.p4.5),[§3\.1](https://arxiv.org/html/2605.06870#S3.SS1.p3.2),[§3\.1](https://arxiv.org/html/2605.06870#S3.SS1.p3.8),[§3\.2](https://arxiv.org/html/2605.06870#S3.SS2.p1.3),[§3](https://arxiv.org/html/2605.06870#S3.p1.1)\.
- C\. H\. Taal, R\. C\. Hendriks, R\. Heusdens, and J\. Jensen \(2010\)A short\-time objective intelligibility measure for time\-frequency weighted noisy speech\.InProceedings of the IEEE International Conference on Acoustics, Speech, and Signal Processing \(ICASSP\),pp\. 4214–4217\.Cited by:[§B\.2](https://arxiv.org/html/2605.06870#A2.SS2.SSS0.Px6.p1.1),[§2](https://arxiv.org/html/2605.06870#S2.p2.1)\.
- Y\. Takida, M\. Okada, and N\. Murata \(2022\)SQ\-vae: variational bayes on discrete representation with self\-annealed stochastic quantization\.InInternational Conference on Machine Learning \(ICML\),Cited by:[§4\.1](https://arxiv.org/html/2605.06870#S4.SS1.SSS0.Px1.p1.4)\.
- Y\. Tian, D\. Krishnan, and P\. Isola \(2020\)Contrastive multiview coding\.InComputer Vision – ECCV 2020: 16th European Conference, Glasgow, UK, August 23–28, 2020, Proceedings, Part XI,Berlin, Heidelberg,pp\. 776–794\.External Links:ISBN 978\-3\-030\-58620\-1,[Link](https://doi.org/10.1007/978-3-030-58621-8_45),[Document](https://dx.doi.org/10.1007/978-3-030-58621-8%5F45)Cited by:[§B\.1](https://arxiv.org/html/2605.06870#A2.SS1.SSS0.Px4.p1.1)\.
- A\. van den Oord, O\. Vinyals, and K\. Kavukcuoglu \(2017\)Neural discrete representation learning\.InAdvances in Neural Information Processing Systems \(NeurIPS\),Vol\.30\.Cited by:[§1](https://arxiv.org/html/2605.06870#S1.p1.3),[§2](https://arxiv.org/html/2605.06870#S2.p1.11),[§2](https://arxiv.org/html/2605.06870#S2.p1.8)\.
- Y\. Wang, Z\. Lin, Y\. Teng, Y\. Zhu, S\. Ren, J\. Feng, and X\. Liu \(2025\)Bridging continuous and discrete tokens for autoregressive visual generation\.arXiv preprint arXiv:2503\.16430\.Cited by:[§1](https://arxiv.org/html/2605.06870#S1.p5.1)\.
- N\. Zeghidour, A\. Luebs, A\. Omber, M\. Tagliasacchi, and J\. V\. Borber \(2021\)SoundStream: an end\-to\-end neural audio codec\.IEEE/ACM Transactions on Audio, Speech, and Language Processing30,pp\. 495–507\.Cited by:[§1](https://arxiv.org/html/2605.06870#S1.p1.3),[§5](https://arxiv.org/html/2605.06870#S5.p1.1)\.
- H\. Zen, V\. Dang, R\. Clark, Y\. Zhang, R\. J\. Weiss, Y\. Jia, Z\. Chen, and Y\. Wu \(2019\)LibriTTS: a corpus derived from librispeech for text\-to\-speech\.InProc\. Interspeech,pp\. 1526–1530\.Cited by:[§4\.2](https://arxiv.org/html/2605.06870#S4.SS2.SSS0.Px1.p1.1)\.
- J\. Zhang, Y\. Shen, G\. Chen, L\. Song, and E\. Xing \(2025\)Dimensional collapse in VQVAEs: evidence and remedies\.InAdvances in Neural Information Processing Systems \(NeurIPS\),Cited by:[§1](https://arxiv.org/html/2605.06870#S1.p1.3),[§5](https://arxiv.org/html/2605.06870#S5.p1.1)\.
- R\. Zhang, P\. Isola, A\. A\. Efros, E\. Shechtman, and O\. Wang \(2018\)The unreasonable effectiveness of deep features as a perceptual metric\.InProceedings of the IEEE Conference on Computer Vision and Pattern Recognition \(CVPR\),pp\. 586–595\.Cited by:[§2](https://arxiv.org/html/2605.06870#S2.p2.1)\.
- X\. Zhang, D\. Zhang, S\. Li, Y\. Zhou, and X\. Qiu \(2024\)SpeechTokenizer: unified speech tokenizer for speech large language models\.InInternational Conference on Learning Representations \(ICLR\),Cited by:[§5](https://arxiv.org/html/2605.06870#S5.p1.1)\.
- W\. Zhao, Q\. Zou, R\. Shah, Y\. Wu, Z\. Lin, and D\. Liu \(2026\)Early quantization shrinks codebook: a simple fix for diversity\-preserving tokenization\.arXiv preprint arXiv:2603\.17052\.Cited by:[§1](https://arxiv.org/html/2605.06870#S1.p5.1)\.
- W\. Zhao, Z\. Liu, S\. Chen, X\. Cao, and G\. He \(2024\)Representation collapsing problems in vector quantization\.arXiv preprint arXiv:2411\.16550\.Cited by:[§1](https://arxiv.org/html/2605.06870#S1.p5.1)\.
- C\. Zheng and A\. Vedaldi \(2023\)Online clustered codebook\.InProceedings of the IEEE/CVF International Conference on Computer Vision \(ICCV\),Cited by:[§1](https://arxiv.org/html/2605.06870#S1.p1.3)\.
- L\. Zhu, F\. Liao, Y\. Li, J\. Jia, and L\. Xie \(2024\)Scaling the codebook size of VQGAN to 100,000 with a utilization rate of 99%\.InAdvances in Neural Information Processing Systems \(NeurIPS\),Cited by:[§1](https://arxiv.org/html/2605.06870#S1.p1.3)\.
- Y\. Zhu, B\. Li, Y\. Xin, Z\. Xia, and L\. Xu \(2025\)Addressing representation collapse in vector quantized models with one linear layer\.InProceedings of the IEEE/CVF International Conference on Computer Vision \(ICCV\),Note:arXiv:2411\.02038Cited by:[§1](https://arxiv.org/html/2605.06870#S1.p1.3),[§5](https://arxiv.org/html/2605.06870#S5.p1.1)\.

## Appendix AProofs for Section[3](https://arxiv.org/html/2605.06870#S3)

### A\.1Dense RD\-AE flow

We derive \([6](https://arxiv.org/html/2605.06870#S3.E6)\)–\([7](https://arxiv.org/html/2605.06870#S3.E7)\) using the notation of Section[3\.1](https://arxiv.org/html/2605.06870#S3.SS1)\. All channel parameters\(cj,τj2,Dj,D⋆\)\(c\_\{j\},\\tau\_\{j\}^\{2\},D\_\{j\},D^\{\\star\}\)are treated as constant throughout, consistent with the assumption that the quantizer remains optimal for the current latent statistics\.

##### Channel in the eigenbasis\.

We denote the covariance matrix of the encoder latent space asΣz\\Sigma\_\{z\}, and write the diagonal decompositionΣz=W1​Σ​W1⊤=U​Λ​U⊤\\Sigma\_\{z\}=W\_\{1\}\\Sigma W\_\{1\}^\{\\top\}=U\\Lambda U^\{\\top\}\. In coordinatesz~=U⊤​z\\tilde\{z\}=U^\{\\top\}z, the componentsz~j∼𝒩​\(0,λj\)\\tilde\{z\}\_\{j\}\\sim\\mathcal\{N\}\(0,\\lambda\_\{j\}\)are independent, so the rate\-RRchannel of Section[3\.1](https://arxiv.org/html/2605.06870#S3.SS1)applies coordinatewise with parameters\(cj,τj2,Dj\)\(c\_\{j\},\\tau\_\{j\}^\{2\},D\_\{j\}\)as in \([9](https://arxiv.org/html/2605.06870#S3.E9)\)\.

Returning to the original basis,𝔼​\[zq∣z\]=M​z\\mathbb\{E\}\[z\_\{q\}\\mid z\]=MzwithM:=U​diag⁡\(c1,…,cd\)​U⊤M:=U\\operatorname\{diag\}\(c\_\{1\},\\ldots,c\_\{d\}\)U^\{\\top\}\. Sincezq,jz\_\{q,j\}is the MMSE estimate ofz~j\\tilde\{z\}\_\{j\}, the residualn:=z−zqn:=z\-z\_\{q\}satisfieszq⟂nz\_\{q\}\\perp nby Gaussianity, withn∼𝒩​\(0,U​diag⁡\(D1,…,Dd\)​U⊤\)n\\sim\\mathcal\{N\}\(0,U\\operatorname\{diag\}\(D\_\{1\},\\ldots,D\_\{d\}\)U^\{\\top\}\)\.

##### Joint moments\.

The decompositionz=zq\+nz=z\_\{q\}\+nwithzq⟂nz\_\{q\}\\perp ngives

𝔼\[zqzq⊤\]=Udiag\(cjλj\)U⊤=:Γq,𝔼\[zzq⊤\]=Γq,𝔼\[xzq⊤\]=ΣW1⊤M,\\mathbb\{E\}\[z\_\{q\}z\_\{q\}^\{\\top\}\]=U\\operatorname\{diag\}\(c\_\{j\}\\lambda\_\{j\}\)U^\{\\top\}=:\\Gamma\_\{q\},\\qquad\\mathbb\{E\}\[zz\_\{q\}^\{\\top\}\]=\\Gamma\_\{q\},\\qquad\\mathbb\{E\}\[xz\_\{q\}^\{\\top\}\]=\\Sigma W\_\{1\}^\{\\top\}M,\(18\)whereM=U​diag⁡\(c1,…,cd\)​U⊤M=U\\operatorname\{diag\}\(c\_\{1\},\\ldots,c\_\{d\}\)U^\{\\top\}\. The first two follow becausezq⟂nz\_\{q\}\\perp nimplies𝔼​\[z​zq⊤\]=𝔼​\[zq​zq⊤\]\\mathbb\{E\}\[zz\_\{q\}^\{\\top\}\]=\\mathbb\{E\}\[z\_\{q\}z\_\{q\}^\{\\top\}\]\. The third useszq=M​z\+\(noise\)z\_\{q\}=Mz\+\(\\text\{noise\}\)andz=W1​xz=W\_\{1\}x, withMMsymmetric\.

##### Decoder gradient\.

FromLrec=𝔼​‖x−W2​zq‖2=tr⁡Σ−2​tr⁡\(W2​𝔼​\[zq​x⊤\]\)\+tr⁡\(W2⊤​W2​Γq\)L\_\{\\mathrm\{rec\}\}=\\mathbb\{E\}\\\|x\-W\_\{2\}z\_\{q\}\\\|^\{2\}=\\operatorname\{tr\}\\Sigma\-2\\operatorname\{tr\}\(W\_\{2\}\\mathbb\{E\}\[z\_\{q\}x^\{\\top\}\]\)\+\\operatorname\{tr\}\(W\_\{2\}^\{\\top\}W\_\{2\}\\Gamma\_\{q\}\), we treat the channel parameters\(M,Γq\)\(M,\\Gamma\_\{q\}\)as constant and differentiate:

∇W2Lrec=−2​Σ​W1⊤​M\+2​W2​Γq\.\\nabla\_\{W\_\{2\}\}L\_\{\\mathrm\{rec\}\}=\-2\\Sigma W\_\{1\}^\{\\top\}M\+2W\_\{2\}\\Gamma\_\{q\}\.\(19\)
LcomL\_\{\\text\{com\}\}does not depend onW2W\_\{2\}, so∇W2Lcom=0\\nabla\_\{W\_\{2\}\}L\_\{\\text\{com\}\}=0\.

##### Encoder reconstruction gradient under STE\.

The straight\-through estimator uses∂zq/∂z:=I\\partial z\_\{q\}/\\partial z:=Iin differentiation\. Equivalently, replacezqz\_\{q\}by the surrogatezqsurr=z\+sg​\(zq−z\)=W1​x\+sz\_\{q\}^\{\\mathrm\{surr\}\}=z\+\\mathrm\{sg\}\(z\_\{q\}\-z\)=W\_\{1\}x\+swheres:=zq−W1​xs:=z\_\{q\}\-W\_\{1\}xis treated as a constant\. The reconstruction loss becomes

LrecSTE=𝔼​‖x−W2​W1​x−W2​s‖2,L\_\{\\mathrm\{rec\}\}^\{\\mathrm\{STE\}\}\\;=\\;\\mathbb\{E\}\\\|x\-W\_\{2\}W\_\{1\}x\-W\_\{2\}s\\\|^\{2\},\(20\)identical in value toLrecL\_\{\\mathrm\{rec\}\}but withW1W\_\{1\}\-dependence routed only throughW1​xW\_\{1\}x\. Differentiating inW1W\_\{1\}withssconstant,

∇W1Lrec\|STE=−2W2⊤𝔼\[\(x−W2zq\)x⊤\]=−2W2⊤Σ\+2W2⊤W2𝔼\[zqx⊤\],\\nabla\_\{W\_\{1\}\}L\_\{\\mathrm\{rec\}\}\|^\{\\mathrm\{STE\}\}\\;=\\;\-2W\_\{2\}^\{\\top\}\\mathbb\{E\}\[\(x\-W\_\{2\}z\_\{q\}\)x^\{\\top\}\]\\;=\\;\-2W\_\{2\}^\{\\top\}\\Sigma\+2W\_\{2\}^\{\\top\}W\_\{2\}\\mathbb\{E\}\[z\_\{q\}x^\{\\top\}\],\(21\)where the residualx−W2​W1​x−W2​s=x−W2​zqx\-W\_\{2\}W\_\{1\}x\-W\_\{2\}s=x\-W\_\{2\}z\_\{q\}uses the truezqz\_\{q\}\. Substituting𝔼​\[zq​x⊤\]=M​W1​Σ\\mathbb\{E\}\[z\_\{q\}x^\{\\top\}\]=MW\_\{1\}\\Sigma,

∇W1Lrec\|STE=2\(W2⊤W2MW1−W2⊤\)Σ\.\\nabla\_\{W\_\{1\}\}L\_\{\\mathrm\{rec\}\}\|^\{\\mathrm\{STE\}\}\\;=\\;2\(W\_\{2\}^\{\\top\}W\_\{2\}MW\_\{1\}\-W\_\{2\}^\{\\top\}\)\\Sigma\.\(22\)

##### Encoder commitment gradient\.

The commitment termLcom=𝔼​‖z−sg​\(zq\)‖2L\_\{\\mathrm\{com\}\}=\\mathbb\{E\}\\\|z\-\\mathrm\{sg\}\(z\_\{q\}\)\\\|^\{2\}treatszqz\_\{q\}as constant during differentiation, so it only depends onW1W\_\{1\}throughz=W1​xz=W\_\{1\}x:

Lcom=tr⁡\(W1​Σ​W1⊤\)−2​𝔼​\[z⊤​sg​\(zq\)\]\+const\.L\_\{\\mathrm\{com\}\}\\;=\\;\\operatorname\{tr\}\(W\_\{1\}\\Sigma W\_\{1\}^\{\\top\}\)\-2\\mathbb\{E\}\[z^\{\\top\}\\mathrm\{sg\}\(z\_\{q\}\)\]\+\\mathrm\{const\}\.\(23\)Differentiating,∇W1tr⁡\(W1​Σ​W1⊤\)=2​W1​Σ\\nabla\_\{W\_\{1\}\}\\operatorname\{tr\}\(W\_\{1\}\\Sigma W\_\{1\}^\{\\top\}\)=2W\_\{1\}\\Sigmaand∇W1𝔼​\[z⊤​sg​\(zq\)\]=𝔼​\[sg​\(zq\)​x⊤\]\\nabla\_\{W\_\{1\}\}\\mathbb\{E\}\[z^\{\\top\}\\mathrm\{sg\}\(z\_\{q\}\)\]=\\mathbb\{E\}\[\\mathrm\{sg\}\(z\_\{q\}\)x^\{\\top\}\]has population meanM​W1​ΣMW\_\{1\}\\Sigma, giving

∇W1Lcom=2​\(W1−M​W1\)​Σ\.\\nabla\_\{W\_\{1\}\}L\_\{\\mathrm\{com\}\}\\;=\\;2\(W\_\{1\}\-MW\_\{1\}\)\\Sigma\.\(24\)

##### Assembling\.

Combining \([19](https://arxiv.org/html/2605.06870#A1.E19)\), \([22](https://arxiv.org/html/2605.06870#A1.E22)\), \([24](https://arxiv.org/html/2605.06870#A1.E24)\), we arrive at the flow in \([6](https://arxiv.org/html/2605.06870#S3.E6)\)–\([7](https://arxiv.org/html/2605.06870#S3.E7)\):

W˙1\\displaystyle\\dot\{W\}\_\{1\}\\;=−2​\(W2⊤​W2​M​W1−W2⊤\)​Σ−2​β​\(W1−M​W1\)​Σ,\\displaystyle=\\;\-2\(W\_\{2\}^\{\\top\}W\_\{2\}MW\_\{1\}\-W\_\{2\}^\{\\top\}\)\\Sigma\\;\-\\;2\\beta\(W\_\{1\}\-MW\_\{1\}\)\\Sigma,\(25\)W˙2\\displaystyle\\dot\{W\}\_\{2\}\\;=2​Σ​W1⊤​M−2​W2​Γq\.\\displaystyle=\\;2\\Sigma W\_\{1\}^\{\\top\}M\-2W\_\{2\}\\Gamma\_\{q\}\.\(26\)

#### A\.1\.1Dense water\-level evolution

In the main text, we examined the water\-level dynamics under diagonal initialization ofW1W\_\{1\}andW2W\_\{2\}\(Eq\. \([13](https://arxiv.org/html/2605.06870#S3.E13)\), derivation in §[A\.3](https://arxiv.org/html/2605.06870#A1.SS3)\), which characterizes howD∗D\*rises during RD\-AE training to remove modes from the active set \. We derive the water\-level dynamics under dense initialization here\.

Recall that the water\-fill constraint∑j∈A12​log2⁡\(λj/D⋆\)=R\\sum\_\{j\\in A\}\\tfrac\{1\}\{2\}\\log\_\{2\}\(\\lambda\_\{j\}/D^\{\\star\}\)=RcouplesD⋆D^\{\\star\}to the eigenvalues ofΣz​\(t\)\\Sigma\_\{z\}\(t\), which themselves evolve throughW1​\(t\)W\_\{1\}\(t\)\. To trackλ˙j\\dot\{\\lambda\}\_\{j\}, differentiate the eigenvalue identityΣz​uj=λj​uj\\Sigma\_\{z\}u\_\{j\}=\\lambda\_\{j\}u\_\{j\}and the normalizationuj⊤​uj=1u\_\{j\}^\{\\top\}u\_\{j\}=1:

Σ˙z​uj\+Σz​u˙j=λ˙j​uj\+λj​u˙j,uj⊤​u˙j=0\.\\dot\{\\Sigma\}\_\{z\}\\,u\_\{j\}\+\\Sigma\_\{z\}\\dot\{u\}\_\{j\}\\;=\\;\\dot\{\\lambda\}\_\{j\}u\_\{j\}\+\\lambda\_\{j\}\\dot\{u\}\_\{j\},\\qquad u\_\{j\}^\{\\top\}\\dot\{u\}\_\{j\}=0\.\(27\)Left\-multiplying byuj⊤u\_\{j\}^\{\\top\}and usinguj⊤​Σz=λj​uj⊤u\_\{j\}^\{\\top\}\\Sigma\_\{z\}=\\lambda\_\{j\}u\_\{j\}^\{\\top\}gives

λ˙j=uj⊤​Σ˙z​uj,Σ˙z=W˙1​Σ​W1⊤\+W1​Σ​W˙1⊤\.\\dot\{\\lambda\}\_\{j\}\\;=\\;u\_\{j\}^\{\\top\}\\dot\{\\Sigma\}\_\{z\}\\,u\_\{j\},\\qquad\\dot\{\\Sigma\}\_\{z\}\\;=\\;\\dot\{W\}\_\{1\}\\Sigma W\_\{1\}^\{\\top\}\+W\_\{1\}\\Sigma\\dot\{W\}\_\{1\}^\{\\top\}\.\(28\)The water\-level evolution follows by differentiating the rate constraint at fixed active set:

D˙⋆D⋆=1k​∑j∈Aλ˙jλj=1k​∑j∈Auj⊤​Σ˙z​ujλj\.\\frac\{\\dot\{D\}^\{\\star\}\}\{D^\{\\star\}\}\\;=\\;\\frac\{1\}\{k\}\\sum\_\{j\\in A\}\\frac\{\\dot\{\\lambda\}\_\{j\}\}\{\\lambda\_\{j\}\}\\;=\\;\\frac\{1\}\{k\}\\sum\_\{j\\in A\}\\frac\{u\_\{j\}^\{\\top\}\\dot\{\\Sigma\}\_\{z\}u\_\{j\}\}\{\\lambda\_\{j\}\}\.\(29\)Between water\-crossings,AAandkkare constant, and \([25](https://arxiv.org/html/2605.06870#A1.E25)\)–\([29](https://arxiv.org/html/2605.06870#A1.E29)\) close the dynamics of\(W1,W2,D⋆\)\(W\_\{1\},W\_\{2\},D^\{\\star\}\)\. At a crossing,AAupdates discretely,kkdecrements by one, andD⋆D^\{\\star\}is continuous\.

### A\.2RD\-AE flow derivation \(diagonal case\)

We derive the diagonal RD\-AE flow \([12](https://arxiv.org/html/2605.06870#S3.E12)\) stated in Section[3\.2](https://arxiv.org/html/2605.06870#S3.SS2)\. Throughout we assume thatW1=diag⁡\(uj\)W\_\{1\}=\\operatorname\{diag\}\(u\_\{j\}\),W2=diag⁡\(vj\)W\_\{2\}=\\operatorname\{diag\}\(v\_\{j\}\)\. The RD channel of rateRRbits acts coordinatewise onz=W1​xz=W\_\{1\}x\.

The channel parameterscj,τj2,D⋆c\_\{j\},\\tau\_\{j\}^\{2\},D^\{\\star\}are treated as constant while differentiating the loss, as in §[A\.1](https://arxiv.org/html/2605.06870#A1.SS1)\. However,cjc\_\{j\}andD⋆D^\{\\star\}change over time as the latent variancesλj=uj2​σj2\\lambda\_\{j\}=u\_\{j\}^\{2\}\\sigma\_\{j\}^\{2\}evolve, and this coupling is derived separately in §[A\.3](https://arxiv.org/html/2605.06870#A1.SS3)\.

##### Key moments\.

In the diagonal case, the dense decomposition of §[A\.1](https://arxiv.org/html/2605.06870#A1.SS1)reduces tozj=zq,j\+njz\_\{j\}=z\_\{q,j\}\+n\_\{j\}withnj∼𝒩​\(0,Dj\)n\_\{j\}\\sim\\mathcal\{N\}\(0,D\_\{j\}\)independent ofzq,jz\_\{q,j\}, giving

𝔼​\[zq,j2\]=cj​λj,𝔼​\[zj​zq,j\]=cj​λj,𝔼​\[xj​zq,j\]=cj​uj​σj2\.\\mathbb\{E\}\[z\_\{q,j\}^\{2\}\]=c\_\{j\}\\lambda\_\{j\},\\qquad\\mathbb\{E\}\[z\_\{j\}\\,z\_\{q,j\}\]=c\_\{j\}\\lambda\_\{j\},\\qquad\\mathbb\{E\}\[x\_\{j\}\\,z\_\{q,j\}\]=c\_\{j\}u\_\{j\}\\sigma\_\{j\}^\{2\}\.\(30\)

##### Decoder gradient\.

The decoder direct gradient updates\. Fromx^j=vj​zq,j\\hat\{x\}\_\{j\}=v\_\{j\}z\_\{q,j\},

∂Lrec∂vj=−2​𝔼​\[\(xj−vj​zq,j\)​zq,j\]=−2​\(𝔼​\[xj​zq,j\]−vj​𝔼​\[zq,j2\]\)=−2​cj​uj​σj2​\(1−uj​vj\),\\frac\{\\partial L\_\{\\mathrm\{rec\}\}\}\{\\partial v\_\{j\}\}\\;=\\;\-2\\,\\mathbb\{E\}\\bigl\[\(x\_\{j\}\-v\_\{j\}z\_\{q,j\}\)z\_\{q,j\}\\bigr\]\\;=\\;\-2\\bigl\(\\mathbb\{E\}\[x\_\{j\}z\_\{q,j\}\]\-v\_\{j\}\\mathbb\{E\}\[z\_\{q,j\}^\{2\}\]\\bigr\)\\;=\\;\-2c\_\{j\}u\_\{j\}\\sigma\_\{j\}^\{2\}\(1\-u\_\{j\}v\_\{j\}\),\(31\)usingcj​λj=cj​uj2​σj2c\_\{j\}\\lambda\_\{j\}=c\_\{j\}u\_\{j\}^\{2\}\\sigma\_\{j\}^\{2\}in the last step\. The commitment term does not depend onvjv\_\{j\}, so∂Lcom/∂vj=0\\partial L\_\{\\mathrm\{com\}\}/\\partial v\_\{j\}=0\.

##### Encoder reconstruction gradient under STE\.

The straight\-through estimator redefines the backward pass to treat∂zq,j/∂zj:=1\\partial z\_\{q,j\}/\\partial z\_\{j\}:=1, while the forward pass uses the actualzq,jz\_\{q,j\}from the channel\. Concretely, STE replaceszq,jz\_\{q,j\}with a surrogatezq,jsurr=zj\+sg​\(zq,j−zj\)z\_\{q,j\}^\{\\mathrm\{surr\}\}=z\_\{j\}\+\\mathrm\{sg\}\(z\_\{q,j\}\-z\_\{j\}\), which equalszq,jz\_\{q,j\}in value but has identity Jacobian inzjz\_\{j\}\. Under this surrogate, the reconstructionx^j=vj​zq,jsurr\\hat\{x\}\_\{j\}=v\_\{j\}z\_\{q,j\}^\{\\mathrm\{surr\}\}satisfies∂x^j/∂uj=vj​∂zj/∂uj=vj​xj\\partial\\hat\{x\}\_\{j\}/\\partial u\_\{j\}=v\_\{j\}\\,\\partial z\_\{j\}/\\partial u\_\{j\}=v\_\{j\}x\_\{j\}, and

∂Lrec∂uj\|STE=−2​vj​𝔼​\[\(xj−vj​zq,j\)​xj\]=−2​vj​\(σj2−vj​𝔼​\[xj​zq,j\]\)=−2​vj​σj2​\(1−cj​uj​vj\)\.\\frac\{\\partial L\_\{\\mathrm\{rec\}\}\}\{\\partial u\_\{j\}\}\\bigg\|\_\{\\mathrm\{STE\}\}\\;=\\;\-2v\_\{j\}\\,\\mathbb\{E\}\\bigl\[\(x\_\{j\}\-v\_\{j\}z\_\{q,j\}\)x\_\{j\}\\bigr\]\\;=\\;\-2v\_\{j\}\\bigl\(\\sigma\_\{j\}^\{2\}\-v\_\{j\}\\mathbb\{E\}\[x\_\{j\}z\_\{q,j\}\]\\bigr\)\\;=\\;\-2v\_\{j\}\\sigma\_\{j\}^\{2\}\(1\-c\_\{j\}u\_\{j\}v\_\{j\}\)\.\(32\)The value ofzq,jz\_\{q,j\}inside the residual uses the true channel output, not the surrogate — only the backward Jacobian is replaced\.

##### Encoder commitment gradient\.

The commitment termLcom=𝔼​‖z−sg​\(zq\)‖2L\_\{\\mathrm\{com\}\}=\\mathbb\{E\}\\\|z\-\\mathrm\{sg\}\(z\_\{q\}\)\\\|^\{2\}treatszqz\_\{q\}as constant during differentiation\. Sincezj=uj​xjz\_\{j\}=u\_\{j\}x\_\{j\}, differentiating with respect touju\_\{j\}gives:

∂Lcom∂uj=2​𝔼​\[\(zj−zq,j\)​xj\]=2​\(uj​σj2−𝔼​\[xj​zq,j\]\)=2​uj​σj2​\(1−cj\)=2​Djuj,\\frac\{\\partial L\_\{\\mathrm\{com\}\}\}\{\\partial u\_\{j\}\}\\;=\\;2\\,\\mathbb\{E\}\\bigl\[\(z\_\{j\}\-z\_\{q,j\}\)x\_\{j\}\\bigr\]\\;=\\;2\\bigl\(u\_\{j\}\\sigma\_\{j\}^\{2\}\-\\mathbb\{E\}\[x\_\{j\}z\_\{q,j\}\]\\bigr\)\\;=\\;2u\_\{j\}\\sigma\_\{j\}^\{2\}\(1\-c\_\{j\}\)\\;=\\;\\frac\{2D\_\{j\}\}\{u\_\{j\}\},\(33\)where the last equality usesDj=λj​\(1−cj\)=uj2​σj2​\(1−cj\)D\_\{j\}=\\lambda\_\{j\}\(1\-c\_\{j\}\)=u\_\{j\}^\{2\}\\sigma\_\{j\}^\{2\}\(1\-c\_\{j\}\)\.

##### Assembling the flow\.

The gradient flow isu˙j=−\(∂Lrec/∂uj\|STE\+β∂Lcom/∂uj\)\\dot\{u\}\_\{j\}=\-\\bigl\(\\partial L\_\{\\mathrm\{rec\}\}/\\partial u\_\{j\}\|^\{\\mathrm\{STE\}\}\+\\beta\\,\\partial L\_\{\\mathrm\{com\}\}/\\partial u\_\{j\}\\bigr\)andv˙j=−∂Lrec/∂vj\\dot\{v\}\_\{j\}=\-\\partial L\_\{\\mathrm\{rec\}\}/\\partial v\_\{j\}\. Substituting gives

u˙j=2​σj2​vj​\(1−cj​uj​vj\)−2​β​Djuj,v˙j=2​σj2​cj​uj​\(1−uj​vj\),\\dot\{u\}\_\{j\}\\;=\\;2\\sigma\_\{j\}^\{2\}\\,v\_\{j\}\(1\-c\_\{j\}u\_\{j\}v\_\{j\}\)\\;\-\\;\\frac\{2\\beta D\_\{j\}\}\{u\_\{j\}\},\\qquad\\dot\{v\}\_\{j\}\\;=\\;2\\sigma\_\{j\}^\{2\}\\,c\_\{j\}u\_\{j\}\(1\-u\_\{j\}v\_\{j\}\),\(34\)as stated in the main text\.

### A\.3Water\-level dynamics in the diagonal case

Here, we derive \([14](https://arxiv.org/html/2605.06870#S3.E14)\) from Section[3\.2](https://arxiv.org/html/2605.06870#S3.SS2), which characterizes how the reverse water\-filling variance threshold rises throughout RD\-AE training\. As described in the main text, this is extremely consequential: it progressively decreases the active set of modes which are represented by the quantizer\.

Given a rate budgetRR, a classical result\[Polyanskiy and Wu,[2025](https://arxiv.org/html/2605.06870#bib.bib2), Ex\. I\.9\]relates the rate\-distortion water levelD⋆D^\{\\star\}to the eigenvaluesλj\\lambda\_\{j\}via the constraint:

∑j12​log2\+⁡\(λj/D⋆\)=R,\\textstyle\\sum\_\{j\}\\tfrac\{1\}\{2\}\\log\_\{2\}^\{\+\}\(\\lambda\_\{j\}/D^\{\\star\}\)\\;=\\;R,\(35\)We define the set of active modes to beA:=\{i:λi\>D∗\}A:=\\\{i:\\lambda\_\{i\}\>D\*\\\}\. The sum \([36](https://arxiv.org/html/2605.06870#A1.E36)\) then reduces to

∑i∈A12​log2⁡\(λi/D⋆\)=R\.\\sum\_\{i\\in A\}\\tfrac\{1\}\{2\}\\log\_\{2\}\(\\lambda\_\{i\}/D^\{\\star\}\)\\;=\\;R\.\(36\)We derive its dependence onuiu\_\{i\}\(equivalently, onλi=ui2​σi2\\lambda\_\{i\}=u\_\{i\}^\{2\}\\sigma\_\{i\}^\{2\}\) by implicit differentiation\.

Rewriting \([36](https://arxiv.org/html/2605.06870#A1.E36)\) as∑i∈Alog2⁡λi−k​log2⁡D⋆=2​R\\sum\_\{i\\in A\}\\log\_\{2\}\\lambda\_\{i\}\-k\\log\_\{2\}D^\{\\star\}=2R, wherek=\|A\|k=\|A\|, and differentiating inuℓu\_\{\\ell\}forℓ∈A\\ell\\in A:

1λℓ​ln⁡2⋅∂λℓ∂uℓ−kD⋆​ln⁡2⋅∂D⋆∂uℓ=0\.\\frac\{1\}\{\\lambda\_\{\\ell\}\\ln 2\}\\cdot\\frac\{\\partial\\lambda\_\{\\ell\}\}\{\\partial u\_\{\\ell\}\}\\;\-\\;\\frac\{k\}\{D^\{\\star\}\\ln 2\}\\cdot\\frac\{\\partial D^\{\\star\}\}\{\\partial u\_\{\\ell\}\}\\;=\\;0\.\(37\)Using∂λℓ/∂uℓ=2​uℓ​σℓ2=2​λℓ/uℓ\\partial\\lambda\_\{\\ell\}/\\partial u\_\{\\ell\}=2u\_\{\\ell\}\\sigma\_\{\\ell\}^\{2\}=2\\lambda\_\{\\ell\}/u\_\{\\ell\}, this gives

∂D⋆∂uℓ=2​D⋆k​uℓfor​ℓ∈A\.\\frac\{\\partial D^\{\\star\}\}\{\\partial u\_\{\\ell\}\}\\;=\\;\\frac\{2D^\{\\star\}\}\{k\\,u\_\{\\ell\}\}\\qquad\\text\{for \}\\ell\\in A\.\(38\)Forℓ∉A\\ell\\notin A, modeℓ\\elldoes not enter \([36](https://arxiv.org/html/2605.06870#A1.E36)\) and∂D⋆/∂uℓ=0\\partial D^\{\\star\}/\\partial u\_\{\\ell\}=0\.

The time\-derivative form follows by the chain rule along a flow trajectory:

D˙⋆D⋆=∑ℓ∈A1D⋆​∂D⋆∂uℓ​u˙ℓ=2k​∑ℓ∈Au˙ℓuℓ\.\\frac\{\\dot\{D\}^\{\\star\}\}\{D^\{\\star\}\}\\;=\\;\\sum\_\{\\ell\\in A\}\\frac\{1\}\{D^\{\\star\}\}\\frac\{\\partial D^\{\\star\}\}\{\\partial u\_\{\\ell\}\}\\dot\{u\}\_\{\\ell\}\\;=\\;\\frac\{2\}\{k\}\\sum\_\{\\ell\\in A\}\\frac\{\\dot\{u\}\_\{\\ell\}\}\{u\_\{\\ell\}\}\.\(39\)
As active modes grow,D⋆D^\{\\star\}rises via \([36](https://arxiv.org/html/2605.06870#A1.E36)\)\. WhenD⋆D^\{\\star\}reachesλj\\lambda\_\{j\}for the weakest active mode, that mode transitions to inactive\. Onceinactive, it remains so: atβ=1\\beta=1, inactive modes are frozen \(r˙j=0\\dot\{r\}\_\{j\}=0from \([13](https://arxiv.org/html/2605.06870#S3.E13)\)\), soλj\\lambda\_\{j\}cannot grow back aboveD⋆D^\{\\star\}\.

### A\.4Per\-mode reduction atβ=1\\beta=1and balanced init

We specialize the flow toβ=1\\beta=1and balanced initializationuj​\(0\)=vj​\(0\)=εu\_\{j\}\(0\)=v\_\{j\}\(0\)=\\varepsilon\. We show that \(i\) balance is preserved on active modes, \(ii\) the active\-mode activationrj=uj​vj=uj2r\_\{j\}=u\_\{j\}v\_\{j\}=u\_\{j\}^\{2\}obeys a Saxe\-like logistic, and \(iii\) inactive modes at balance are frozen\.

##### Balance preservation on active modes\.

From the flow withβ=1\\beta=1andj∈Aj\\in A:

u˙j−v˙j=2​σj2​vj​\(1−cj​uj​vj\)−2​Djuj−2​σj2​cj​uj​\(1−uj​vj\)\.\\dot\{u\}\_\{j\}\-\\dot\{v\}\_\{j\}\\;=\\;2\\sigma\_\{j\}^\{2\}v\_\{j\}\(1\-c\_\{j\}u\_\{j\}v\_\{j\}\)\-\\frac\{2D\_\{j\}\}\{u\_\{j\}\}\-2\\sigma\_\{j\}^\{2\}c\_\{j\}u\_\{j\}\(1\-u\_\{j\}v\_\{j\}\)\.\(40\)Evaluating atuj=vj=:aju\_\{j\}=v\_\{j\}=:a\_\{j\}and usingDj=D⋆=\(1−cj\)​λj=\(1−cj\)​aj2​σj2D\_\{j\}=D^\{\\star\}=\(1\-c\_\{j\}\)\\lambda\_\{j\}=\(1\-c\_\{j\}\)a\_\{j\}^\{2\}\\sigma\_\{j\}^\{2\}on active modes,

u˙j−v˙j\|uj=vj=2​σj2​aj​\(1−cj​aj2\)−2​σj2​\(1−cj\)​aj−2​σj2​cj​aj​\(1−aj2\)=0\.\\dot\{u\}\_\{j\}\-\\dot\{v\}\_\{j\}\\big\|\_\{u\_\{j\}=v\_\{j\}\}\\;=\\;2\\sigma\_\{j\}^\{2\}a\_\{j\}\(1\-c\_\{j\}a\_\{j\}^\{2\}\)\-2\\sigma\_\{j\}^\{2\}\(1\-c\_\{j\}\)a\_\{j\}\-2\\sigma\_\{j\}^\{2\}c\_\{j\}a\_\{j\}\(1\-a\_\{j\}^\{2\}\)\\;=\\;0\.\(41\)Henceuj​\(t\)=vj​\(t\)u\_\{j\}\(t\)=v\_\{j\}\(t\)for alltton active modes; writeaj​\(t\)a\_\{j\}\(t\)for this common value\.

##### Active\-mode logistic\.

At balance,u˙j=v˙j\\dot\{u\}\_\{j\}=\\dot\{v\}\_\{j\}, so we compute either\. Fromv˙j=2​σj2​cj​uj​\(1−uj​vj\)\\dot\{v\}\_\{j\}=2\\sigma\_\{j\}^\{2\}c\_\{j\}u\_\{j\}\(1\-u\_\{j\}v\_\{j\}\)withuj=vj=aju\_\{j\}=v\_\{j\}=a\_\{j\}:

a˙j=2​σj2​cj​aj​\(1−aj2\)\.\\dot\{a\}\_\{j\}\\;=\\;2\\sigma\_\{j\}^\{2\}c\_\{j\}a\_\{j\}\(1\-a\_\{j\}^\{2\}\)\.\(42\)Withrj=aj2r\_\{j\}=a\_\{j\}^\{2\},r˙j=2​aj​a˙j\\dot\{r\}\_\{j\}=2a\_\{j\}\\dot\{a\}\_\{j\}, giving

r˙j=4​σj2​cj​rj​\(1−rj\)for​j∈A,\\dot\{r\}\_\{j\}\\;=\\;4\\sigma\_\{j\}^\{2\}c\_\{j\}\\,r\_\{j\}\(1\-r\_\{j\}\)\\qquad\\text\{for \}j\\in A,\(43\)which is Saxe’s logistic \([4](https://arxiv.org/html/2605.06870#S2.E4)\) with timescale1/\(4​σj2\)1/\(4\\sigma\_\{j\}^\{2\}\)stretched to1/\(4​σj2​cj\)1/\(4\\sigma\_\{j\}^\{2\}c\_\{j\}\)by the channel gain\.

##### Inactive modes\.

Forj∉Aj\\notin A,cj=τj2=0c\_\{j\}=\\tau\_\{j\}^\{2\}=0andDj=λj=uj2​σj2D\_\{j\}=\\lambda\_\{j\}=u\_\{j\}^\{2\}\\sigma\_\{j\}^\{2\}\. The flow reduces to

u˙j=2​σj2​vj−2​uj2​σj2uj=2​σj2​\(vj−uj\),v˙j=0\.\\dot\{u\}\_\{j\}\\;=\\;2\\sigma\_\{j\}^\{2\}v\_\{j\}\-\\frac\{2u\_\{j\}^\{2\}\\sigma\_\{j\}^\{2\}\}\{u\_\{j\}\}\\;=\\;2\\sigma\_\{j\}^\{2\}\(v\_\{j\}\-u\_\{j\}\),\\qquad\\dot\{v\}\_\{j\}\\;=\\;0\.\(44\)At balanceuj=vju\_\{j\}=v\_\{j\},u˙j=0\\dot\{u\}\_\{j\}=0andv˙j=0\\dot\{v\}\_\{j\}=0, so inactive modes at balance are stationary\.

### A\.5Plain\-AE mode activation \(closed form\)

We derive the closed\-form solution of the plain\-AE flow fromSaxeet al\.\[[2014](https://arxiv.org/html/2605.06870#bib.bib27)\]for completeness and reference in the warm\-up theorem proof\. The plain\-AE gradient flow from the decoupled loss is

u˙j=2​σj2​vj​\(1−uj​vj\),v˙j=2​σj2​uj​\(1−uj​vj\)\.\\dot\{u\}\_\{j\}\\;=\\;2\\sigma\_\{j\}^\{2\}v\_\{j\}\(1\-u\_\{j\}v\_\{j\}\),\\qquad\\dot\{v\}\_\{j\}\\;=\\;2\\sigma\_\{j\}^\{2\}u\_\{j\}\(1\-u\_\{j\}v\_\{j\}\)\.\(45\)Fromu˙j−v˙j=−2​σj2​\(1−uj​vj\)​\(uj−vj\)\\dot\{u\}\_\{j\}\-\\dot\{v\}\_\{j\}=\-2\\sigma\_\{j\}^\{2\}\(1\-u\_\{j\}v\_\{j\}\)\(u\_\{j\}\-v\_\{j\}\), the balance gapuj−vju\_\{j\}\-v\_\{j\}decays to zero from any initial condition; in particular, balanced initializationuj​\(0\)=vj​\(0\)=εu\_\{j\}\(0\)=v\_\{j\}\(0\)=\\varepsilonis preserved\. Writingaj​\(t\)a\_\{j\}\(t\)for the common value andrj=aj2r\_\{j\}=a\_\{j\}^\{2\},

a˙j=2​σj2​aj​\(1−aj2\),r˙j=4​σj2​rj​\(1−rj\)\.\\dot\{a\}\_\{j\}\\;=\\;2\\sigma\_\{j\}^\{2\}a\_\{j\}\(1\-a\_\{j\}^\{2\}\),\\qquad\\dot\{r\}\_\{j\}\\;=\\;4\\sigma\_\{j\}^\{2\}r\_\{j\}\(1\-r\_\{j\}\)\.\(46\)Separating variables,d​rj/\[rj​\(1−rj\)\]=4​σj2​d​t\\mathrm\{d\}r\_\{j\}/\[r\_\{j\}\(1\-r\_\{j\}\)\]=4\\sigma\_\{j\}^\{2\}\\,\\mathrm\{d\}t, and integrating fromrj​\(0\)=ε2r\_\{j\}\(0\)=\\varepsilon^\{2\}gives

rj​\(t\)=\(1\+1−ε2ε2​e−4​σj2​t\)−1\.r\_\{j\}\(t\)\\;=\\;\\Bigl\(1\+\\tfrac\{1\-\\varepsilon^\{2\}\}\{\\varepsilon^\{2\}\}e^\{\-4\\sigma\_\{j\}^\{2\}t\}\\Bigr\)^\{\-1\}\.\(47\)

### A\.6Proof of Theorem[1](https://arxiv.org/html/2605.06870#Thmtheorem1)

The proof has three steps: \(i\) identify the active set at the switch time, \(ii\) show warmup\-inactive modes stay frozen, \(iii\) derive the loss bound\.

##### State at the switch\.

Plain\-AE training from balanced initializationW1​\(0\)=W2​\(0\)=ε​IW\_\{1\}\(0\)=W\_\{2\}\(0\)=\\varepsilon Ipreserves diagonality and balance \(§[A\.5](https://arxiv.org/html/2605.06870#A1.SS5)\), so at timeTwuT\_\{\\mathrm\{wu\}\}the weights areW1=W2=diag⁡\(gj​\(Twu\)\)W\_\{1\}=W\_\{2\}=\\operatorname\{diag\}\(\\sqrt\{g\_\{j\}\(T\_\{\\mathrm\{wu\}\}\)\}\)and the latent eigenvalues areλj​\(Twu\)=gj​\(Twu\)​σj2\\lambda\_\{j\}\(T\_\{\\mathrm\{wu\}\}\)=g\_\{j\}\(T\_\{\\mathrm\{wu\}\}\)\\,\\sigma\_\{j\}^\{2\}\. Sinceσj2\\sigma\_\{j\}^\{2\}is strictly decreasing injjandgj​\(Twu\)g\_\{j\}\(T\_\{\\mathrm\{wu\}\}\)is strictly decreasing injj\(largerσj2\\sigma\_\{j\}^\{2\}gives faster Saxe activation\),λj​\(Twu\)\\lambda\_\{j\}\(T\_\{\\mathrm\{wu\}\}\)is strictly decreasing injj\.

The active set at the switch is therefore a top prefix\{1,…,m\}\\\{1,\\ldots,m\\\}, wheremmis the maximal integer satisfying the water\-fill consistency conditionλm​\(Twu\)\>D0⋆\\lambda\_\{m\}\(T\_\{\\mathrm\{wu\}\}\)\>D^\{\\star\}\_\{0\}, withD0⋆=\(∏i≤mλi​\(Twu\)\)1/m​2−2​R/mD^\{\\star\}\_\{0\}=\\bigl\(\\prod\_\{i\\leq m\}\\lambda\_\{i\}\(T\_\{\\mathrm\{wu\}\}\)\\bigr\)^\{1/m\}2^\{\-2R/m\}\. Rearranging yieldsR\>12​∑i<mlog2⁡\(σi2​gi/σm2​gm\)R\>\\tfrac\{1\}\{2\}\\sum\_\{i<m\}\\log\_\{2\}\(\\sigma\_\{i\}^\{2\}g\_\{i\}/\\sigma\_\{m\}^\{2\}g\_\{m\}\), som=mwu​\(Twu,R\)m=m^\{\\mathrm\{wu\}\}\(T\_\{\\mathrm\{wu\}\},R\)\.

##### Warmup\-inactive modes are frozen forever\.

Forj\>mwuj\>m^\{\\mathrm\{wu\}\}, modejjis inactive at the switch \(cj=τj2=0c\_\{j\}=\\tau\_\{j\}^\{2\}=0\) and balanced\. By §[A\.4](https://arxiv.org/html/2605.06870#A1.SS4),u˙j=v˙j=0\\dot\{u\}\_\{j\}=\\dot\{v\}\_\{j\}=0at balance, soλj\\lambda\_\{j\}is constant\. The water levelD⋆D^\{\\star\}is non\-decreasing along the flow \(§[A\.3](https://arxiv.org/html/2605.06870#A1.SS3)\), soλj<D⋆\\lambda\_\{j\}<D^\{\\star\}for allt≥Twut\\geq T\_\{\\mathrm\{wu\}\}, and modejjremains inactive\.

##### The active set shrinks as a prefix\.

Active modes haver˙j=4​σj2​cj​rj​\(1−rj\)≥0\\dot\{r\}\_\{j\}=4\\sigma\_\{j\}^\{2\}c\_\{j\}r\_\{j\}\(1\-r\_\{j\}\)\\geq 0\(§[A\.4](https://arxiv.org/html/2605.06870#A1.SS4)\), soλj\\lambda\_\{j\}is non\-decreasing on the active set\. Combined with monotonicD⋆D^\{\\star\}, modes can only exit the active set, not enter\. Sinceλj\\lambda\_\{j\}remains ordered \(largerσj2\\sigma\_\{j\}^\{2\}drives faster growth, preserving the decreasing sequence\), the active set at any later time is a prefix\{1,…,k\}\\\{1,\\ldots,k\\\}withk≤mwuk\\leq m^\{\\mathrm\{wu\}\}\. Any surviving mode converges torj=1r\_\{j\}=1, givingA∞=\{1,…,k∞\}A\_\{\\infty\}=\\\{1,\\ldots,k\_\{\\infty\}\\\}withk∞≤mwuk\_\{\\infty\}\\leq m^\{\\mathrm\{wu\}\}andλj=σj2\\lambda\_\{j\}=\\sigma\_\{j\}^\{2\}onA∞A\_\{\\infty\}\.

##### Loss bound\.

At final state, the reconstruction loss is

Lrec∞=k∞​D∞⋆\+∑j\>k∞σj2=k∞​\(∏i≤k∞σi2\)1/k∞​2−2​R/k∞\+∑j\>k∞σj2,L\_\{\\mathrm\{rec\}\}^\{\\infty\}\\;=\\;k\_\{\\infty\}\\,D^\{\\star\}\_\{\\infty\}\\;\+\\;\\sum\_\{j\>k\_\{\\infty\}\}\\sigma\_\{j\}^\{2\}\\;=\\;k\_\{\\infty\}\\bigl\(\\textstyle\\prod\_\{i\\leq k\_\{\\infty\}\}\\sigma\_\{i\}^\{2\}\\bigr\)^\{1/k\_\{\\infty\}\}2^\{\-2R/k\_\{\\infty\}\}\\;\+\\;\\sum\_\{j\>k\_\{\\infty\}\}\\sigma\_\{j\}^\{2\},\(48\)since each surviving mode pays water\-fill distortionD∞⋆=δk∞​\(R\)D^\{\\star\}\_\{\\infty\}=\\delta\_\{k\_\{\\infty\}\}\(R\)and each inactive mode paysσj2\\sigma\_\{j\}^\{2\}\(the decoder receives zero on inactive coordinates\)\. This is the Shannon distortion of an idealized scheme that activates the topk∞k\_\{\\infty\}modes and ignores the rest\. Sincek∞≤mwuk\_\{\\infty\}\\leq m^\{\\mathrm\{wu\}\}and the Shannon optimal scheme for the topmwum^\{\\mathrm\{wu\}\}source at rateRRactivates allmwum^\{\\mathrm\{wu\}\}modes \(with distortionmwu​δmwu​\(R\)\+∑j\>mwuσj2m^\{\\mathrm\{wu\}\}\\delta\_\{m^\{\\mathrm\{wu\}\}\}\(R\)\+\\sum\_\{j\>m^\{\\mathrm\{wu\}\}\}\\sigma\_\{j\}^\{2\}\), thek∞k\_\{\\infty\}\-mode scheme is suboptimal, and

Lrec∞≥mwu​δmwu​\(R\)\+∑j\>mwuσj2\.L\_\{\\mathrm\{rec\}\}^\{\\infty\}\\;\\geq\\;m^\{\\mathrm\{wu\}\}\\,\\delta\_\{m^\{\\mathrm\{wu\}\}\}\(R\)\\;\+\\;\\sum\_\{j\>m^\{\\mathrm\{wu\}\}\}\\sigma\_\{j\}^\{2\}\.\(49\)∎

#### A\.6\.1Numerical validation of Theorem[1](https://arxiv.org/html/2605.06870#Thmtheorem1)

![Refer to caption](https://arxiv.org/html/2605.06870v1/figures/fig_theory_warmup.png)Figure 6:Warmup theorem: predictions vs\. simulation\.Balanced initializationW1​\(0\)=W2​\(0\)=ε​IW\_\{1\}\(0\)=W\_\{2\}\(0\)=\\varepsilon Iwithε=0\.1\\varepsilon=0\.1,β=1\\beta=1,d=64d=64,σj2=j−1\\sigma\_\{j\}^\{2\}=j^\{\-1\}, three ratesR∈\{12,14,16\}R\\in\\\{12,14,16\\\}bits\.\(a\)Active\-mode count at convergence\. Solid lines: predictedmwu​\(Twu,R\)m^\{\\mathrm\{wu\}\}\(T\_\{\\mathrm\{wu\}\},R\)from \([16](https://arxiv.org/html/2605.06870#S3.E16)\)\. Dashed lines: observedk∞k\_\{\\infty\}from simulated RD\-AE flow\.\(b\)Reconstruction loss\. Solid lines: loss bound from \([17](https://arxiv.org/html/2605.06870#S3.E17)\)\. Dashed lines: observedLrec∞L\_\{\\mathrm\{rec\}\}^\{\\infty\}\. Dotted horizontal lines mark the Shannon floorD​\(R\)D\(R\)for each rate\. AsTwuT\_\{\\mathrm\{wu\}\}grows, predictions and observations both approach the Shannon limit; at smallTwuT\_\{\\mathrm\{wu\}\}the bound is vacuous and the observed collapse is severe\.Figure[6](https://arxiv.org/html/2605.06870#A1.F6)compares the theorem to simulation across warm\-up durations in the linear RD\-AE setting\. Left panel:mwu​\(Twu,R\)m^\{\\mathrm\{wu\}\}\(T\_\{\\mathrm\{wu\}\},R\)from \([16](https://arxiv.org/html/2605.06870#S3.E16)\) against the observed final active\-set sizek∞k\_\{\\infty\}in RD\-AE simulation, forR∈\{10,14,18\}R\\in\\\{10,14,18\\\}bits\. Right panel: the loss bound \([17](https://arxiv.org/html/2605.06870#S3.E17)\) against the observed finalLrec∞L\_\{\\mathrm\{rec\}\}^\{\\infty\}, with the rate\-distortion floorD​\(R\)D\(R\)as reference\. The figure showsTwuT\_\{\\mathrm\{wu\}\}in the regime where the bound is informative; atTwu=0T\_\{\\mathrm\{wu\}\}=0the bound reduces toD​\(R\)D\(R\)and the observed loss exceeds it\. The bound tightens asTwuT\_\{\\mathrm\{wu\}\}grows, and bothk∞k\_\{\\infty\}andLrec∞L\_\{\\mathrm\{rec\}\}^\{\\infty\}approach the Shannon limit\.

## Appendix BExperiment Details

![Refer to caption](https://arxiv.org/html/2605.06870v1/figures/vqgan_latent_dim_comparison_broken.png)\(a\)VQGAN on ImageNet\-100 \(20k\)
![Refer to caption](https://arxiv.org/html/2605.06870v1/figures/wav_latent_dim_comparison_broken.png)\(b\)WavTokenizer on LibriTTS\.

Figure 7:Without VQ, real architectures exhibit Saxe\-like sequential activation when trained with Adam\. VQ causes dimensional collapse\.Latent effective dimension \(deffd\_\{\\mathrm\{eff\}\},99%99\\%\-variance threshold\) during autoencoder and standard VQ\-VAE training\. In both \(a\) VQGAN and \(b\) WavTokenizer, standard VQ training pins latent dimension at 2\-4, while dimension rises as modes are learned sequentially when training as an autoencoder\.### B\.1VQGAN Experiments

##### Architecture and loss\.

We use the f16 VQGAN encoder–decoder fromEsseret al\.\[[2021](https://arxiv.org/html/2605.06870#bib.bib5)\]with theimagenet\_vqgan\.yamlconfiguration \(128 base channels, multipliers\[1,1,2,2,4\]\[1,1,2,2,4\], two residual blocks per level, self\-attention at16×1616\\times 16, 256\-channel bottleneck, 256\-dimensional codebook embedding space\)\. At128×128128\\times 128input resolution each image is tokenized to an8×88\\times 8grid of 64 discrete tokens\. As per the original VQGAN training recipe, the discriminator is disabled for the first 250k training steps, the remaining loss isℒ=ℒL1\+ℒLPIPS\+λcb​ℒVQ\\mathcal\{L\}=\\mathcal\{L\}\_\{L\_\{1\}\}\+\\mathcal\{L\}\_\{\\text\{LPIPS\}\}\+\\lambda\_\{\\text\{cb\}\}\\,\\mathcal\{L\}\_\{\\text\{VQ\}\}withλcb=1\.0\\lambda\_\{\\text\{cb\}\}=1\.0and commitment weightβ=0\.25\\beta=0\.25, followingEsseret al\.\[[2021](https://arxiv.org/html/2605.06870#bib.bib5)\]\.

##### Training\.

All runs use the default VQGAN optimizer settings: Adam with\(β1,β2\)=\(0\.5,0\.9\)\(\\beta\_\{1\},\\beta\_\{2\}\)=\(0\.5,0\.9\), no weight decay, constant effective learning rate2\.88×10−42\.88\\times 10^\{\-4\}, and 100,000 steps\. We emphasize that any “warm\-up” referenced in this paper is the AE phase on the model, not an LR schedule\. Runs are distributed across4×4\\timesA100 40 GB GPUs at batch size 16/GPU \(effective batch size 64\), completing in approximately 5\.5 hours of wall\-clock time\.

##### Codebook handling\.

*Vanilla VQGAN*initializes the codebook from𝒰​\(−1/K,1/K\)\\mathcal\{U\}\(\-1/K,1/K\)and updates it via straight\-through gradients\.*VQGAN w/ Respawn*and all*AE Warm\-up*recipes usekk\-means initialization on encoder outputs \(100 Lloyd iterations over the full training set\) with EMA updates, and respawn any code whose EMA usage \(decay 0\.99\) falls below 0\.01 from a random encoder output in the current batch\. Respawn masks are broadcast from rank 0 to prevent NCCL desynchronization\.

##### Data\.

We start with the ImageNet\-100 dataset fromTianet al\.\[[2020](https://arxiv.org/html/2605.06870#bib.bib37)\]and sample 20,000 training images uniformly at random\. We retain the entire 5,000\-image validation set\. Images are center\-cropped and resized to128×128128\\times 128\.

##### Evaluation\.

All metrics in Table[1](https://arxiv.org/html/2605.06870#S4.T1)are computed on the 5,000\-image validation set at step 100,000\. L1 and LPIPS \(VGG\-16 backbone\) are the standard pixel and perceptual reconstruction terms; Rec Loss is their sum\. CB Util% is the fraction of codes receiving at least one assignment on the validation set, and CB Dim is the 99%\-variance PCA effective dimension of the codebook entries\. rFID uses the InceptionV3 backbone viatorchmetrics\.

### B\.2WavTokenizer Experiments

##### Architecture\.

We use the official WavTokenizer 75 FPS architecture fromJiet al\.\[[2025](https://arxiv.org/html/2605.06870#bib.bib36)\]with a 512\-dimensional pre\-quantization latent and a single codebook\. Unless otherwise noted, all architectural and loss hyperparameters are inherited from the released 75 FPS checkpoint configuration, including the convolutional encoder/decoder backbone, the multi\-period and multi\-scale STFT discriminators, and the full reconstruction loss suite \(time\-domain, mel\-spectrogram, feature matching, adversarial\)\. We vary only the codebook sizeK∈\{4096,8192,16384,65536\}K\\in\\\{4096,8192,16384,65536\\\}and the training regime \(cold\-start vs\. AE warm\-up\)\.

##### Training\.

All runs train for 100 epochs on2×2\\timesH100 GPUs with a per\-GPU batch size of 40; at 75 FPS this corresponds to 9,000 latent vectors per GPU per batch\. The optimizer, learning\-rate schedule, discriminator schedule, and loss weightings follow the released WavTokenizer recipe for the 75FPS model with codebook size 4096\. All runs use the same random seed \(the WavTokenizer default\); we did not run multiple seeds due to compute constraints–each run takes approximately 7 days on 2XH100s\.

##### AE warm\-up protocol\.

For warm\-start runs, we first train the encoder and decoder as a plain autoencoder for 43 epochs withβ=0\\beta=0andQ​\(z\)=zQ\(z\)=z,retaining the full non\-VQ loss suite \(reconstruction \+ discriminator\)\. This single warm\-up checkpoint is reused across all codebook sizes\. At epoch 43, we introduce VQ: we initialize the codebook via scikit\-learnkk\-means on encoder outputs forK∈\{4096,8192,16384\}K\\in\\\{4096,8192,16384\\\}, and fall back to random initialization forK=65536K=65536, wherekk\-means becomes prohibitive at this codebook count\. Training then proceeds for the remaining 57 epochs under the standard recipe with EMA codebook updates \(decay 0\.99\) and respawn enabled\. We did not sweep over warm\-up duration due to compute constraints; 43 epochs was chosen as the point at which the latent effective dimensiondeffd\_\{\\text\{eff\}\}visibly plateaus in Figure[7\(b\)](https://arxiv.org/html/2605.06870#A2.F7.sf2), matching the stopping rule advocated in Section[4](https://arxiv.org/html/2605.06870#S4)\.

##### Codebook respawn\.

We respawn codes whose EMA usage falls below a codebook\-size\-dependent threshold, aggregating usage counts across both GPUs each epoch for stability at largeKK\. Cold\-start baselines use the WavTokenizer default threshold of 1\.0\. Warm\-start runs scale the threshold inversely withKKto keep the expected number of respawned codes per epoch roughly constant:

##### Data\.

We train on both train\-clean\-360 and train\-other\-500 sampled at 24 kHz, with audio clips of 3s, following the original WavTokenizer training recipe\. Validation and model selection use 100 randomly sampled batches of combined dev\-clean and dev\-other; final evaluation reports metrics on both LibriTTS test\-clean and test\-other\. For each run we select the checkpoint with the best validation loss\.

##### Evaluation\.

We report five reconstruction metrics on test\-clean and test\-other: mel\-spectrogram L1 loss \(Mel\), UTMOS\[Saekiet al\.,[2022](https://arxiv.org/html/2605.06870#bib.bib30)\]as a perceptual MOS predictor, wide\-band PESQ\[Rixet al\.,[2001](https://arxiv.org/html/2605.06870#bib.bib32)\], STOI\[Taalet al\.,[2010](https://arxiv.org/html/2605.06870#bib.bib33)\], and voiced/unvoiced F1 \(V/UV F1\) computed following the methodology ofMorrisonet al\.\[[2022](https://arxiv.org/html/2605.06870#bib.bib34)\], using CREPE\[Kimet al\.,[2018](https://arxiv.org/html/2605.06870#bib.bib35)\]for pitch and periodicity extraction as in the reference CARGAN implementation\. With the exception of L1 loss, these match the metrics reported inJiet al\.\[[2025](https://arxiv.org/html/2605.06870#bib.bib36)\]\. Codebook utilization is computed over the full evaluation set\. and codebook effective dimensiondim\(𝒞\)\\dim\(\\mathcal\{C\}\)is the number of PCA components needed to explain 99% of the variance across codebook entries, matching the definition used for VQGAN in Appendix[C](https://arxiv.org/html/2605.06870#A3)\.

## Appendix CAdditional results on VQGAN

![Refer to caption](https://arxiv.org/html/2605.06870v1/figures/ref128_comparison_100k_loss_dim_v3.png)Figure 8:VQGAN on ImageNet\-100: codebook size×\\timeswarm\-up duration\.Training runs spanningK∈\{1​k,16​k,65​k\}K\\in\\\{1\\text\{k\},\\,16\\text\{k\},65\\text\{k\}\\\}andTwu∈\{0,1​k,5​k,10​k,20​k,30​k,40​k\}T\_\{\\mathrm\{wu\}\}\\in\\\{0,\\,1\\text\{k\},\\,5\\text\{k\},\\,10\\text\{k\},\\,20\\text\{k\},\\,30\\text\{k\},\\,40\\text\{k\}\\\}steps of AE Warm\-up, all with k\-means init and respawning of the codebook \(solid lines\), compared with the vanilla VQGAN training recipe with random codebook init and no respawn \(dashed line\)\.*Top row:*validation reconstruction loss \(L1 \+ LPIPS\) versus training step\.*Bottom row:*codebook effective dimensiondeffd\_\{\\mathrm\{eff\}\}\(99%99\\%variance threshold\)\. Vanilla VQGAN collapses todeff=2d\_\{\\mathrm\{eff\}\}=2and hits aKK\-independent loss floor regardless of codebook size\. Adding respawn \(No Warm\-up\) partially rescues codebook utilization and modestly raisesdeffd\_\{\\mathrm\{eff\}\}, but still saturates well above the warm\-up curves\. Each additional AE warm\-up phase further increases the codebook effective dimension and lowers the reconstruction floor, with the gains amplifying asKKgrows — consistent with the prediction that warm\-up, notKK, controls how many modes survive the VQ transition\.![Refer to caption](https://arxiv.org/html/2605.06870v1/figures/ae_warmup_scatter_v3_allsamples_nox_recloss.png)Figure 9:AE warm\-up vs\. codebook\-size scaling in VQGAN\.Final reconstruction metrics on the ImageNet\-100 validation set at step100,000100\{,\}000—L1 \(left\), LPIPS \(center\), and rFID \(right\)—as a function of AE warm\-up lengthTwu∈\{0,1​k,5​k,10​k,15​k,20​k,30​k,40​k\}T\_\{\\mathrm\{wu\}\}\\in\\\{0,1\\text\{k\},5\\text\{k\},10\\text\{k\},15\\text\{k\},20\\text\{k\},30\\text\{k\},40\\text\{k\}\\\}, for three codebook sizesK∈\{210,214,216\}K\\in\\\{2^\{10\},2^\{14\},2^\{16\}\\\}\.Twu=0T\_\{\\mathrm\{wu\}\}=0corresponds to the VQGAN w/ Respawn baseline \(kk\-means initialization plus dead\-code respawn, no AE phase\)\. AtTwu=0T\_\{\\mathrm\{wu\}\}=0the three curves are tightly clustered: increasingKKby64×64\\timesyields only a modest reduction in any of the three metrics\. AsTwuT\_\{\\mathrm\{wu\}\}grows, the curves separate into the expected orderingK=216<214<210K=2^\{16\}<2^\{14\}<2^\{10\}and the gap between them widens, indicating that warm\-up not only lowers the reconstruction floor but also restores the scaling benefit of a larger codebook\. The effect is cleanest in LPIPS, which decreases monotonically withTwuT\_\{\\mathrm\{wu\}\}at everyKK; L1 and rFID show the same overall trend but with more run\-to\-run noise\. Returns diminish pastTwu≈20​kT\_\{\\mathrm\{wu\}\}\\approx 20\\text\{k\}, matching the plateau of the AE\-phase effective\-dimension trajectory in Figure[7](https://arxiv.org/html/2605.06870#A2.F7)\.[Figure˜8](https://arxiv.org/html/2605.06870#A3.F8)shows the training trajectories for the experiments across all codebook sizes and warm\-up runs, and[Figure˜9](https://arxiv.org/html/2605.06870#A3.F9)shows the effect of the number of AE warm\-up steps on final reconstruction loss after 100k total training steps\. We observe that for all codebook sizes, the effect of adding warm\-up steps plateaus at around 20k steps, and increasing warm\-up duration past this point does not improve final loss\.

### C\.1Water\-filling upper\-bounds final VQ\-VAE

[Theorem˜1](https://arxiv.org/html/2605.06870#Thmtheorem1)bounds the surviving codebook dimension bymwu​\(Twu,R\)m\_\{\\text\{wu\}\}\(T\_\{\\text\{wu\}\},R\), the number of latent PCA components that lie above the Shannon water level at rateR=log2⁡KR=\\log\_\{2\}K\. The theorem is stated for the linear\-Gaussian RD\-AE; here we test whether the same construction transfers to nonlinear VQGAN\.

![Refer to caption](https://arxiv.org/html/2605.06870v1/figures/ae_1k_variances_waterfilling_all_cb_pca.png)1k steps
![Refer to caption](https://arxiv.org/html/2605.06870v1/figures/ae_20k_variances_waterfilling_all_cb_pca.png)10k steps
![Refer to caption](https://arxiv.org/html/2605.06870v1/figures/ae_40k_variances_waterfilling_all_cb_pca.png)40k steps

Figure 10:AE latent PCA spectrum and water levels at three warmup durations\.Gray bars: PCA eigenvalues of the pre\-quantization latent at AE warmup steps 1k, 10k, and 40k \(left to right\), computed over the full ImageNet\-100 \(20k\) training set\. Dashed horizontal lines: Shannon water levelD⋆D^\{\\star\}at ratesR=log2⁡KR=\\log\_\{2\}Kfor codebook sizesK∈\{210,214,216\}K\\in\\\{2^\{10\},2^\{14\},2^\{16\}\\\}, computed by reverse water\-filling \(Eq\.[8](https://arxiv.org/html/2605.06870#S3.E8)\) on the displayed spectrum\. The number of PCA components above each line is the predicted active\-mode countmwu​\(Twu,R\)m\_\{\\text\{wu\}\}\(T\_\{\\text\{wu\}\},R\)from[Theorem˜1](https://arxiv.org/html/2605.06870#Thmtheorem1)\. Two trends visible across panels: asTwuT\_\{\\text\{wu\}\}grows, more modes lift above each line \(weak modes activate during AE training\); at fixedTwuT\_\{\\text\{wu\}\}, largerKKpushes the water line lower and admits more modes\. The annotations report the resultingmwum\_\{\\text\{wu\}\}for each\(Twu,K\)\(T\_\{\\text\{wu\}\},K\)pair\.![Refer to caption](https://arxiv.org/html/2605.06870v1/figures/ae_warmup_scatter_v3_allsamples_nox_cbdim.png)Figure 11:Water\-filling on the AE latent spectrum upper\-bounds the trained VQ\-VAE codebook dimension\.For each AE warm\-up checkpointTwu∈\{0,1​k,5​k,10​k,15​k,20​k,30​k,40​k\}T\_\{\\text\{wu\}\}\\in\\\{0,1\\text\{k\},5\\text\{k\},10\\text\{k\},15\\text\{k\},20\\text\{k\},30\\text\{k\},40\\text\{k\}\\\}and each codebook sizeK∈\{210,214,216\}K\\in\\\{2^\{10\},2^\{14\},2^\{16\}\\\}\(one panel perKK\), we compare the water\-filling predictionmwu​\(Twu,R\)m\_\{\\text\{wu\}\}\(T\_\{\\text\{wu\}\},R\)from the AE PCA spectrum \(light, “WF theory”\) against the codebook effective dimension of the corresponding VQ\-VAE trained from that checkpoint to 100k total steps \(dark, “Trained”\)\. The predicted dimension lies above the trained dimension at every point and across all three codebook sizes: water\-filling on the AE latent is an empirical upper bound on what the downstream VQ\-VAE achieves\. Both quantities increase monotonically withTwuT\_\{\\text\{wu\}\}, and the gap between them is roughly constant — consistent with hard VQ paying a fixed efficiency cost relative to the rate\-distortion\-optimal channel\.##### Procedure\.

For each AE warm\-up checkpoint, we compute the PCA spectrum of the pre\-quantization latent over the full ImageNet\-100 \(20k\) training set — concretely, the eigenvalues of the centered latent covariance, obtained via SVD of the centered latent matrix\. We then run reverse water\-filling at ratesR∈\{10,14,16\}R\\in\\\{10,14,16\\\}bits, matching the codebook sizes used in[Table˜1](https://arxiv.org/html/2605.06870#S4.T1)\. The number of PCA components whose eigenvalue exceedsD⋆D^\{\\star\}gives the predicted active\-mode countmwum\_\{\\text\{wu\}\}\.[Figure˜10](https://arxiv.org/html/2605.06870#A3.F10)visualizes the spectrum and water lines at three checkpoints;[Figure˜11](https://arxiv.org/html/2605.06870#A3.F11)comparesmwum\_\{\\text\{wu\}\}to the trained VQ\-VAE codebook effective dimension at every\(Twu,K\)\(T\_\{\\text\{wu\}\},K\)pair\.

##### Findings\.

The water\-filling prediction sits strictly above the trained codebook dimension in all 24 settings, validating the theorem’s bound qualitatively in the nonlinear regime\. The bound also captures the right qualitative behavior in both axes: the predicted dimension grows withTwuT\_\{\\text\{wu\}\}\(weak modes activate during AE training and lift above the water line\) and withKK\(higher rate lowers the water line, admitting more modes\)\. Both trends are reproduced by the trained VQ\-VAEs\.

The gap between prediction and observation is roughly constant inTwuT\_\{\\text\{wu\}\}and shrinks slightly withKK, suggesting the looseness comes primarily from the hard\-VQ\-vs\-Shannon\-channel efficiency gap rather than from a breakdown of the linear analysis\.

### C\.2VQGAN withβ=0\\beta=0

Table 3:ImageNet\-100 \(20k\) warm\-up ablation, withβ=0\\beta=0baselines\.L1, LPIPS, and rFID: lower is better\. Dim: codebookdeffd\_\{\\text\{eff\}\}\(higher = more of the codebook span is used\)\. Bold indicates the best value within each codebook size\. Vanilla VQGAN uses uniform codebook init and no respawn; VQGAN w/ Respawn adds k\-means init and respawn;Twu=T\_\{\\text\{wu\}\}=\{1k, …, 40k\} additionally pre\-trains the encoder–decoder as a plain autoencoder for the indicated number of steps before introducing the codebook\. Theβ=0\\beta=0rows remove the commitment loss; without it, vanilla VQGAN collapses to near\-zero codebook utilization while the encoder is free to span many latent dimensions, and respawn alone partially compensates for utilization but still falls short of warm\-up at largeKK\. All runs share the same 100k\-step budget, so warm\-up trades VQ steps for AE steps\.![Refer to caption](https://arxiv.org/html/2605.06870v1/figures/ref128_comparison_100k_loss_dim_util_v3_withbeta0.png)Figure 12:VQGAN on ImageNet\-100 \(20k\) with the commitment term removed\.Training curves atβ=0\\beta=0\(vanilla and respawn variants, dashed\) overlaid on the standardβ\>0\\beta\>0runs from[Figure˜8](https://arxiv.org/html/2605.06870#A3.F8)\. Rows, top to bottom: validation reconstruction loss \(L1 \+ LPIPS\), codebook effective dimension, codebook utilization, latent effective dimension\. Columns: codebook sizesK∈\{210,214,216\}K\\in\\\{2^\{10\},2^\{14\},2^\{16\}\\\}\. Vanilla VQGAN atβ=0\\beta=0collapses to near\-zero codebook utilization at everyKK— almost all codewords are dead from initialization onward, since without the commitment term pulling the encoder toward the codebook, randomly initialized codes are never assigned\. Respawn atβ=0\\beta=0keeps codes alive and admits a higher latent effective dimension at smallKK\(consistent with the linear\-theory escape route in[Section˜3\.2](https://arxiv.org/html/2605.06870#S3.SS2)\), but does not translate to a lower reconstruction floor\.The linear analysis in[Section˜3](https://arxiv.org/html/2605.06870#S3)identifiedβ=0\\beta=0as a degenerate case in which inactive modes are not strictly frozen: atβ=0\\beta=0, the encoder ODE readsu˙j=2​σj2​vj\\dot\{u\}\_\{j\}=2\\sigma\_\{j\}^\{2\}v\_\{j\}, which is positive for anyvj\>0v\_\{j\}\>0, so inactive modes can in principle grow back into the active set\. The text claims that this escape is closed in practice by a combination of weight decay, normalization, and nonlinearity\.[Figure˜12](https://arxiv.org/html/2605.06870#A3.F12)and[Table˜3](https://arxiv.org/html/2605.06870#A3.T3)test this directly by removing the commitment term from VQGAN\.

##### Vanilla VQGAN atβ=0\\beta=0fails catastrophically\.

Without commitment loss, codebook utilization drops to 3\.7%, 0\.3%, and 0\.0% atK=210,214,216K=2^\{10\},2^\{14\},2^\{16\}respectively\. The codebook effective dimension is reported as 248–254, but this is a degenerate artifact: random codebook initialization spreads codewords across the embedding space, and with no commitment term to pull encoder outputs toward existing codes, almost no codes are ever assigned\. The unassigned codes remain at their random init values, inflating the PCA span without contributing to reconstruction\. L1 and LPIPS confirm the failure \(rFID climbs to 368 atK=216K=2^\{16\}\)\.

##### Respawn atβ=0\\beta=0: partial escape, no win\.

Withkk\-means initialization and dead\-code respawn,β=0\\beta=0does produce a higher latent effective dimension than the standard recipe at smallKK— codebook dim 22 atK=210K=2^\{10\}, versus 3 for theβ=0\.25\\beta=0\.25baseline\. This is the predicted escape: without the commitment term pulling the encoder toward the codebook’s current span, the encoder does spread into more directions\. However, the higher latent dimension does not translate into better reconstruction \(L1 0\.1566 vs\. 0\.1484 atK=210K=2^\{10\}\)\. At largerKK, even the dimensional advantage disappears; codebook dim reverts to∼\\sim5, comparable to the commitment\-on respawn baseline\.

##### Reading\.

The commitment term plays two coupled roles\. It*suppresses*weak modes in the encoder \(the mechanism behind collapse\) but it also*stabilizes*the encoder–codebook relationship by anchoring encoder outputs to codeword positions, which is what allows respawn to actually inject codewords into the encoder’s distribution\. Removing it eliminates the suppression but breaks the stabilization, and the net effect on reconstruction is negative\. AE warm\-up sidesteps this trade\-off entirely: weak modes activate*before*VQ is introduced, so neither role of the commitment term is in tension with collapse avoidance during the critical phase\. This is why warm\-up dominatesβ=0\\beta=0across every codebook size in[Table˜3](https://arxiv.org/html/2605.06870#A3.T3)\.

## Appendix DAdditional results on WavTokenizer

![Refer to caption](https://arxiv.org/html/2605.06870v1/figures/metrics_3panel.png)Figure 13:WavTokenizer on LibriTTS: cold\-start training fails to benefit from larger codebooks; AE Warm\-up restores codebook scaling\.Cold\-start WavTokenizer \(blue curves,K∈\{4​k,8​k,16​k\}K\\in\\\{4\\text\{k\},8\\text\{k\},16\\text\{k\}\\\}\) versus AE\-VQ \(warm curves,K∈\{4​k,8​k,16​k,65​k\}K\\in\\\{4\\text\{k\},8\\text\{k\},16\\text\{k\},65\\text\{k\}\\\}\)\.*From left:*validation mel loss, PESQ, codebookdeffd\_\{\\text\{eff\}\}\(99% variance\)\. AE Warm\-up runs are trained for 43 epochs before introducing VQ\. Cold\-start runs show flat loss and low effective dimension regardless ofKK; AE\-VQ runs show monotone improvement withKKand substantially higher effective dimension in both latent and codebook spaces\.Figure[13](https://arxiv.org/html/2605.06870#A4.F13)shows training trajectories for the WavTokenizer runs reported in Table[2](https://arxiv.org/html/2605.06870#S4.T2)\. Two observations are visible in the curves but not in the final\-checkpoint metrics:

##### The VQ transition is recoverable\.

At epoch 43, when AE Warm\-up runs introduce the quantizer, PESQ both degrades sharply for several epochs as the codebook is initialized and the encoder adapts to the quantization channel\. All AE Warm\-up runs recover within roughly 10 epochs and then continue to improve monotonically through the remaining budget, indicating that the cost of the warm\-to\-VQ transition is small relative to the gains it unlocks\.

##### Cold\-start improvement saturates early\.

Cold\-start runs reach their best validation Mel loss plateau after the first 20–30 epochs\. AE Warm\-up runs continue to improve Mel loss throughout post\-VQ training and select checkpoints from late in the run\. The additional training budget is productive only when the latent has enough effective dimension to absorb it\.

## NeurIPS Paper Checklist

1. 1\.Claims
2. Question: Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?
3. Answer:\[Yes\]
4. Justification: The abstract and Section 1 \(Contributions\) state our four claims—demonstration of dimensional collapse in VQGAN/WavTokenizer \(Section 4\), the AE Warm\-up fix \(Section 3\.3, Section 4\), the RD\-AE theoretical model \(Section 3\.1–3\.2\), and the warm\-up duration bound \(Theorem 1, Section 3\.3\)—each of which is directly supported by the corresponding section\.
5. Guidelines: - •The answer\[N/A\]means that the abstract and introduction do not include the claims made in the paper\. - •The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations\. A\[No\]or\[N/A\]answer to this question will not be perceived well by the reviewers\. - •The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings\. - •It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper\.
6. 2\.Limitations
7. Question: Does the paper discuss the limitations of the work performed by the authors?
8. Answer:\[Yes\]
9. Justification: Section 5 \(Discussion and Limitations\) discusses the linearity and Gaussian\-input assumptions of our theory, the gap between RD\-AE and hard VQ, and open questions about combining warm\-up with other codebook\-geometry techniques\. Appendix C\.2 additionally discusses theβ=0\\beta=0degenerate case where the linear theory’s predictions diverge from practice\.
10. Guidelines: - •The answer\[N/A\]means that the paper has no limitation while the answer\[No\]means that the paper has limitations, but those are not discussed in the paper\. - •The authors are encouraged to create a separate “Limitations” section in their paper\. - •The paper should point out any strong assumptions and how robust the results are to violations of these assumptions \(e\.g\., independence assumptions, noiseless settings, model well\-specification, asymptotic approximations only holding locally\)\. The authors should reflect on how these assumptions might be violated in practice and what the implications would be\. - •The authors should reflect on the scope of the claims made, e\.g\., if the approach was only tested on a few datasets or with a few runs\. In general, empirical results often depend on implicit assumptions, which should be articulated\. - •The authors should reflect on the factors that influence the performance of the approach\. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting\. Or a speech\-to\-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon\. - •The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size\. - •If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness\. - •While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren’t acknowledged in the paper\. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community\. Reviewers will be specifically instructed to not penalize honesty concerning limitations\.
11. 3\.Theory assumptions and proofs
12. Question: For each theoretical result, does the paper provide the full set of assumptions and a complete \(and correct\) proof?
13. Answer:\[Yes\]
14. Justification: All assumptions are stated explicitly in Section 3\.1–3\.2 and in the statement of Theorem 1\. Full proofs are provided in Appendix A, with a proof sketch given in the main text\. Numerical validation of Theorem 1 against simulation appears in Appendix A\.6\.1\.
15. Guidelines: - •The answer\[N/A\]means that the paper does not include theoretical results\. - •All the theorems, formulas, and proofs in the paper should be numbered and cross\-referenced\. - •All assumptions should be clearly stated or referenced in the statement of any theorems\. - •The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition\. - •Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material\. - •Theorems and Lemmas that the proof relies upon should be properly referenced\.
16. 4\.Experimental result reproducibility
17. Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper \(regardless of whether the code and data are provided or not\)?
18. Answer:\[Yes\]
19. Justification: Appendix B\.1 and Appendix B\.2 provide complete training details for VQGAN and WavTokenizer respectively\. The simulation details for the RD\-AE numerical experiments are stated in the captions of Figure 4 and Figure 6\. An anonymized code repository is also provided \(footnote 1\)\.
20. Guidelines: - •The answer\[N/A\]means that the paper does not include experiments\. - •If the paper includes experiments, a\[No\]answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not\. - •If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable\. - •Depending on the contribution, reproducibility can be accomplished in various ways\. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model\. In general\. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model \(e\.g\., in the case of a large language model\), releasing of a model checkpoint, or other means that are appropriate to the research performed\. - •While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution\. For example 1. \(a\)If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm\. 2. \(b\)If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully\. 3. \(c\)If the contribution is a new model \(e\.g\., a large language model\), then there should either be a way to access this model for reproducing the results or a way to reproduce the model \(e\.g\., with an open\-source dataset or instructions for how to construct the dataset\)\. 4. \(d\)We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility\. In the case of closed\-source models, it may be that access to the model is limited in some way \(e\.g\., to registered users\), but it should be possible for other researchers to have some path to reproducing or verifying the results\.
21. 5\.Open access to data and code
22. Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?
23. Answer:\[Yes\]
25. Guidelines: - •The answer\[N/A\]means that paper does not include experiments requiring code\. - • - •While we encourage the release of code and data, we understand that this might not be possible, so\[No\]is an acceptable answer\. Papers cannot be rejected simply for not including code, unless this is central to the contribution \(e\.g\., for a new open\-source benchmark\)\. - •The instructions should contain the exact command and environment needed to run to reproduce the results\. See the NeurIPS code and data submission guidelines \([https://neurips\.cc/public/guides/CodeSubmissionPolicy](https://neurips.cc/public/guides/CodeSubmissionPolicy)\) for more details\. - •The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc\. - •The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines\. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why\. - •At submission time, to preserve anonymity, the authors should release anonymized versions \(if applicable\)\. - •Providing as much information as possible in supplemental material \(appended to the paper\) is recommended, but including URLs to data and code is permitted\.
26. 6\.Experimental setting/details
27. Question: Does the paper specify all the training and test details \(e\.g\., data splits, hyperparameters, how they were chosen, type of optimizer\) necessary to understand the results?
28. Answer:\[Yes\]
29. Justification: Section 4\.1 and 4\.2 give the high\-level setup \(architecture, codebook sizes, training budget, baselines\) for VQGAN and WavTokenizer respectively\. Full details are provided in Appendix B\.1 \(VQGAN\) and Appendix B\.2 \(WavTokenizer\)\.
30. Guidelines: - •The answer\[N/A\]means that the paper does not include experiments\. - •The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them\. - •The full details can be provided either with the code, in appendix, or as supplemental material\.
31. 7\.Experiment statistical significance
32. Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
33. Answer:\[Yes\]
34. Justification: RD\-AE simulations report the median over 128 seeds with cross\-seed variance below10−310^\{\-3\}\(Figure 4 caption\)\. VQGAN and WavTokenizer results are single\-seed due to compute cost \(5\.5 hours and 7 days per run, respectively\); however, the consistent monotonic trends across 3 codebook sizes×\\times7 warm\-up durations for VQGAN \(21 runs total\) and across 4 codebook sizes for WavTokenizer establish the empirical pattern, and qualitative agreement across two modalities corroborates the theoretical predictions\.
35. Guidelines: - •The answer\[N/A\]means that the paper does not include experiments\. - •The authors should answer\[Yes\]if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper\. - •The factors of variability that the error bars are capturing should be clearly stated \(for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions\)\. - •The method for calculating the error bars should be explained \(closed form formula, call to a library function, bootstrap, etc\.\) - •The assumptions made should be given \(e\.g\., Normally distributed errors\)\. - •It should be clear whether the error bar is the standard deviation or the standard error of the mean\. - •It is OK to report 1\-sigma error bars, but one should state it\. The authors should preferably report a 2\-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified\. - •For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range \(e\.g\., negative error rates\)\. - •If error bars are reported in tables or plots, the authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text\.
36. 8\.Experiments compute resources
37. Question: For each experiment, does the paper provide sufficient information on the computer resources \(type of compute workers, memory, time of execution\) needed to reproduce the experiments?
38. Answer:\[Yes\]
39. Justification: Appendix B\.1 specifies the hardware used for VQGAN runs and training times\. Appendix B\.2 reports the same for WavTokenizer\.
40. Guidelines: - •The answer\[N/A\]means that the paper does not include experiments\. - •The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage\. - •The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute\. - •The paper should disclose whether the full research project required more compute than the experiments reported in the paper \(e\.g\., preliminary or failed experiments that didn’t make it into the paper\)\.
41. 9\.Code of ethics
43. Answer:\[Yes\]
44. Justification: The research conforms with the NeurIPS Code of Ethics\. The work uses publicly available datasets \(ImageNet\-100, LibriTTS\), does not involve human subjects or personally identifiable information, and the submission is anonymized\.
45. Guidelines: - •The answer\[N/A\]means that the authors have not reviewed the NeurIPS Code of Ethics\. - •If the authors answer\[No\], they should explain the special circumstances that require a deviation from the Code of Ethics\. - •The authors should make sure to preserve anonymity \(e\.g\., if there is a special consideration due to laws or regulations in their jurisdiction\)\.
46. 10\.Broader impacts
47. Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
48. Answer:\[N/A\]
49. Justification: This is foundational research on the training dynamics of vector\-quantized autoencoders\. The contribution is an analytical framework and training\-recipe improvement for an existing class of models; it does not introduce new immediate societal impacts beyond those associated with improving generative modeling more broadly\.
50. Guidelines: - •The answer\[N/A\]means that there is no societal impact of the work performed\. - •If the authors answer\[N/A\]or\[No\], they should explain why their work has no societal impact or why the paper does not address societal impact\. - •Examples of negative societal impacts include potential malicious or unintended uses \(e\.g\., disinformation, generating fake profiles, surveillance\), fairness considerations \(e\.g\., deployment of technologies that could make decisions that unfairly impact specific groups\), privacy considerations, and security considerations\. - •The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments\. However, if there is a direct path to any negative applications, the authors should point it out\. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate Deepfakes for disinformation\. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster\. - •The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from \(intentional or unintentional\) misuse of the technology\. - •If there are negative societal impacts, the authors could also discuss possible mitigation strategies \(e\.g\., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML\)\.
51. 11\.Safeguards
52. Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse \(e\.g\., pre\-trained language models, image generators, or scraped datasets\)?
53. Answer:\[N/A\]
54. Justification: The paper does not release any new pretrained models, generative systems, or scraped datasets\.
55. Guidelines: - •The answer\[N/A\]means that the paper poses no such risks\. - •Released models that have a high risk for misuse or dual\-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters\. - •Datasets that have been scraped from the Internet could pose safety risks\. The authors should describe how they avoided releasing unsafe images\. - •We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort\.
56. 12\.Licenses for existing assets
57. Question: Are the creators or original owners of assets \(e\.g\., code, data, models\), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
58. Answer:\[Yes\]
59. Justification: All existing assets used—VQGAN architecture, WavTokenizer architecture, ImageNet\-100 split, LibriTTS, LPIPS, PESQ, STOI, UTMOS, and CREPE—are cited in the references and used within their respective licenses for academic research\. Architectural and training details inherited from the released codebases are noted explicitly in Appendix B\.1 and B\.2\.
60. Guidelines: - •The answer\[N/A\]means that the paper does not use existing assets\. - •The authors should cite the original paper that produced the code package or dataset\. - •The authors should state which version of the asset is used and, if possible, include a URL\. - •The name of the license \(e\.g\., CC\-BY 4\.0\) should be included for each asset\. - •For scraped data from a particular source \(e\.g\., website\), the copyright and terms of service of that source should be provided\. - •If assets are released, the license, copyright information, and terms of use in the package should be provided\. For popular datasets,[paperswithcode\.com/datasets](https://arxiv.org/html/2605.06870v1/paperswithcode.com/datasets)has curated licenses for some datasets\. Their licensing guide can help determine the license of a dataset\. - •For existing datasets that are re\-packaged, both the original license and the license of the derived asset \(if it has changed\) should be provided\. - •If this information is not available online, the authors are encouraged to reach out to the asset’s creators\.
61. 13\.New assets
62. Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
63. Answer:\[Yes\]
64. Justification: The new asset introduced is the experimental code, released via an anonymized repository \(footnote 1\)\. The repository contains training scripts, configuration files, and evaluation code, with documentation describing how to reproduce the VQGAN and WavTokenizer experiments reported in Section 4\.
65. Guidelines: - •The answer\[N/A\]means that the paper does not release new assets\. - •Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates\. This includes details about training, license, limitations, etc\. - •The paper should discuss whether and how consent was obtained from people whose asset is used\. - •At submission time, remember to anonymize your assets \(if applicable\)\. You can either create an anonymized URL or include an anonymized zip file\.
66. 14\.Crowdsourcing and research with human subjects
67. Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation \(if any\)?
68. Answer:\[N/A\]
69. Justification: The paper does not involve crowdsourcing or research with human subjects\. All evaluations use automated metrics on publicly available datasets\.
70. Guidelines: - •The answer\[N/A\]means that the paper does not involve crowdsourcing nor research with human subjects\. - •Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper\. - •According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector\.
71. 15\.Institutional review board \(IRB\) approvals or equivalent for research with human subjects
72. Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board \(IRB\) approvals \(or an equivalent approval/review based on the requirements of your country or institution\) were obtained?
73. Answer:\[N/A\]
74. Justification: The paper does not involve human subjects research, so IRB approval is not applicable\.
75. Guidelines: - •The answer\[N/A\]means that the paper does not involve crowdsourcing nor research with human subjects\. - •Depending on the country in which research is conducted, IRB approval \(or equivalent\) may be required for any human subjects research\. If you obtained IRB approval, you should clearly state this in the paper\. - •We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution\. - •For initial submissions, do not include any information that would break anonymity \(if applicable\), such as the institution conducting the review\.
76. 16\.Declaration of LLM usage
77. Question: Does the paper describe the usage of LLMs if it is an important, original, or non\-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does*not*impact the core methodology, scientific rigor, or originality of the research, declaration is not required\.
78. Answer:\[N/A\]
79. Justification: LLMs were not used as a component of the core methodology\. The theoretical analysis, RD\-AE framework, AE Warm\-up method, and experimental design do not involve LLMs\.
80. Guidelines: - •The answer\[N/A\]means that the core method development in this research does not involve LLMs as any important, original, or non\-standard components\. - •Please refer to our LLM policy in the NeurIPS handbook for what should or should not be described\.