Chapter 11: Matrix Calculus. In the chapter mind map this icon labels STE & Subgradients. The discussion below is excerpted and lightly edited from § Subgradients and the Straight-Through Estimator in Mathematics for AI and Machine Learning.

Discrete operations, such as quantization in Vector Quantized Variational Autoencoders (VQ-VAE), present a fundamental challenge: they are not differentiable anywhere in their domain. The Straight-Through Estimator (STE) is a gradient approximation technique that enables backpropagation through such operations.

Consider a quantization operation $Q : \mathbb{R} \to \mathcal{Q}$ that maps continuous values to a discrete codebook $\mathcal{Q} = \{q_0, q_1, \ldots, q_{k-1}\}$:

The forward pass uses the quantized value $\hat{x}$. During backpropagation, the STE approximates the gradient by treating the quantization operation as the identity function:

This corresponds to setting $\frac{\partial \hat{x}}{\partial x} = 1$ in the chain rule, even though the true derivative is zero (or undefined) everywhere.

The STE is motivated by the observation that, while quantization is discrete, the loss function $\mathcal{L}(\hat{x})$ is typically smooth with respect to the quantized values. The gradient $\frac{\partial \mathcal{L}}{\partial \hat{x}}$ provides directional information for updating $x$ to reduce the loss, despite the non-differentiability of the quantization step itself. This heuristic has proven effective in practice for training models with discrete bottlenecks.

For a vector quantization operation $\mathbf{Q} : \mathbb{R}^N \to \mathcal{Q}^N$, the STE approximation extends to: