Fake Quantization and the Straight-Through Estimator

Post-training quantization chooses a representation for a network that was trained without one. The network’s weights are whatever the optimizer left them at, and the quantizer must map them onto a grid. When the grid is coarse enough, that mapping is not a small perturbation: a weight that needs to move by a third of a step to land on a better code will not move, because the training has ended.

Quantization-aware training moves the grid into the training loop. The forward pass computes with quantized values; the backward pass updates the same floating-point parameters as before. The parameters stay continuous and the network learns to place them where the grid can represent them well. The mechanism that makes this possible is a substitution for the derivative of rounding, and the consequences of that substitution are what this article is about.

Rounding has no gradient worth having

The forward function of an affine quantizer is $q = \mathrm{round}(x/s) + z$ followed by clamping. Differentiated with respect to $x$, the round function has derivative zero almost everywhere and undefined at the half-integer points. A training step that used the true derivative would send no gradient to any parameter that passes through the quantizer, and the layers below it would never update.

A softer version of the same problem appears with the scale. The scale is a parameter that training might want to adjust, and the derivative of $q$ with respect to $s$ is also zero almost everywhere under exact rounding, because the code does not change as the scale moves within an interval that keeps every value on the same side of its nearest half-integer. The scale would receive no gradient at all, and a network that needed a wider range would have no way to signal it.

Both problems have the same shape: the operation is piecewise constant, and the information a gradient carries is about infinitesimal change, which a piecewise-constant function discards.

The substitution

The straight-through estimator replaces the derivative of the quantizer with the identity. In the forward direction the quantized value is used, exactly as it will be at inference. In the backward direction the gradient arriving at the quantizer’s output is passed through unchanged to its input, as though the quantizer were the identity function.

Formally, if $y = Q(x)$ in the forward pass, the backward pass computes $\partial L/\partial x = \partial L/\partial y$. The quantizer is treated as having derivative one everywhere, including in the clipped region where the forward function is genuinely flat.

The justification is not that the derivative is correct; it is not. It is that the identity is the closest thing to a useful direction. A small change in $x$ usually does not change $Q(x)$, but when it does, the change in the output is one step. The identity says the expected sensitivity is about one output unit per input unit, which overstates the true sensitivity by the ratio of the step to the typical distance between values, and understates the effect of a change large enough to cross a boundary. It is a biased estimator with the right sign in the common case, and training with it converges in practice for 8-bit and 4-bit quantization.

import torch

class FakeQuant(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, scale, zero_point, qmin, qmax):
        q = torch.clamp(torch.round(x / scale) + zero_point, qmin, qmax)
        return (q - zero_point) * scale          # dequantized, same dtype as x

    @staticmethod
    def backward(ctx, grad_out):
        # Identity in the backward direction; scale and zero point get no gradient here.
        return grad_out, None, None, None, None

fake_quant = FakeQuant.apply

The forward returns the dequantized value rather than the integer, so that the rest of the network continues to compute in floating point. The rounding error is present in the forward value, which is the point: the loss the network sees includes the error the quantizer will introduce, and the optimizer reduces it.

What the estimator does to each parameter

The weight update is straightforward. A weight’s gradient passes through the quantizer unchanged, so the weight moves in the direction that would reduce the loss if it were unquantized. Whether the move changes the quantized value depends on whether it crosses a rounding boundary, and the step size relative to the grid spacing determines how often that happens. Training with a learning rate much smaller than the grid spacing produces many updates that change nothing in the forward pass, which wastes steps but does not corrupt the result. Training with a learning rate much larger than the spacing produces updates that jump several codes, which overshoots.

The scale update is where the substitution is least satisfactory. With the identity applied to the path from $x$ to the output, the scale receives gradient only through the dequantization multiply, not through the rounding decision. In the forward function written above, the scale appears twice: inside the round and again in the final multiply. The derivative through the final multiply is well defined; the derivative through the rounding is zero. So the scale is updated by a signal that says “make the dequantized value track the gradient” but not “change the grid to fit the data”.

