This was the output for a question about encouraging a VQ-VAE to use more codes in a codebook
Why do only ≈60 of 512 codes become “alive”? The encoder is free to collapse to any subset of vectors that minimizes the reconstruction loss. If it can achieve a low error with 60 vectors, nothing in the objective forces it to use the rest, and once a code has not been selected for a few thousand steps its embedding stops moving (the gradient through the straight-through estimator is zero), so it never comes back to life.
Below are the main levers practitioners use to fight this collapse and obtain a larger, more informative codebook. They can be mixed and matched.
──────────────────────────────────────────────────────────────────────────────
Use the EMA (“VQ-VAE-v2”) update instead of the gradient update
──────────────────────────────────────────────────────────────────────────────
With the original loss-based update, an embedding only moves when it is chosen,
so dead codes stay dead. The exponential–moving-average update keeps the codebook
closer to the running K-means solution and makes it much harder for vectors to die.
N_i ← γ N_i + (1-γ) n_i # running counts
m_i ← γ m_i + (1-γ) Σ_enc(z)_i
e_i ← m_i / N_i # codebook entry
Simply switching to the EMA variant usually boosts “perplexity” (number of effective
codes) dramatically.
────────────────────────────────────────────────────────────────────────────── 2. Re-initialise or “respawn” dead codes ────────────────────────────────────────────────────────────────────────────── Every K steps, check how many times each code has been used in the last window. If usage < τ (e.g. <10), sample a random bottleneck vector from the current minibatch (or from a large buffer) and overwrite that code with it. This keeps the codebook size fixed but continuously recycles unused slots.
────────────────────────────────────────────────────────────────────────────── 3. Increase the commitment cost β (or the overall VQ loss weight) ────────────────────────────────────────────────────────────────────────────── Loss = L_recon + β ‖sg[z_e] – e‖² + ‖z_e – sg[e]‖²
• A larger β pushes the encoder outputs closer to the nearest embedding, so different inputs must move to different embeddings to avoid clashes. • Monitor perplexity while sweeping β; values 0.25–2.0 are common, but audio often needs β≈2–4 before the codebook fills up.
────────────────────────────────────────────────────────────────────────────── 4. Add an explicit “diversity” or entropy regulariser ────────────────────────────────────────────────────────────────────────────── L_div = λ · KL( q(c) || U(K) ) where q(c) is the empirical code frequency from the current batch and U(K) is the uniform distribution over K codes. A lighter-weight alternative is −λ·H(q), i.e. maximise batch-wise entropy.
This directly rewards using many codes; λ must be tuned so that reconstruction quality does not degrade too much.
────────────────────────────────────────────────────────────────────────────── 5. Increase the batch size (or accumulate code-usage statistics) ────────────────────────────────────────────────────────────────────────────── With a 32-frame spectrogram and batch size 16, you expose only 512 bottleneck vectors per step. The probability that every code is observed at least once in the early phase is tiny. Double the batch or accumulate N_i, m_i over several steps before you apply the EMA update.
────────────────────────────────────────────────────────────────────────────── 6. Curriculum / grow-the-codebook training ────────────────────────────────────────────────────────────────────────────── a) Train with K_small (e.g. 64) until convergence. b) Copy those 64 vectors into a K_big=512 table, jitter them with small noise, and continue training. The already-useful vectors stay alive and the free slots can specialise.
────────────────────────────────────────────────────────────────────────────── 7. Lower the embedding dimensionality, then raise it again ────────────────────────────────────────────────────────────────────────────── When d (embedding dim) is large, many codes become almost collinear, so the nearest-neighbour search picks the same few. Train first with a small d (e.g. 32 instead of 128), then increase d and finetune. This acts as a form of capacity bottleneck that forces dispersion.
────────────────────────────────────────────────────────────────────────────── 8. Better initialisation (K-means seeding) ────────────────────────────────────────────────────────────────────────────── Run a few thousand encoder passes without the decoder and collect latent vectors; run K-means on them and use the centroids to initialise the codebook. You start with all codes “in range” of the encoder output distribution, which sharply reduces early death.
────────────────────────────────────────────────────────────────────────────── 9. Architectural tricks ────────────────────────────────────────────────────────────────────────────── • Split or Grouped Codebooks: Instead of one K=b512 table, use G groups of size K_g (e.g. 8 groups × 64 codes). The product space has K_g^G possible combinations, but each group is easier to populate. (See “Residual Vector Quantisation”, “Decomposed VQ”, Jukebox, etc.)
• Hierarchical VQ (VQ-VAE-2): a high-level codebook with coarse information and a lower-level fine codebook. Each level operates on simpler distributions and tends to use more codes.
────────────────────────────────────────────────────────────────────────────── 10. Train longer, with smaller learning rate for the codebook ────────────────────────────────────────────────────────────────────────────── If the encoder/decoder learn much faster than the codebook, they may settle early and freeze most vectors out. A slower but longer update schedule lets latecomer vectors find a niche.
────────────────────────────────────────────────────────────────────────────── Practical recipe to try first ──────────────────────────────────────────────────────────────────────────────
Switch to EMA codebook updates (if not already).
Raise β until perplexity keeps growing; start with β=2.0.
Use a batch size ≥256 time steps (or gradient-accumulate).
Re-spawn codes with usage<5 every 2000 steps.
If still below ~200 active codes after 100k steps, add an entropy loss.
This usually takes a 60/512 situation to >300/512 without harming reconstruction quality, and from there you can decide whether even more capacity is beneficial for your downstream task.