Gumbel-Softmax or a vector-quantized codebook for a discrete latent: how do you choose?

## What the decision is actually about A discrete latent admits no pathwise gradient, so every option buys trainability by giving something up. There are three families, and picking between them is a systems call, not a formula. 1. **Relax the sample.** Add parameter-free noise to the logits and replace the argmax with a tempered softmax. Low-variance gradients, biased objective, a temperature schedule to maintain. 2. **Quantize against a codebook.** Keep the forward pass discrete and route the gradient around the non-differentiable step. 3. **Keep the sample exact and use a score-function estimator.** Unbiased on the real discrete objective, but the variance is severe enough that a baseline and other variance reduction become part of the model, not an optimisation. ## How the codebook route works The encoder emits a continuous vector. You hold a table of learnable embedding vectors and replace the encoder's output with the nearest entry by Euclidean distance; the decoder sees only that entry. Backward, the lookup is a hard assignment with no useful derivative, so the decoder's gradient at the chosen entry is copied unchanged onto the encoder's output — a straight-through pass around the lookup. This is deliberately the gradient of a different function than the one computed, and it works because the chosen entry is, by construction, close to the encoder's output. Two extra terms hold the arrangement together. A codebook term pulls the chosen entry toward the encoder output that selected it, with the encoder output treated as a constant in that term; many implementations replace it with an exponential moving average of the encoder outputs assigned to each entry, which is more stable. A commitment term pulls the encoder output toward its chosen entry, with the entry treated as a constant, weighted by a coefficient. Without commitment, encoder outputs drift away from the table faster than the table can follow, and the straight-through approximation stops being justified. **Its failure mode is codebook collapse.** Entries that are never the nearest to anything receive no update and stay where they were initialised, so usage concentrates on a small subset and the effective code space is far smaller than the table you paid for. Detect it by logging, per epoch, the fraction of entries selected at least once and the entropy of the usage histogram. The standard mitigations are exponential-moving-average updates, re-initialising dead entries from recent encoder outputs, initialising the table by clustering encoder outputs, normalising codes, and reducing the code dimension so distances are better behaved. ## The axes that actually decide it **Must the forward pass be discrete?** If the decoder or any downstream consumer must never see a blend, the codebook is discrete from the first step, whereas the relaxation only approaches discreteness. A straight-through relaxation partly closes that gap, but at the cost of a gradient still further from the computed function. **How many classes?** A relaxation needs noise and a normalisation over every class. Over a 10,000-item catalogue, an early soft sample is a weighted average across all items — a vector resembling nothing real. A codebook of a few hundred to a few thousand entries sidesteps this entirely, because nearest-neighbour lookup returns one real vector regardless of table size. **Who consumes the code?** If the plan is to fit a prior over the latent, integer indices from a codebook are directly modellable; relaxed simplex vectors are not. This is usually the argument that ends the discussion. **Which failure would you rather operate?** Relaxation risks a schedule that sharpens too fast, freezing the assignment before the model learns, and a quality drop when inference switches to a hard choice. Codebooks risk collapse, which is silent unless you instrument usage. Both are recoverable; the question is which one your team can see. **How much bias can the objective tolerate?** If the task genuinely requires an unbiased gradient of the discrete objective — because the objective is a hard, non-differentiable measurement rather than a reconstruction — neither approach qualifies and the score-function estimator with a baseline is the honest choice, with the variance budget planned in. ## A defensible default For a perceptual model whose codes feed a later sequence model, start with a codebook, use moving-average table updates, instrument entry usage from day one, and treat collapse as a monitored condition rather than a surprise. For a small discrete choice inside a larger differentiable pipeline — a handful of routes, a modest vocabulary — start with a relaxation, because it is fewer moving parts and the annealing schedule is the only new hyperparameter. Reverse either call when the evidence you instrumented for shows up: a collapsed table, or a training-to-inference quality cliff.