Different implementations handle this differently, and the difference is visible in the trained result. One common approach computes the scale and zero point from running statistics of the tensor rather than by gradient descent, so they are not parameters at all: the quantizer observes the minimum and maximum of the tensor during training, smooths them with a moving average, and uses the resulting range. Another treats the scale as a learnable parameter with its own gradient path and a separate, often smaller, learning rate.

The statistics approach has a failure mode: the moving average lags the parameter movement. As training shifts a weight distribution, the observed range follows it with a delay set by the averaging constant, and the quantizer spends part of training operating on a range that no longer matches the tensor. The effect is a transient accuracy dip that resolves if training continues long enough, and it is the reason learned scales are sometimes preferred for the last part of training.

Clipping and the gradient that should stop

The identity substitution is applied uniformly, including where the value is clipped. In the clipped region the forward function is constant, so the true derivative is zero and the identity is maximally wrong: it tells a weight that its value matters when the quantizer has already discarded it.

The consequence is that weights pushed beyond the range keep receiving gradient that pushes them further, because the loss does not change when they move. They drift outward, the range statistics widen to follow them, and the quantizer’s effective resolution degrades. The mechanism is a positive feedback loop between clipping and range estimation, and it is why implementations often disable the gradient through the clamp separately, returning zero for elements that were clipped in the forward pass.

Doing that requires the backward function to know which elements were clipped, which means saving a mask from the forward pass. The extra memory is one bit per element, and the extra complexity is the reason some implementations skip it and accept the drift.

Where the quantizer sits in the graph

A quantized inference graph has quantizers at specific places: after each weight tensor, after each activation that feeds a convolution or a matrix multiply, and sometimes after the accumulator. Training must simulate the same placements, or the network trains against a different function than it will run.

The activation quantizer is the harder one, because it sits on a path that processes every token and its range depends on the input. A weight quantizer sees a fixed tensor whose distribution moves slowly; an activation quantizer sees a distribution that changes with each batch. Simulating it with a per-tensor range computed from a calibration set is what most post-training flows do, and the training-time analogue is a range that is estimated online from the batches seen so far.

The difference matters at inference when the range is computed dynamically per batch rather than statically. A network trained with a static simulated range and then run with a dynamic per-batch range sees a different quantizer than the one it trained against, and the mismatch is a source of degradation that no amount of training fixes. The remedy is to simulate the dynamic behaviour during training, which means the simulated quantizer must compute its range from the current batch, exactly as inference will.

What quantization-aware training buys and what it costs

The gain is largest at low bit widths and for small models. At 8 bits, post-training quantization with a well-chosen range usually recovers most of the accuracy of the floating-point model, and training adds little. At 4 bits and below, the grid is coarse enough that the placement of weights relative to grid points matters, and training moves weights toward positions the grid can represent. The gain is not free: the training run is comparable in cost to the original training, it requires a representative dataset, and the resulting model is tied to one specific quantizer configuration.

That last point is the operational cost. A model trained with a per-channel 4-bit weight quantizer and a per-tensor 8-bit activation quantizer cannot be re-quantized to a different configuration without retraining, because the weights have been placed for the grid they were trained against. Post-training methods have the opposite property: they can be applied at several bit widths from one checkpoint, and the choice can be made at deployment time.

The interaction with the rest of the cluster is direct. Training-time simulation must use the same rounding rule as the inference kernel, or the error the network optimized against is not the error it will experience; the requantization epilogue defines that rule. The simulated range must be chosen by the same criterion as the deployed calibration, or the network learns to tolerate clipping at a threshold it will not be given. And for the very low bit widths where training helps most, the rounding direction itself becomes a design variable rather than a fixed nearest-neighbour rule, which is where error-compensated rounding departs from the estimator described here.

The estimator is a stand-in for a derivative that does not exist. It works because the loss surface, viewed at the resolution of the grid, is close enough to the loss surface viewed continuously that a gradient computed on the latter points downhill on the former. That approximation degrades as the grid coarsens, and the point at which it stops working is where quantization becomes a search problem over discrete assignments rather than an optimization over continuous ones.