Softmax Reparameterization for Output-Head Quantization

1 INTRODUCTION

• For vocabulary size V and hidden width d, the dense output projection contains V d weights and evaluates the full vocabulary at every decoding step.Unlike input embeddings, which retrieve selected rows, the output projection reads the full head.Source: §1 · p. 1
• The method searches a scalar family of exactly equivalent output-head parameterizations and selects the coefficient by validation KL separately for each base quantizer.The search includes the original head and ordinary mean-centering, uses pretrained weights, and leaves the decoder unchanged.Source: §1 · p. 2
• Function-preserving quantization methods include channel scaling and rotations, while this contribution is a post-training scalar search within the output head’s additive softmax equivalence class.The paper contrasts this search with fixed mean-centering and concurrent symmetry-based methods.Source: §1 · p. 2
• Reparameterization improves nearly every model–quantizer combination across RTN, AW-MSE, and GPTQ, with W2 benefits extending across nearly the full model–quantizer matrix.At W4, gains are largest where baseline quantization substantially distorts predictions; already faithful heads change little.Source: §1 · p. 3
• In packed Phi deployment, W4 head compression reduces batch-one latency by 10.8% and batch-16 latency by 9.4%, while reparameterization recovers quality without sacrificing the speedup.The selected representative can have greater logit reconstruction error while placing error where the softmax responds less.Source: §1 · p. 3

2.1 EXACT OUTPUT-HEAD EQUIVALENCE

• Subtracting the same vector from every vocabulary row preserves the linear-softmax distribution because it adds only a common logit shift.The identity is independent of the hidden-state distribution, and an unchanged output bias preserves it.Source: §2.1 · pp. 3–4
• The induced logit error is E_a h, and exact equivalence alone does not guarantee that quantization improves fidelity.Quantization changes the residual even when the source distribution is preserved.Source: §2.1 · p. 4

2.2 FIDELITY-SELECTED REPARAMETERIZATION

• The method searches shifts along the vocabulary-row mean direction and selects the coefficient by validation KL using a disjoint selection set.The finite search includes zero; the selected coefficient is guaranteed no worse in validation KL only for the quantizer used in selection, not on unseen data or another quantizer.Source: §2.2 · p. 4
• The mean direction gives a low-cost one-dimensional path, but the authors do not claim it contains the optimal representative.Full-precision predictions remain unchanged across shifts while quantized KL can vary substantially.Source: §2.2 · p. 4

2.3 WHY FIXED MEAN-CENTERING IS INSUFFICIENT

• Mean-centering (t = 1) minimizes the full-precision weight norm, but it need not minimize post-quantization KL.On Phi, validation KL selects t = 4; the scan also reports that raw and projected weight MSE select t = 1 while diagnostic KL favors other coefficients.Source: §2.3 · p. 4
• For shift-compatible paths, the selected shift is folded into the packed head without an added inference operation; nonlinear soft-capping requires a rank-one correction.Tied models retain the source input embedding and quantize a separate output copy.Source: §2.3 · p. 5

3 EXPERIMENTS

• The decoder stays fixed while seven heads are tested with W4 G128 RTN, AW-MSE, and full-Hessian GPTQ, selecting coefficients by validation KL separately for each quantizer.The evaluation compares each quantized head with its own source model; the search uses a 14-point grid over [−2, 8].Source: §3 · p. 5, §3.2 · p. 6

3.2 ABLATION STUDIES AND VALIDATION

• Table 2 compares W4 test KL for unshifted, mean-centered, and validation-selected heads; selected shifts beat centering on Phi, BLOOM, and BLOOMZ, while XGLM selects t⋆=1.Lower KL is better.Source: Table 2 · p. 6, §3.2 · p. 6
• On Phi, increasing GPTQ calibration to 65,536 states lowers W4 test KL from 0.160 to 0.089, and reparameterization further reduces it to 0.035 (61%).The range initializer is held fixed.Source: §3.2 · p. 6
• On 26 previously unused WikiText articles, reparameterization reduces Phi AW-MSE KL by 73.6% (0.943 to 0.249) and BLOOM AW-MSE KL by 74.5% (0.558 to 0.142).The grid, selection rule, quantizers, and code were frozen before this evaluation.Source: §3.2 · p. 6
• Frozen WikiText coefficients outperform fixed mean-centering in all 18 C4 and OpenWebMath comparisons where the coefficient differs from 1, and tie the remaining six.Larger validation sets from those domains can select coefficients different from WikiText.Source: §3.2 · p. 6

3.3 PACKED INFERENCE EVALUATION

