Softmax Reparameterization for Output-Head Quantization

Asim Kadav; Christian Flores; Chirag Arora; Varun Kotte; Hongbo Zheng; Lan Yan; Priya Shanmugasundaram; Tracy Holloway King

cs.LGSep 25, 2026

Softmax Reparameterization for Output-Head Quantization

Authors: Asim Kadav, Christian Flores, Chirag Arora, Varun Kotte, Hongbo Zheng, Lan Yan, Priya Shanmugasundaram, Tracy Holloway King

Organizations: Adobe Search, Discovery & ContentAI

Abstract

Large vocabularies make output heads a substantial inference cost in small language models. We introduce softmax reparameterization, a post-training method that searches over functionally equivalent output heads before quantization. The method subtracts a scalar multiple of the vocabulary-row mean from every output row and selects the coefficient by validation KL. For linear-softmax heads, these shifts preserve full-precision predictions exactly and require no decoder retraining; a rank-one correction extends the construction to nonlinear logit paths. Across seven output heads and three quantizers, W4 gains are largest where baseline quantization substantially distorts predictions: test KL falls by 93% on XGLM under RTN and by 73--77% on Phi, BLOOM, and BLOOMZ under activation-weighted MSE. Heads with low baseline error change little; at W2, used as a compression stress test, benefits extend more broadly. On Phi, the gains persist under stronger GPTQ calibration; a separate untouched holdout reproduces the improvements on Phi and BLOOM. Frozen WikiText-selected coefficients also transfer without retuning to C4 and OpenWebMath. Residual analysis on Phi shows how fidelity can improve despite greater total logit error: the selected representative reduces error on likely outputs and lowers its Fisher-weighted cost. For shift-compatible heads, the shift adds no inference operation. With the decoder held in BF16, a packed W4 Phi output head reduces batch-one generation latency by 10.8%, and reparameterization preserves this speedup.

Figures & tables

	W4 KL			W2 KL
Model	RTN	AW-MSE	GPTQ	RTN	AW-MSE	GPTQ
Phi-4-mini	1.23 $\to$ 0.351	0.936 $\to$ 0.256	0.158 $\to$ 0.059	98.2 $\to$ 74.2	41.6 $\to$ 16.2	3.08 $\to$ 1.16
Gemma 3	0.049 $\to$ 0.046	0.041 $\to$ 0.038	0.033 $\to$ 0.033	2.66 $\to$ 2.21	0.822 $\to$ 0.714	0.508 $\to$ 0.491
Gemma 4 †	0.012 $\to$ 0.010	0.008 $\to$ 0.008	0.008 $\to$ 0.008	0.911 $\to$ 0.911	0.212 $\to$ 0.199	0.165 $\to$ 0.162
Qwen3.5	0.021 $\to$ 0.019	0.013 $\to$ 0.013	0.011 $\to$ 0.010	1.84 $\to$ 1.74	0.312 $\to$ 0.285	0.200 $\to$ 0.197
BLOOM-1.7B	1.12 $\to$ 0.531	0.594 $\to$ 0.136	0.037 $\to$ 0.028	16.7 $\to$ 6.15	4.46 $\to$ 1.46	0.427 $\to$ 0.335

Table 1: Output-head quantization fidelity relative to the BF16 source model. Seven heads at W4 and W2 under RTN, AW-MSE and full-Hessian GPTQ. Each cell reports test KL before (

t=0

)

\to

after reparameterization (

t^{\star}

); lower is better. The shift is selected separately for each quantizer by validation KL. Gemma 4 (

\dagger

) uses the rank-one soft-cap correction.

	RTN			AW-MSE			GPTQ
Head	$t{=}0$	$t{=}1$	$t^{\star}$	$t{=}0$	$t{=}1$	$t^{\star}$	$t{=}0$	$t{=}1$	$t^{\star}$
Phi-4-mini	1.229	0.811	0.351	0.936	0.577	0.256	0.158	0.117	0.059
BLOOM-1.7B	1.117	0.978	0.531	0.594	0.451	0.136	0.037	0.036	0.028
BLOOMZ-1.7B	1.119	0.859	0.483	0.662	0.399	0.158	0.042	0.039	0.032
XGLM-1.7B	2.129	0.143	0.143	0.584	0.095	0.095	0.009	0.007	0.007

