Softmax Reparameterization for Output-Head Quantization

A post-training method called softmax reparameterization selects a functionally equivalent output head before quantization to cut the inference cost of large-vocabulary output heads in small language models, according to the researchers behind the proposal. The technique subtracts a scalar multiple of the vocabulary-row mean from every output-head row, leaving the softmax distribution unchanged while making the head easier to quantize.

Large vocabularies make output heads a substantial inference cost in small language models. We propose softmax reparameterization, a post-training method that selects a functionally equivalent output head before quantization. The method subtracts a scalar multiple of the vocabulary-row mean from eve