• On matched WikiText tokens, packed AW-MSE shifting lowers perplexity from 31.65 to 15.40 and KL to BF16 from 0.97 to 0.28; GPTQ lowers perplexity from 13.67 to 12.23 and KL from 0.15 to 0.06.Table 3 reports prediction fidelity and generation latency for packed W4 Phi-4-mini heads.Source: Table 3 · p. 7, §3.3 · p. 7

3.2 ABLATION STUDIES AND VALIDATION

• A higher-capacity group-wise extension adds gains under W4 RTN and AW-MSE in a nine-model study and a separate SmolLM3 evaluation.Source: §3.2 · p. 7

3.3 PACKED INFERENCE EVALUATION

• For Phi’s linear-softmax head, folding the shift into weights adds no inference operation, and shifted AW-MSE has 10.8% lower batch-one latency than BF16.GPTQ latency was not measured; the tied deployment retains a BF16 input embedding and uses a separate packed output head.Source: §3.3 · p. 7

4.1 SOFTMAX SENSITIVITY TO LOGIT ERROR

• For small logit perturbations, the Fisher quadratic approximates KL by weighting relative-logit errors according to source probabilities; it is a local approximation with higher-order terms.For nonlinear pre-softmax transformations such as soft-capping, the quadratic uses the post-transformation residual after the rank-one correction.Source: §4.1 · pp. 7–8
• The Fisher matrix’s common-shift invariance motivates projecting out the vocabulary-wide mean; the analysis diagnoses quantizer residuals, while representative selection uses actual validation KL.For RTN and AW-MSE, the selected shift also coincides with a larger softmax-invariant component.Source: §4.1 · p. 8, §4.2 · p. 9

4.2 EMPIRICAL ANALYSIS OF QUANTIZATION RESIDUALS

• From t=0 to t=4, raw logit-error energy rises 3.0× under RTN and 2.7× under AW-MSE, while Fisher-weighted error falls by 72–73%.Across all three quantizers, the selected representative lowers the residual’s distributional cost.Source: §4.2 · pp. 8–9
• Under full-Hessian GPTQ, t=4 raises source-normalized logit-error energy Dz from 0.025% to 0.067% (2.66×), while Fisher/2 falls from 0.0888 to 0.0347 and test KL from 0.0894 to 0.0349.The Fisher quadratic matches actual KL within 0.7% at both coefficients, even after GPTQ’s second-order error compensation.Source: §4.2 · p. 8
• Under AW-MSE, error energy on entries with p<10^-6 rises from 92.4% to 98.2%, while absolute error on entries with p≥0.01 falls by more than half.The low-probability entries carry 0.20% of probability mass; the p≥0.01 entries carry 84.8%.Source: §4.2 · p. 8

Quantization and equivalent parameterizations

• Unlike channel scaling, rotations, and training-time symmetry methods, this approach searches a post-training scalar family specific to the additive symmetry of a frozen linear-softmax output head.The input representation remains unchanged.Source: §5 · p. 9

Softmax invariance and output centering

• Unlike fixed centering of attention keys or pretraining-time output-row centering, the method selects a static output-weight shift by post-quantization predictive fidelity.Neither softmax symmetry nor mean-centering is claimed as a new contribution.Source: §5 · p. 9

Reconstruction and distributional objectives

• Unlike reconstruction-objective methods and curvature-based rate allocation, the paper uses curvature diagnostically and measured KL to select a representative before applying a base quantizer.Source: §5 · p. 9

Large-vocabulary prediction and execution

• The intervention keeps a dense head and reparameterizes it before an existing quantizer, rather than reorganizing vocabulary prediction or compressing the head with low-rank structure.The related work also identifies tied-weight separation and packed execution as established deployment tools.Source: §5 · p. 9

• W4 gains concentrate on heads with substantial baseline error, while W2 gains broaden across nearly the full model–quantizer matrix; packed inference preserves the latency benefit.Matched residual analysis finds fidelity can improve despite greater reconstruction error when the residual is less costly to the predictive distribution.Source: §6 · p. 10

REPRODUCIBILITY STATEMENT

• The protocol records article partitions, token limits, scoring precision, candidate-selection rules, quantizer conventions, and model dimensions for reproduction.The reported evaluation uses validation to select the coefficient and separate evaluation states for KL, top-1 agreement, and perplexity.Source: REPRODUCIBILITY STATEMENT · p. 10, §A.1 · pp. 13–14

APPENDIX