Table 2: Fixed mean-centering versus validation-selected reparameterization at W4. Test KL relative to the BF16 source model for the unshifted head (

t=0

), mean-centered head (

t=1

), and validation-selected head (

t^{\star}

), using the Table 1 evaluation split. Lower is better.

Head	KL to BF16	WikiText PPL	B1 (ms)	B16 (ms)
BF16	0.00	11.65	1136.6	1335.8
W4 min–max	1.23	40.57	1015.3	1211.8
W4 min–max + shift	0.34	16.28	1014.1	1210.1
W4 AW-MSE	0.97	31.65	1014.8	1210.2
W4 AW-MSE + shift	0.28	15.40	1013.4	1210.3
W4 GPTQ	0.15	13.67	—	—

Table 3: Prediction fidelity and generation latency with packed W4 output heads on Phi-4-mini. Quality is evaluated on identical WikiText tokens. KL is measured relative to BF16 using FP32 readout; perplexity uses vLLM-packed execution. A10G latency covers 64-token greedy generation at batch sizes 1 and 16, with prefix caching disabled. Dashes indicate unmeasured GPTQ latency.

Quantizer	$t$	$\\|E_{t}\\|/\\|W_{t}\\|$	$D_{z}$ (%)	Common (%)	Fisher/2	Actual KL
RTN	0	0.141	0.239	0.03	1.271	1.223
RTN	1	0.137	0.189	0.00	0.720	0.821
RTN	4	0.150	0.707	5.67	0.347	0.352
AW-MSE	0	0.125	0.249	4.06	0.861	0.935
AW-MSE	1	0.121	0.191	0.13	0.532	0.583
AW-MSE	4	0.136	0.677	27.51	0.240	0.256

Table 4: Lower KL despite greater logit error on Phi. Moving from

t=0

t=4

increases

D_{z}

while reducing Fisher/2 and KL.

D_{z}

is source-normalized logit-error energy; Common is the fraction removed by vocabulary centering. Fisher/2 denotes

\tfrac{1}{2}\mathbb{E}[e^{\top}F(p)e]

. Metrics use matched heads and states.

Source probability	Mass (%)	Visible logit error, $t=0\to 4$	Fisher/2, $0\to 4$	Drop (%)
$p\geq 10^{-2}$	84.766	$35.00\to 13.48$	$0.60557\to 0.15943$	71.8
$10^{-6}\leq p<10^{-2}$	15.037	$10,812.73\to 5,279.47$	$0.25264\to 0.07930$	27.9
$p<10^{-6}$	0.196	$132,569.58\to 288,852.93$	$0.00300\to 0.00116$	0.3

Table 5: Probability-bin decomposition of the Fisher quadratic for Phi W4 AW-MSE. Means use the Table 4 test states. Visible logit error is squared logit error after removing the vocabulary-wide common offset; Fisher/2 uses probability centering. Mass is source probability mass; Drop is each bin’s share of the total Fisher/2 reduction from

t=0

t=4

Appendix figures & tables27 assets

Supplementary material from the paper’s appendix.

Appendix

Comparison	Codes	Scale selection/storage	Reconstruction
Matrix/breadth RTN	Symmetric	FP32	FP32
Matrix/breadth AW-MSE	Signed	Score and store BF16	BF16, then FP32 products
GPTQ probe	Signed	BF16 initial scales	BF16, then FP32 products
Matched Phi residuals	Signed	Score and store BF16	BF16, then FP32 products
Packed Phi RTN/AW-MSE	Signed	Float candidates; store BF16	Packed W4 serving

Appendix

Table 6: Numerical conventions for the reported comparisons. “Signed” is

[-2^{b-1},k_{b}]

; “symmetric” is

[-k_{b},k_{b}]

. Reconstruction precision describes the weights presented to the scoring path.

Head	Quantizer	KL	PPL	BF16 PPL
Phi	AW-MSE	$0.943\to 0.249$	$23.41\to 11.93$	9.33
Phi	RTN	$1.204\to 0.348$	$31.00\to 12.93$	9.33
BLOOM	AW-MSE	$0.558\to 0.142$	—	—
BLOOM	GPTQ	$0.037\to 0.028$	—	—

Appendix

Table 7: KL reductions persist on untouched holdouts. Pairs are original

\to

shifted; BF16 PPL uses the same source and evaluation path. Dashes denote unreported values.

Model	Language	RTN	AW-MSE	GPTQ
BLOOM-1.7B	English	$1.586599\!\to\!0.665851$	$0.920758\!\to\!0.144382$	$0.055268\!\to\!0.038390$
	French	$0.946322\!\to\!0.556638$	$0.845735\!\to\!0.208255$	$0.713031\!\to\!0.250820$
	Spanish	$2.074753\!\to\!0.762232$	$1.565073\!\to\!0.273351$	$0.815738\!\to\!0.297205$
	Arabic	$\mathbf{0.523881}\!\to\!\mathbf{0.876482}$	$0.150436\!\to\!0.148087^{\dagger}$	$0.236476\!\to\!0.209441$
	Hindi	$\mathbf{1.629840}\!\to\!\mathbf{1.847837}$	$\mathbf{0.230221}\!\to\!\mathbf{0.420559}$	$0.951959\!\to\!0.447360$
	Chinese	$0.641800\!\to\!0.625781$	$0.083138\!\to\!0.067635$	$\mathbf{0.118631}\!\to\!\mathbf{0.119211}^{\dagger}$

Appendix

Table 8: English-selected shifts lower all 18 XGLM FLORES KL values but can regress on BLOOM. Entries are raw

\to

shifted KL (nats). BLOOM uses

t=(6,4,4)

for (RTN, AW-MSE, GPTQ); XGLM uses

t=1

. GPTQ retains the original 1,024-state Hessian and ranges. Bold marks increased KL;

\dagger

marks paired 95% intervals containing zero: 5,000 bootstrap resamples of 75 source-article URL clusters, without multiplicity adjustment. Arabic is Modern Standard Arabic; Chinese uses Simplified script.

Head	Quantizer	Domain	$t_{\rm WT}$	Raw KL	Centered KL	Frozen KL	$\Delta$ KL: 95% CI
Phi-4-mini	RTN	C4	$4$	1.372879	0.899519	0.393513	$[-0.517935,-0.493749]$
Phi-4-mini	RTN	OpenWebMath	$4$	1.105418	0.675704	0.342138	$[-0.348787,-0.318388]$
Phi-4-mini	AW-MSE	C4	$4$	1.039478	0.657195	0.304324	$[-0.359930,-0.345035]$
Phi-4-mini	AW-MSE	OpenWebMath	$4$	0.764193	0.498817	0.252586	$[-0.256702,-0.235430]$
Phi-4-mini	GPTQ	C4	$4$	0.111289	0.078885	0.041984	$[-0.037799,-0.035971]$
Phi-4-mini	GPTQ	OpenWebMath	$4$	0.101269	0.072217	0.044054	$[-0.029566,-0.026734]$

Appendix

Table 9: Frozen WikiText coefficients beat centering in 18 transfer comparisons and tie in six. C4/OpenWebMath test KL to BF16 (nats): unshifted (

t=0

), centered (

t=1

), and frozen (

t=t_{\rm WT}

). Frozen-minus-centered KL intervals use 10,000 paired document-bootstrap resamples; they are pointwise, conditional on frozen coefficients, and unadjusted for multiple comparisons.

Quantizer	WikiText (16)	C4 (2,048)	OpenWebMath (2,048)
RTN	$6$	$4$	$4$
AW-MSE	$4$	$4$	$5$
GPTQ	$5$	$4$	$4$

Appendix

Table 10: BLOOM’s selected coefficient varies across validation domains. Column headings give validation sizes in articles/documents. WikiText coefficients are from the main experiment; C4 and OpenWebMath selections use separate 2,048-document pools.

