Softmax Reparameterization for Output-Head Quantization

Softmax Reparameterization for Output-Head Quantization. By Asim Kadav and seven others, from Adobe. Published on the twenty-fifth of September. No code repository is linked to it.

In small language models, large vocabularies make the output head a substantial part of inference cost. Quantizing that head is the obvious lever, but baseline quantization can substantially distort the model's predictions, and some heads suffer much more than others.

Softmax reparameterization is a post-training method that picks a functionally equivalent output head before quantization happens. It subtracts a scalar multiple of the vocabulary-row mean from every output row. That shift preserves the full-precision softmax distribution, so in full precision nothing about the model's predictions changes, and the trained decoder is left as it was. What the choice of coefficient does affect is the head that the quantizer is handed.

The coefficient is chosen by a one-dimensional search scored by validation KL, run separately for each of three quantizers: RTN, activation-weighted MSE, and full-Hessian GPTQ. The search includes the original head, and it includes fixed mean-centering, so both remain options the search can land on. Where a model routes logits through a nonlinear path, such as soft-capping, a rank-one correction handles it.

Across seven heads at W4, the gains concentrate where baseline quantization substantially distorts predictions. On Phi-4-mini, with activation-weighted MSE, KL falls from zero point nine three six to zero point two five six. The gains survive stronger GPTQ calibration, and they remain complementary to exact per-channel scaling and to affine quantization.

The coefficients also transfer. Selected on WikiText and then frozen, they carry over to C4 and OpenWebMath. Across four heads and three W4 quantizers, they outperform mean-centering in all eighteen comparisons where the frozen coefficient differs from one, and they match it in the remaining six. At W2, used as a compression stress test, the benefits broaden across nearly the full matrix of models and quantizers.

A matched residual analysis adds a subtlety. Improved fidelity can come with greater logit reconstruction error, while the residual's Fisher-weighted cost goes down. So a larger raw error need not mean a worse head. On efficiency, for shift-compatible heads the reparameterization adds no inference operation and keeps packed W4 execution. With the decoder held in BF16, quantizing the Phi output head cuts batch-one generation latency by ten point eight percent, relative to the baseline with a BF16 head.

The method rests on a freedom the softmax already has. Shifting every output row by a multiple of the row mean leaves the distribution alone, so there is a whole line of heads that behave identically before quantization. The method spends that freedom on surviving quantization, searches a single number to do it, and for heads compatible with the shift, costs nothing at inference.