• The appendices specify separate protocols for holdout, transfer, stronger quantization controls, mechanism analysis, packed inference, and grouped search.The paper states that code, machine-readable experiment records, and scripts for regenerating principal result tables will be released upon publication.Source: REPRODUCIBILITY STATEMENT · p. 10, APPENDIX · p. 13

A.1 DATA, SELECTION AND EVALUATION

• Validation selects each quantizer’s coefficient from a fixed grid, while evaluation uses disjoint articles and states; earlier pilot use of the diagnostic pool is disclosed.The diagnostic pool had been used in earlier pilot exploration; the separate untouched holdout is reported in Appendix C.1.Source: §A.1 · p. 13
• The search spans 14 ordered coefficients, including the original head and mean-centering, and reuses fitting moments or the same Hessian across candidates.The four-head RTN/AW-MSE search takes about five minutes on one A100; Phi’s 14-point GPTQ sweep takes about twelve minutes.Source: §A.1 · p. 13

A.2 BASE QUANTIZERS

• RTN fixes c = 1; GPTQ propagates rounding error using the full fitting second moment, with 1% damping, block size 128, and no activation ordering.AW-MSE uses fitting activation moments as weights; unweighted MSE sets those weights to 1.Source: §A.2 · p. 14

A.3 NUMERICAL CONVENTIONS BY EXPERIMENT

• Numerical conventions differ by experiment: the packed Phi AW-MSE path uses the first ten clipping factors through 0.7, while RTN uses c = 1.The reconstruction scan is a separate symmetric-RTN diagnostic; raw/shifted pairs share conventions, while cross-quantizer comparisons can differ in scale precision and integer range.Source: §A.3 · p. 14, Table 5 · p. 15

B.1 SHIFT-COMPATIBLE LOGIT PATHS

• A shared shift preserves softmax equivalence when the complete pre-softmax logit path preserves vocabulary-wide additive shifts; tanh soft-capping does not.Identity logits, temperature scaling, and an unchanged additive bias satisfy the shift condition.Source: §B.1 · p. 15, B NONLINEAR LOGIT PATHS AND INFERENCE… · p. 15

B.2 RANK-ONE CORRECTION

• For nonlinear logit paths, the method quantizes the centered head and restores the removed shared component before the nonlinearity as a rank-one correction.The correction costs one d-dimensional dot product and a scalar broadcast; for a pure softmax path it may be omitted.Source: §B.2 · p. 15
• The rank-one correction reproduces the unquantized source distribution on Gemma-4 to KL ≈7 × 10^-10, but its packed runtime cost was not measured.The packed timing results concern Phi and do not measure this correction.Source: §B.2 · p. 15

C.1 UNTOUCHED HOLDOUT REPLICATION

• On the untouched 26-article holdout, Phi KL reductions are 72.9–74.4% under AW-MSE and 70.4–71.7% under RTN; BLOOM reductions are 72.4–76.2% and 48.7–54.3%, respectively.BLOOM GPTQ improves by 23.4%, without a reported bootstrap interval; the table reports original-to-shifted pairs.Source: Table 6 · p. 16, §C.1 · p. 16

C.2 FROZEN ENGLISH-SELECTED SHIFTS ON FLORES

• The English-selected shifts lower all 18 XGLM FLORES KL values but can increase BLOOM KL, including significant regressions for Arabic/RTN, Hindi/RTN, and Hindi/AW-MSE.This short-context fidelity test does not evaluate translation quality or guarantee multilingual transfer.Source: §C.2 · p. 16
• Frozen English-selected shifts lower XGLM KL across six FLORES languages but produce several BLOOM regressions, so transfer outcomes vary by model and language.The table reports raw-to-shifted KL in nats and marks paired intervals containing zero.Source: §C.2 · p. 16

C.3 CROSS-DOMAIN TRANSFER AND COEFFICIENT STABILITY

• Across C4 and OpenWebMath, frozen WikiText coefficients beat mean-centering in all 18 comparisons where tWT ≠ 1 and tie in the six XGLM cases where tWT = 1.All frozen-coefficient KL point estimates are below the unshifted baseline; the table reports unshifted, centered, and frozen-coefficient KL.Source: §C.3 · pp. 16–17, Table 8 · p. 17
• The transfer studies do not isolate distribution shift from differences in validation-set size and quantized-tensor reconstruction.AW-MSE and GPTQ tensors were rebuilt for transfer; comparisons within each study use matched reconstructed candidates.Source: §C.3 · p. 18

C.4 LOWER-PRECISION EXTENSION (W3 AND W2)