Quantizer	Domain	Selection at $n=16$	Selection at $n=1{,}024$	Test KL: $16\to 1{,}024$
RTN	C4	$t=4$ : 6/10; $t=6$ : 4/10	$t=4$ : 10/10	$0.572704\to 0.563011$
RTN	OpenWebMath	$t=4$ : 9/10; $t=6$ : 1/10	$t=4$ : 10/10	$0.532576\to 0.531784$
AW-MSE	C4	$t=4$ : 10/10	$t=4$ : 10/10	$0.146926\to 0.146926$
AW-MSE	OpenWebMath	$t=4$ : 3/10; $t=5$ : 7/10	$t=5$ : 10/10	$0.126947\to 0.126493$
GPTQ	C4	$t=3$ : 3/10; $t=4$ : 7/10	$t=4$ : 10/10	$0.033628\to 0.033356$
GPTQ	OpenWebMath	$t=3$ : 5/10; $t=4$ : 5/10	$t=4$ : 10/10	$0.054443\to 0.053538$

Appendix

Table 11: At 1,024 validation documents, all ten seeds agree on a non-unit BLOOM coefficient per setting. Counts show selection frequencies. Test KL averages seeds on a fixed, disjoint 1,024-document set unused for selection.

Bits	Model	RTN	AW-MSE	GPTQ
W2	Phi	98.2 $\to$ 74.2	41.6 $\to$ 16.2	3.08 $\to$ 1.16
W2	Gemma 3	2.66 $\to$ 2.21	0.822 $\to$ 0.714	0.508 $\to$ 0.491
W2	Gemma 4 †	0.911 $\to$ 0.911	0.212 $\to$ 0.199	0.165 $\to$ 0.162
W2	Qwen3.5	1.84 $\to$ 1.74	0.312 $\to$ 0.285	0.2 $\to$ 0.197
W2	BLOOM-1.7B	16.7 $\to$ 6.15	4.46 $\to$ 1.46	0.427 $\to$ 0.335
W2	BLOOMZ-1.7B	16.2 $\to$ 6.22	3.99 $\to$ 2.12	0.456 $\to$ 0.367

Appendix

Table 12: Shifting lowers or preserves KL at W3 and W2, but W2 distortion remains large. Each cell is raw (

t=0

)

\to

reparameterized (

t^{\star}

) test KL to the source, with per-quantizer validation selection over the same 14-point grid. The four main heads appear above each rule; three very-large-vocabulary heads appear below. Gemma 4 (

\dagger

) uses the rank-one soft-cap correction.

Head	BF16	RTN	AW-MSE	GPTQ
Phi	9.73	$3.98{\times}10^{6}$ $\to$ 111.5	1002 $\to$ 28.47	16.86 $\to$ 11.97
Gemma 3	38.49	44.34 $\to$ 44.34	38.54 $\to$ 38.89	40.09 $\to$ 40.09
Gemma 4 †	66.41	71.11 $\to$ 70.83	68.71 $\to$ 69.06	69.05 $\to$ 68.85
Qwen3.5	8.98	9.851 $\to$ 9.793	9.503 $\to$ 9.446	9.351 $\to$ 9.291
BLOOM-1.7B	18.47	282.3 $\to$ 100.7	50.78 $\to$ 26.89	20.14 $\to$ 20.04
BLOOMZ-1.7B	22.06	285.3 $\to$ 124.9	56.08 $\to$ 31.42	23.8 $\to$ 23.9

Appendix

Table 13: KL-selected shifts can substantially reduce W3 perplexity, but gains are not uniform. Each quantizer cell is raw (

t=0

)

\to

reparameterized (

t^{\star}

) perplexity. BF16 is the unquantized source head on the same 8,176-position test split (FP32 readout). Gemma 4 (

\dagger

) uses the rank-one soft-cap correction.

Condition	Test KL	Test PPL	KL reduction
GPTQ: 1,024 states	$0.16041\to 0.06051$	$11.2883\to 10.2858$	$62.3\%$
GPTQ: 65,536, original ranges	$0.08936\to 0.03491$	$10.5608\to 10.0542$	$60.9\%$
GPTQ: 65,536, refitted ranges	$0.09247\to 0.03491$	$10.6275\to 10.0182$	$62.2\%$
RTN	$1.23002\to 0.35058$	$34.1129\to 13.4611$	$71.5\%$
AW-MSE	$0.93940\to 0.25551$	$24.8709\to 12.4856$	$72.8\%$
Scaled RTN	$0.33659\to 0.24234$	$13.5410\to 12.2376$	$28.0\%$

Appendix

Table 14: Phi shift gains persist under stronger calibration, scaling and affine quantization. W4/G128 entries: baseline

\to

shifted. Scaled baselines include selected channel scaling. BF16 PPL is 9.7192.

Direction	$t^{\star}$	Test KL	Test PPL	$t^{\star}$	Test KL	Test PPL
	RTN			AW-MSE
Vocabulary mean	$4$	0.35058	13.4611	$4$	0.25551	12.4856
Coordinate-wise median	$4$	0.73863	20.1241	$4$	0.55894	17.0431
Random seed 0	$-0.5$	1.05676	27.7650	$-1$	0.77292	21.8672
Random seed 1	$2$	1.06612	28.1215	$1$	0.86884	22.6155
Random seed 2	$-1$	1.04971	27.0636	$-0.5$	0.84131	22.1993

Appendix

Table 15: The mean direction gives the lowest KL among the tested alternatives on Phi. Median and random directions are norm-matched to the mean and use the same validation scalar search. Random shifts also improve on the raw head.

Condition	$\Delta$ test KL	$95\%$ interval
GPTQ: 1,024 states	$-0.099903$	$[-0.106154,-0.093710]$
GPTQ: 65,536, original ranges	$-0.054452$	$[-0.059731,-0.049574]$
GPTQ: 65,536, refitted ranges	$-0.057559$	$[-0.062531,-0.052901]$
RTN	$-0.879442$	$[-0.933053,-0.824182]$
AW-MSE	$-0.683890$	$[-0.726657,-0.637061]$
Scaled RTN	$-0.094253$	$[-0.104851,-0.084514]$

Appendix

Table 16: All eight Phi KL-reduction intervals exclude zero. Negative differences favor shifting. Paired 95% intervals use 10,000 bootstrap resamples of 16 test articles, conditional on validation selection.

Model	$t^{\star}$	Raw KL	Reparam KL	$t^{\star}$	Raw KL	Reparam KL
	RTN			AW-MSE
Phi	4	1.229	0.351	4	0.936	0.256
Gemma 3	0.5	0.049	0.046	$-1$	0.041	0.038
Gemma 4 †	1	0.012	0.010	0.5	0.008	0.008
Qwen3.5	$-0.5$	0.021	0.019	$-0.5$	0.013	0.013

Appendix

Table 17: Phi benefits most from W4 shifts; the other three heads change little. RTN and AW-MSE use G128 and independently select coefficients by validation KL. Entries report raw and reparameterized test KL against each head’s source distribution. Gemma 4 (

\dagger

) uses the rank-one soft-cap correction.

Model	Bits	MSE	Proj. MSE	KL	At MSE min	At KL min
		Best sampled $t$			Diagnostic KL
Gemma 3	4	1	1	0.5	0.0580	0.0465
Gemma 3	3	1	1	0.25	0.3161	0.2340
Gemma 3	2	1	1	0	2.8305	2.5793
Qwen3.5	4	1	1	0	0.0291	0.0202
Qwen3.5	3	1	1	0	0.1599	0.1017

Appendix

Table 18: Weight MSE and KL favor different shifts. Each row uses the same symmetric RTN quantizer and diagnostic states at every coefficient. Both raw and projected weight MSE select

t=1

; the KL minimum is descriptive on this diagnostic set, not independently selected on validation.

Model	Shared weight energy (%)	Range ratio	Raw KL	Centered KL
Gemma 3	3.41	0.9612	0.0499	0.0585
Qwen3.5	14.23	0.9029	0.0196	0.0294
Phi	16.40	1.0128	1.2109	0.8087

Appendix

Table 19: Smaller group ranges do not guarantee lower KL. Shared energy and range ratios describe unquantized weights. KL compares raw and mean-centered W4 G128 signed min–max heads on matched diagnostic articles.