• At W3, Phi test KL falls from 12.97 to 2.46 under RTN, 4.62 to 1.12 under AW-MSE, and 0.543 to 0.211 under GPTQ; no head’s test KL increases.Table 11 reports raw-to-reparameterized test KL with per-quantizer validation selection.Source: §C.4 · p. 18, Table 11 · p. 19
• At W2, Phi GPTQ remains degraded at PPL 30.0 versus the matched BF16 reference of 9.73, while Gemma 3 RTN KL falls from 248.3 to 146.3.The table covers seven heads at W3 and W2; W2 distortion remains large despite shifting.Source: §C.4 · p. 18, Table 11 · p. 19
• W3 perplexity also drops sharply for near-degenerate RTN heads: Phi falls from 3.98 × 10^6 to 111.5, and XGLM from 1.5 × 10^14 to 34.5.The same-split BF16 reference for XGLM is 13.8; Phi GPTQ reaches 11.97 versus BF16 9.73.Source: §C.4 · p. 18, Table 12 · p. 19

D COMPLEMENTARITY AND STRONGER CONTROLS ON PHI

• Phi shift gains persist under stronger GPTQ calibration, exact channel scaling, and affine quantization in the reported W4/G128 controls.All three GPTQ conditions select t = 4; the affine condition selects t = 5.Source: D COMPLEMENTARITY AND STRONGER… · pp. 19–20, Table 13 · p. 20
• The mean direction gives the lowest KL among the tested alternatives; norm-matched random shifts also improve on the raw head.The controls support the mean direction’s empirical utility, not its optimality.Source: Table 14 · p. 20, D COMPLEMENTARITY AND STRONGER… · p. 20
• All eight Phi KL-reduction intervals exclude zero, conditional on validation selection and based on only 16 independent test articles.Intervals are paired article-bootstrap intervals from 10,000 resamples and do not establish cross-family generalization.Source: D COMPLEMENTARITY AND STRONGER… · p. 20, Table 15 · p. 21

E.1 FIXED CENTERING VERSUS THE FIDELITY-SELECTED SHIFT

• Across the four original heads, Phi benefits most from W4 shifts while the other three change little under RTN and AW-MSE.Table 16 reports raw and reparameterized test KL, with coefficients selected independently by validation KL.Source: Table 16 · p. 22, §E.1 · p. 22
• Transferring Gemma 3’s RTN-selected coefficient to AW-MSE yields test KL 0.042, versus 0.038 with AW-MSE’s own selection and 0.041 without a shift.This comparison distinguishes quantizer-specific selection from reusing a coefficient selected for another quantizer.Source: §E.1 · p. 22

E.2 RECONSTRUCTION-SELECTED VERSUS FIDELITY-SELECTED REPRESENTATIVE

• On the tested candidates, raw and projected weight MSE both select t = 1, while diagnostic KL favors different shifts for Gemma, Qwen, and Phi.The scan is limited to t ≤ 2 and 32 test article prefixes; it does not establish a general failure of reconstruction objectives.Source: §E.2 · pp. 22–23, Table 17 · p. 22

E.3 RANGE AND SHARED-ROW CONTROLS

• Smaller group ranges do not guarantee lower KL: Gemma and Qwen have smaller ranges but worse KL after centering, whereas Phi has a slightly larger range and better KL.Qwen and Phi have comparable shared weight-energy fractions but opposite fidelity changes.Source: §E.3 · p. 23, Table 18 · p. 23

F MECHANISM DETAILS

• The mechanism analysis compares raw reconstruction error with Fisher-weighted error across source-probability bins to locate where shifted-head errors fall.The binning uses the same probability groups across shifts.Source: F MECHANISM DETAILS · p. 23, §F.1 · p. 23, Table 19 · p. 24

F.1 FISHER AND PROBABILITY-BIN DIAGNOSTICS

• Vocabulary centering removes 4.1% of AW-MSE residual energy at t = 0 and 27.5% at t = 4; for RTN, the corresponding shares are 0.03% and 5.7%.The lowest-probability bin contains 98.75% of vocabulary entries on average, so its large error share alone does not establish preferential concentration.Source: §F.1 · p. 23, Table 19 · p. 24

F.2 RESIDUAL DIAGNOSTICS UNDER STRONGER GPTQ

• With stronger full-Hessian GPTQ calibration, moving from t = 0 to t = 4 raises relative weight error by 9.2% and raw logit-error energy by 2.66×, yet lowers Fisher-weighted error by 60.90% and KL by 60.94%.This matched Phi W4 comparison preserves the qualitative residual pattern, but does not establish a universal mechanism or extend the local quadratic claim to W2.Source: §F.2 · p. 24, Table 20 · p. 24