Quantizer	$t$	Source probability	Mass (%)	Error share (%)	Error energy	Fisher share (%)
RTN	0	$p<10^{-6}$	0.196	93.474	134039.57	0.317
RTN	0	$10^{-6}\leq p<10^{-4}$	2.196	6.024	8637.93	4.015
RTN	0	$10^{-4}\leq p<10^{-2}$	12.842	0.476	682.63	24.223
RTN	0	$p\geq 10^{-2}$	84.766	0.026	36.88	71.444
RTN	1	$p<10^{-6}$	0.196	94.179	106621.80	0.392
RTN	1	$10^{-6}\leq p<10^{-4}$	2.196	5.381	6092.47	4.994

Appendix

Table 20: Phi’s selected shift reduces error on high-probability outputs. Bins use unchanged source probabilities on matched states. Error is vocabulary-centered logit-error energy; Fisher shares use probability-centered residuals.

$t$	$\\|E_{t}\\|/\\|W_{t}\\|$	$D_{z}$ (%)	Common (%)	Fisher/2	Actual KL
0	0.148515	0.025340	0.4644	0.088771	0.089362
4	0.162178	0.067442	4.7268	0.034710	0.034909

Appendix

Table 21: Shifting Phi W4 GPTQ lowers KL despite greater logit-error energy. GPTQ uses a 65,536-state full Hessian and the original range initializer. Definitions follow Table 4 ; test states are matched.

Model	$\mathbb{E}\\|\delta z\\|^{2}$	$D_{z}$	$S_{p}$	$C_{\mathrm{out}}$	KL
Gemma 3 4B	$2.08\times 10^{4}$	$2.93\times 10^{-3}$	$3.89\times 10^{-6}$	$3.94\times 10^{-6}$	0.041
Qwen3.5 4B	$1.23\times 10^{4}$	$7.44\times 10^{-3}$	$2.09\times 10^{-6}$	$2.09\times 10^{-6}$	0.013
Phi-4-mini	$1.47\times 10^{5}$	$2.41\times 10^{-3}$	$1.26\times 10^{-5}$	$1.19\times 10^{-5}$	0.925

Appendix

Table 22: Phi exceeds Qwen in both logit-error energy and sensitivity. Unshifted W4 AW-MSE uses disjoint WikiText articles (16,352 positions/model).

D_{z}

normalizes error by source energy;

S_{p}

measures sensitivity per unit error and

C_{\rm out}

its Fisher approximation on the same residuals.

BF16 head	WikiText PPL	$\Delta$ NLL vs. source
Tied source	11.6477	+0.000000
Untied, $t=0$	11.6455	-0.000194
Untied, $t=4$	11.6363	-0.000983

Appendix

Table 23: Untying and BF16 serialization leave small likelihood differences on Phi. Matched exports test the unshifted and shifted heads; algebraic softmax invariance does not imply bitwise inference equivalence.

Representation	Head payload (GiB)	Relative to BF16
BF16	1.2503	1.0000
W8, G128	0.6349	0.5078
W4, G128	0.3223	0.2578
W4, G32	0.3516	0.2812

Appendix

Table 24: W4 G128 uses 26% of the BF16 head payload. Analytical Gemma 3 payload:

V=262{,}208

d=2{,}560

. An untied packed head adds resident storage when the BF16 input embedding is retained.

Head	B1 latency (ms)	B16 throughput (tokens/s)
BF16	516.6	570.9
W8 AW-MSE	480.4	585.9
W4 AW-MSE	467.2	614.8

Appendix

Table 25: W4 reduces Qwen3.5 4B batch-one latency by 9.6% versus BF16 on A10G. Head-only W4/W8 use G128; timings cover 25-token greedy generation over seven warmed compiled-vLLM repetitions, excluding loading and compilation.

Model	Vocabulary	Width	Head (M)	Nominal share (%)
Gemma 3 4B	262,208	2,560	671	16.8
Gemma 4 E4B	262,144	2,560	671	—
Qwen3.5 4B	248,320	2,560	636	15.9
Phi-4-mini	200,064	3,072	615	16.2
BLOOM-1.7B	250,880	2,048	514	30.2
BLOOMZ-1.7B	250,880	2,048	514	30.2

Appendix

Table 26: Output-projection dimensions and nominal parameter shares. Gemma 4’s 262,144-row configuration and the evaluated Gemma 3 checkpoint’s 262,208-row matrix both round to 671M head weights.

Family	Earlier checkpoint and rows	Selected checkpoint and rows	$d$
Gemma	google/gemma-2b [ 2024 ]; $V=256{,}000$	google/gemma-3-4b-it [ 2025 ]; $V=262{,}208$	2560
XGLM	—	facebook/xglm-1.7B [ 2021 ]; $V=256{,}008$	2048
BLOOM	—	bigscience/bloom-1b7 [ 2022 ]; $V=250{,}880$	2048
Qwen	Qwen/Qwen-7B [ 2023 ]; $V=151{,}936$	Qwen/Qwen3.5-4B [ 2026 ]; $V=248{,}320$	2560
Llama	meta-llama/Llama-2-7b-hf [ 2023 ]; $V=32{,}000$	meta-llama/Llama-4-Scout-17B-16E-Instruct [ 2025 ]; $V=202{,}048$	5120
GPT	openai-community/gpt2 [ 2019 ]; $V=50{,}257$	openai/gpt-oss-20b [ 2025 ]; $V=201{,}088$	2880

Appendix

Table 27: Checkpoints behind Figure 1 , in plotted order. Checkpoint names link to configuration sources and years to release sources.

V

counts output rows;

d

is the selected checkpoint width. A dash indicates no earlier comparison.

Family	Model	Vocabulary
Phi	Phi-2	51,200
	Phi-3-mini	32,064
	Phi-4	100,352
	Phi-4-mini	200,064
Llama	Llama 2	32,000
	Llama 3	128,256

Appendix

Table 28: Vocabulary sizes across selected generations of four families, from each model’s official configuration.

Model	Head (M)	BF16	W4 min–max	W4 AW-MSE	W8 AW-MSE
Gemma 3 4B	671	60.12	60.84	59.46	60.12
Gemma 4 E4B	671	74.97	75.48	75.55	74.95
Qwen3.5 4B	636	10.89	11.19	10.98	10.89
Phi-4-mini	615	11.65	40.58	30.94	11.75

Appendix

Table 29: W4 AW-MSE retains near-BF16 perplexity for three heads but substantially degrades Phi. Head-only WikiText evaluation uses G32 for Gemma 3 W4 and G128 otherwise. Each model uses its own BF16 reference and tokenizer, with matched Hugging Face loading and evaluation paths.

	Nine-model experiment	Matched-budget SmolLM3 study
Initialization	Main 14-point scalar grid	Shared 27-point scalar grid
Scalar budget	14	$27+64=91$ (scalar refinement)
Grouped budget	$14+64=78$	$27+64=91$
Evaluation	Main 16 test articles	16 previously unused articles
Quantizers	RTN, AW-MSE, GPTQ	RTN, AW-MSE

Appendix

Table 30: Grouped-search protocols. Counts are validation evaluations per quantizer. Each grouped search adds 32 updates with two proposals each.

Model	Quantizer	$t^{\star}$	Raw	Scalar	Grouped	Red. (%)	$\Delta$ KL: 95% CI
Phi-4-mini	RTN	4	1.230021	0.350579	0.300249	$+14.36$	$[-0.059382,-0.042322]$
	AW-MSE	4	0.939398	0.255508	0.222474	$+12.93$	$[-0.040437,-0.026940]$
	GPTQ	4	0.160414	0.060511	0.058634	$+3.10$	$[-0.003556,-0.000413]$
Gemma 3	RTN	0.5	0.049208	0.046005	0.038335	$+16.67$	$[-0.009860,-0.005429]$
	AW-MSE	-1	0.040232	0.037523	0.033956	$+9.51$	$[-0.005063,-0.002085]$
	GPTQ	0	0.033609	0.033609	0.032296	$+3.91$	$[-0.003080,+0.000345]$

Appendix

Table 31: Grouped search lowers RTN/AW-MSE KL on all nine models; GPTQ results are mixed. W4 G128 KL uses each model’s source; scalar/grouped budgets are 14/78 validation evaluations. Positive reductions favor grouped search;

\Delta=\mathrm{KL}_{\rm grp}-\mathrm{KL}_{\rm scalar}