F.3 CROSS-MODEL OUTPUT SENSITIVITY

• For unshifted W4 AW-MSE, Phi has roughly 12× Qwen’s absolute logit-error energy and about 6× its per-unit Fisher sensitivity; both factors contribute to Phi’s larger KL.Phi also has a smaller mean top-two logit margin than Gemma, 1.96 versus 2.59, and a larger mean Fisher trace, 0.592 versus 0.452.Source: §F.3 · p. 25, Table 21 · p. 25

G DEPLOYMENT AND SYSTEMS VALIDATION

• The deployment section documents matched packed Phi exports and timing, tied-versus-untied storage accounting, and an auxiliary Qwen serving measurement.The deployment results use their own BF16 references and numerical paths.Source: G DEPLOYMENT AND SYSTEMS VALIDATION · p. 25

G.1 MATCHED PHI PACKED EXPORT AND TIMING

• The deployment retest builds matched Phi BF16 and W4 heads and verifies serialized tensors; shifted AW-MSE shows essentially unchanged W4 latency, not a guarantee of zero overhead elsewhere.GPTQ latency was not measured, and its timing entries are unreported.Source: §G.1 · pp. 25–26

G.2 TIED AND UNTIED STORAGE

• For tied Phi, retaining the BF16 embedding while adding a packed output projection lowers loaded model weights to about 7.17 GiB from 7.47 GiB with the independent packed head, while increasing resident weight memory.Actual memory traffic depends on caching, batch reuse, packing, and kernel execution.Source: §G.2 · p. 26, Table 23 · p. 27

G.3 AUXILIARY QWEN SERVING

• On Qwen3.5 4B, packed W4 reduces batch-one generation latency by 9.6% versus BF16 on A10G.The timing includes the output projection and surrounding model execution, rather than an isolated head microbenchmark.Source: Table 24 · p. 27, §G.3 · p. 27

H MODEL, HEAD AND VOCABULARY DETAILS

• The appendix provides model and vocabulary reference material, including head dimensions and checkpoint details used for the cross-model comparisons.The head-parameter count is a scale indicator, not a fraction of resident memory.Source: H MODEL, HEAD AND VOCABULARY DETAILS · p. 27, §H.1 · pp. 27–28, Table 26 · p. 28

H.1 OUTPUT-HEAD SIZES

• The evaluated 1.7B BLOOM, BLOOMZ, and XGLM heads place roughly 30% of nominal parameters in their 250–256K-vocabulary output heads.The reported nominal share is a scale indicator, not a resident-memory fraction.Source: §H.1 · p. 27

H.2 PRECISION AND PERPLEXITY ACROSS MODELS

• W4 AW-MSE stays within 1.1% of BF16 perplexity for Gemma 3/4 and Qwen3.5, while Phi’s W4 perplexity remains far above its BF16 reference.Among the evaluated settings, W8 AW-MSE preserves near-reference likelihood for Phi; these are practical operating points, not minimum viable precisions.Source: §H.2 · p. 29, Table 28 · p. 30

I CONFIDENCE AND LIKELIHOOD CONTROLS

• Raw perplexity can reward confidence changes rather than source fidelity, so the paper selects representatives by source-to-candidate KL and reports perplexity separately.Matched temperature controls for the shifted heads remain unmeasured.Source: I CONFIDENCE AND LIKELIHOOD CONTROLS · p. 29

J.1 PARAMETERIZATION AND SEARCH

• Grouped coefficients are initialized from the validation-selected shared shift and updated coordinate-wise by validation KL, separately for each base quantizer.The finite searches do not establish an optimal shift.Source: §J.1 · p. 31, Table 29 · p. 31

J.2 EVALUATION ACROSS NINE MODELS

• Grouped search lowers test KL on all nine models under RTN and AW-MSE, with reductions of 3.0–20.3% and 2.3–22.2%, respectively; GPTQ gains are smaller and mixed.For each quantizer, eight of nine paired intervals exclude zero for RTN and AW-MSE, while only Phi’s GPTQ interval excludes zero.Source: §J.2 · p. 32, Table 30 · p. 32

J.3 SMOLLM3 COMPARISON WITH MATCHED CANDIDATE COUNTS

• On previously unused WikiText articles, frozen grouped shifts reduce SmolLM3 W4 G128 KL from 0.8923 to 0.8008 under RTN and from 0.3404 to 0.2848 under AW-MSE, versus refined scalar search.Both search procedures use 91 validation evaluations; this is an exploratory, pointwise comparison for one model.Source: §J.3 · p. 33