. Pointwise paired article-bootstrap 95% intervals use 10,000 resamples, conditional on frozen selections, on previously examined test articles. Gemma 4 (

\dagger

) restores the common offset before soft-capping.

		Scalar $\to$ grouped PPL
Model	Source PPL	RTN	AW-MSE	GPTQ
Phi-4-mini	9.719	13.461 $\to$ 13.042	12.486 $\to$ 12.106	10.286 $\to$ 10.287
Gemma 3	37.977	38.482 $\to$ 38.419	38.086 $\to$ 38.176	38.852 $\to$ 38.613
Gemma 4 †	67.223	68.156 $\to$ 68.011	68.192 $\to$ 68.129	67.646 $\to$ 67.961
Qwen3.5	8.977	9.131 $\to$ 9.143	9.106 $\to$ 9.060	9.065 $\to$ 9.054
Qwen3	12.577	13.020 $\to$ 12.839	12.747 $\to$ 12.617	12.636 $\to$ 12.639

Appendix

Table 32: Lower KL does not always yield lower perplexity. Nine-model scalar

\to

grouped perplexity uses the same states as Table 31 . Source PPL is each model’s BF16 reference.

Explore similar work

SoftWater: Class-Aware Rate Allocation for Softmax Quantization
Aug 12, 2026Joao V. Cavalcanti, Ashia C. WilsonGumbel-Softmax Relaxation Kullback-Leibler Divergence
Department of Electrical Engineering and Computer Science · Massachusetts Institute of Technology · Cambridge, MA 02139, USA
ARCHead: Activation-Metric Residual Correction for Large Language Model Output Heads
Aug 3, 2026Şuayp Talha Kocabay, Talha Rüzgar Akkuş, Kamer Ali YukselLarge Language Model Quantization Matched Fp16 Intermediate
Independent Researcher · aiXplain, Inc.
Variable Bit-width Quantization: Learning Per-Group Precision for "Bigger-but-Smaller" Language Models
Jul 3, 2026Hamish OgilvyLarge Language Model Quantization Autoregressive Decoding
Algolia Research

Aug 12, 2026, cs.LGJ/K move · Enter open · S save

SoftWater: Class-Aware Rate Allocation for Softmax Quantization

Joao V. Cavalcanti, Ashia C. Wilson

Department of Electrical Engineering and Computer Science · Massachusetts Institute of Technology · Cambridge, MA 02139, USA

Gumbel-Softmax Relaxation Kullback-Leibler Divergence

Open paper AskPDF arXiv

Post-training quantization pipelines routinely leave the softmax output layer in high precision. Yet in small LLMs with modern vocabularies, the head holds 15--30% of all parameters, so a nominal ``2-bit'' model with an fp16 head can store several times as many bits per weight. We pose softmax-layer quantization as a rate-distortion problem under the KL divergence between the original and quantized output distributions. A second-order analysis reveals a class-aware geometry: quantization error is weighted jointly by feature covariance and class-specific softmax curvature. A separability approximation replaces the

Kn\times Kn

Cholesky with one

n\times n

factorization rescaled per class, making the lattice encodable by successive interference cancellation, with both statistics from a single forward pass. The resulting method, SoftWater, gives fine grids to frequent, low-variance classes and coarse grids to rare ones, a large gap under Zipfian token distributions. Across five models from 1B to 32B, SoftWater outperforms the released WaterSIC quantizer (near-optimal under linear-layer WMSE but not output KL) at matched head rates on 59 of 60 test points, using none of that pipeline's refinements and cutting head-induced KL by

6.5\times

8.3\times

at 2 bits. On Llama-3.2-1B-Instruct with quantized bodies, a 2-bit head removes 45--60% of stored bytes for a

2.9

3.7\%

perplexity increase. Because the class-side statistic comes from calibration data, matching calibration to the deployment domain gives the lowest KL on that domain throughout. On a tied model, a 4-bit head is near-lossless and a 2-bit head costs under 4% perplexity, making head quantization of such models practical.

Abstract

Explore similar work

SoftWater: Class-Aware Rate Allocation for Softmax Quantization

ARCHead: Activation-Metric Residual Correction for Large Language Model Output Heads

Variable Bit-width Quantization: Learning Per-Group Precision for "Bigger-but-Smaller" Language Models