Pretraining Transformers with Quantized Softmax in Attention

Shangzhen Zhu; Muyan Hu; Tomasz Kozlowski

cs.LGSep 27, 2026

Pretraining Transformers with Quantized Softmax in Attention

Authors: Shangzhen Zhu, Muyan Hu, Tomasz Kozlowski

Organizations: University of Illinois Urbana-Champaign

Abstract

Low-precision Transformer systems increasingly quantize attention matrix multiplications, while softmax often remains at higher precision. During pretraining, an approximate softmax changes the gradients that train the model as well as its forward computation. We study this interaction with K-interval attention, which approximates the exponential using K+1 grid values. We vary per-row grid calibration, interpolation versus hard rounding, and the placement of a straight-through surrogate relative to normalization. We derive the corresponding backward rules, including calibration derivatives, and compare these choices in pretraining experiments matched on model, data, and optimizer. Detaching the row extrema leaves the forward computation unchanged but produces a delayed increase in validation loss. With hard rounding at K=4, min-max calibration and a pre-normalization surrogate incur a large loss gap; changing either choice substantially reduces it. At 124M parameters and 2.5B training tokens, fixed-window calibration with a post-normalization surrogate yields a validation loss gap of +0.019 nats relative to softmax at K=4, and with a pre-normalization surrogate yields +0.004 nats at K=16.

Figures & tables

Figure 1: The operator family (

K{=}4

; three design axes). (a) MinMax calibrates the

K{+}1

exponential grid values (dots) over the observed row range,

h=(M-m)/K

; a lower minimum widens every bin at fixed

M

K

. (b) FWM anchors a fixed-width window at the row maximum,

h=\tau/K

, zero weight below it (shaded tail). LERP (green) interpolates between adjacent grid values; Nearest (step) rounds to the nearest grid center. Strips: Nearest rounding bins (boundaries halfway between grid points; scores in one bin share

w^{N}

). (c, d) Weight-STE applies the straight-through replacement to the unnormalized weights, Prob-STE to the normalized probabilities: the node sg (stop-gradient: identity forward, zero derivative backward) receives both reconstructions (c) or both normalized probabilities (d), and its output is added to the LERP branch (dashed arrow); the hard forward

P=P^{N}

is the same, the normalization Jacobian is evaluated at the hard (c) or relaxed (d) weights.

Name	Calibration	Reconstr.	Backward
softmax	—	exp	autograd
LERP	MinMax	LERP	exact Jacobian ( 2 )
MinMax–Weight	MinMax	Nearest	Weight-STE ( 3 )
MinMax–Prob	MinMax	Nearest	Prob-STE ( 4 )
FWM–Weight	FWM, $\tau{=}6$	Nearest	Weight-STE
FWM–Prob	FWM, $\tau{=}6$	Nearest	Prob-STE

Table 1: Method identity (three design axes) and notation; edge cases in Appendix A .

Suite	Purpose	Model	Tokens	Corpus	Seeds	Runs
A long-horizon	$K$ , horizon, downstream	124M	2.5B	FineWebEdu-3B	1	15
B factorial	calibration $\times$ surrogate $\times K$	124M	250M	FineWebEdu-3B	1	10
C scale probe	ordering at $8\times$ parameters	1B	100M	WikiText-103	1	12
D seeds	$5{\times}5$ grid, seed spread	124M	100M	WikiText-103	5	130

Table 2: The four formal suites and the question each answers. Each suite matches model, data and optimizer protocol (AdamW, cosine decay, 1024-token sequences, 131,072 tokens/step, no learning-rate or method tuning) and varies the operator (suite D also the seed); historical training-stack differences: Tables S3 and S4 .

Figure 2: Same forward, different calibration backward (124M, 100M tokens, WikiText-103). (a) LERP

K{=}32

, seed 1337: logged validation loss minus softmax under full (black), detach (purple dashed) and detach with the zero-sum projection (orange dash-dot); grey band: observed separation. (b) Logged mean row span (log). (c) 100M endpoints, detach minus full, every matched pair (CIs: Tables S19 , S23 ). (d) Single-channel ablation, LERP

K{=}32

: each run minus its own same-seed full run, seeds 7 (

\triangle

) and 42 (

\triangledown

), 95% block-bootstrap CIs over validation blocks (evaluation sampling only; Table S25 ).

Figure 3: Calibration

\times

surrogate interaction for Nearest. (a, b) 124M @ 250M at

K{=}4

and

K{=}16

: final

\Delta\mathrm{NLL}

vs softmax by surrogate placement, one line per calibration, 95% block-bootstrap CIs (evaluation sampling only); dotted: LERP reference, in (a) only. (c) Five-seed mean

\Delta\mathrm{NLL}

K

, all five families (suite D), 95%

t

intervals over seeds. (d) The surrogate effect Prob

-

Weight against

K

under each calibration, paired by seed, 95%

t

intervals (4 df). Filled markers/solid lines: Weight-STE; hollow/dashed: Prob-STE.

Figure 4: 124M @ 2.5B. (a) Final

\Delta\mathrm{NLL}

vs softmax by

K

and family at the measured

K

only (whiskers: 95% block-bootstrap CI, evaluation sampling only). (b) Gap along training for the five

K{=}4

conditions. (c) Nearest minus LERP at matched

K{=}4

; the 95% CI (evaluation sampling only) is narrower than the line width. (a, b): symlog vertical scale, linear within

\pm 0.02

nats; (c): linear.

Figure 5: Geometry and resolution at 2.5B. Geometry statistics are layer means over eight validation blocks on each model’s own scores. (a) Native interval width

h

. (b) Total variation between native and softmax probabilities. (c) Endpoint NLL gap versus native width for all 14 approximate conditions. The cross marks MinMax–Weight

K{=}64

(

h=1.54

\Delta\mathrm{NLL}=0.038

): large score span, moderate interval width. Vertical axis symlog, linear for

|\Delta\mathrm{NLL}|\leq 0.02

Figure 6: Downstream accuracy differences from softmax (pp, native forward), 2.5B endpoints; row labels: training-domain $\Delta\mathrm{NLL}$ ; column labels: softmax’s accuracy. (a) The matched $K{=}4$ conditions ( $\pm 20$ pp scale); (b) the near-baseline group ( $\pm 5$ pp). BLiMP macro-accuracy; HellaSwag/PIQA/ARC-Easy acc_norm ; LAMBADA and SciQ acc . Star: BH $q<0.05$ within the column’s 14-comparison family; all 14 conditions in Appendix E .

Appendix figures & tables17 assets

Supplementary material from the paper’s appendix.

Appendix

Condition	Forward	Backward
Degenerate MinMax row ( $M-m\leq\varepsilon$ ): single valid key or equal valid scores	weights set to 1 ( torch.where ), invalid keys 0, $P$ uniform (softmax’s forward)	zero gradient through where , extremum gradients cancel; finite
Very small positive span ( $10^{-6}>\varepsilon$ )	regular branch, tiny $h$ , normal reconstruction	all gradients finite
FWM: single key or equal scores	$z_{j}=0$ for every valid key, all at the top grid value, $P$ uniform; no degenerate branch	regular formulas
Tied maxima/minima, positive span	unique-extremum formulas	amax / amin split the extremum gradient evenly among ties
Grid-boundary index switch	$r_{j}$ is a floor, piecewise constant	zero derivative a.e.; LERP weight continuous, secant identities on either side
Nearest midpoint tie $t_{j}=\tfrac{1}{2}$	rounds up; hard weight discontinuous	( 5 ) applies away from these points

Appendix

Table S1: Edge cases: forward value and backward convention as implemented (verified finite in fp32 and fp64 for LERP, MinMax–Weight, MinMax–Prob at

K\in\{4,16\}

and for FWM–Weight/FWM–Prob).

Run (suite: condition)	mean $\Delta$	max $\|\Delta\|$
B: MinMax–Prob K4	$3.28\mathrm{e}{-5}$	$2.14\mathrm{e}{-2}$
B: MinMax–Prob K16	$-3.37\mathrm{e}{-5}$	$8.90\mathrm{e}{-3}$
B: FWM–Prob K4	$1.08\mathrm{e}{-4}$	$9.90\mathrm{e}{-3}$
B: FWM–Prob K16	$0$	$0$
A: MinMax–Prob K4	$-9.91\mathrm{e}{-5}$	$1.38\mathrm{e}{-2}$
A: FWM–Prob K4	$2.33\mathrm{e}{-5}$	$1.41\mathrm{e}{-2}$

Appendix

Table S2: Suite D: operator family by training stack, in complete training runs (130 formal runs; supplementary runs excluded and listed in Table S4 ). “Logged”: GPU model recorded by the pod’s nvidia-smi in the master log; “author-confirmed”: rented RTX 5090 pods whose logs did not record the model. Local runs: RTX 3090, torch 2.5.1+cu121.

Run (suite: condition)	mean $\Delta$	max $\|\Delta\|$
B: MinMax–Prob K4	$3.28\mathrm{e}{-5}$	$2.14\mathrm{e}{-2}$
B: MinMax–Prob K16	$-3.37\mathrm{e}{-5}$	$8.90\mathrm{e}{-3}$
B: FWM–Prob K4	$1.08\mathrm{e}{-4}$	$9.90\mathrm{e}{-3}$
B: FWM–Prob K16	$0$	$0$
A: MinMax–Prob K4	$-9.91\mathrm{e}{-5}$	$1.38\mathrm{e}{-2}$
A: FWM–Prob K4	$2.33\mathrm{e}{-5}$	$1.41\mathrm{e}{-2}$

Appendix

Table S3: Secondary frozen-checkpoint audit (historical expressions only): Prob-STE frozen code minus the shared Weight-STE-expression evaluator,

\Delta

per block on the full validation split (976 blocks for suites A and B, 243 for C); mean

\Delta

over blocks and the largest single-block

|\Delta|

. Neither side is a direct hard forward (Table S5 ); this audit alone covers suite C and

K{=}16

. Absolute NLL, mean

|\Delta|

and the first-8-block maximum of the parity check: file tab_rounding_full_split_audit.tex in compact/migrated/ .

Figure S1: 124M @ 100M on five seeds, WikiText-103. (a) Calibration, surrogate and interaction effects by

K

, paired by seed, 95%

t

intervals. (b) Seed SD of each condition’s

\Delta\mathrm{NLL}

against the paired evaluation SE of the same contrast; every point lies above the 1:1 line.

Family	slope	95% $t$ -CI	seed SD
LERP	-0.0222	[-0.024, -0.020]	0.0014
MinMax–Weight	-0.0472	[-0.051, -0.043]	0.0033
MinMax–Prob	-0.0301	[-0.034, -0.027]	0.0028
FWM–Weight	-0.0078	[-0.010, -0.005]	0.0020
FWM–Prob	-0.0030	[-0.007, +0.001]	0.0029

Appendix

Table S8: Five-seed

2{\times}2{\times}K

effects in nats: mean over five paired seed-level observations (seed SD); the 95%

t

interval (4 df) is mean

\pm 1.24\,\mathrm{SD}

(Figure S1 a), and every

K\leq 16

interval excludes zero.

Family	slope	95% $t$ -CI	seed SD
LERP	-0.0222	[-0.024, -0.020]	0.0014
MinMax–Weight	-0.0472	[-0.051, -0.043]	0.0033
MinMax–Prob	-0.0301	[-0.034, -0.027]	0.0028
FWM–Weight	-0.0078	[-0.010, -0.005]	0.0020
FWM–Prob	-0.0030	[-0.007, +0.001]	0.0029

Appendix

Table S9: Five-seed

K

trend: slope of

\Delta\mathrm{NLL}

vs softmax in nats per doubling of

K

(least squares against

\log_{2}K

over

K\in\{4,\dots,64\}

, fitted per seed); mean over seeds, 95%

t

-CI, seed SD.

Figure S2: MinMax–Weight

K

sweep along training at 2.5B ,

\Delta\mathrm{NLL}

vs softmax (symlog).

Figure S3: Reduced-precision exponential baselines (suite-D protocol, 124M, WikiText-103, seed 7): validation NLL minus the same-seed softmax run at the five evaluated checkpoints, all evaluated on one RTX 3090. (a) Full range including LERP

K{=}4

; (b) near-reference trajectories, LERP

K{=}4

excluded, with the historical FWM–Prob

K{=}4

run (different training stack). One run per curve; endpoint evaluation intervals in Table S11 .

Figure S4: Layer-mean score geometry along training at 2.5B ,

K{=}4

conditions, softmax and MinMax–Weight

K{=}64

(dotted), on each model’s own scores over 8 validation blocks. (a) Row span; (b) softmax mass on the keys in the operator’s top set, which is the top rounding bucket for Nearest and the top interval for LERP (these sets differ, so the panel compares each operator with its own top set). The native interval width

h

and the per-row TV of the same six approximate conditions and checkpoints are Figure 5 a,b.

Figure S5: 124M @ 250M,

K{=}4

conditions. (a)

\Delta\mathrm{NLL}

vs softmax with 95% CI at 25M resolution. (b) Row span and (c) TV along training. At 50M the MinMax–Weight/MinMax–Prob NLL gap (0.047) is open while their geometry is nearly identical.

Figure S6: Five-seed condition means: geometry vs endpoint gap. (a) Native interval width

h

and (b) row score span, each averaged over layers and seeds at the 100M endpoint, against the condition’s mean

\Delta\mathrm{NLL}

vs softmax (symlog axis). Spearman correlations are across the 25 approximate conditions of a designed grid: cross-condition associations, not causal estimates.

Figure S7: Downstream accuracy of the fourteen approximate 2.5B endpoints , method minus softmax in percentage points with 95% paired bootstrap CIs: (a) BLiMP macro accuracy, (b) LAMBADA accuracy. Rows are ordered by 2.5B

\Delta\mathrm{NLL}

; the

\Delta\mathrm{NLL}

values on the training domain and the three held-out corpora are in Tables S12 and S15 .

Condition	BLiMP	HellaSwag	PIQA	ARC-E	LAMBADA	SciQ	SciQ-ns
softmax (absolute, %)	79.5	30.0	59.7	45.2	21.6	73.1	51.2
MinMax–Weight K4	-19.4 ∗	-3.9 ∗	-4.1 ∗	-9.6 ∗	-18.4 ∗	-19.3 ∗	-17.6 ∗
MinMax–Weight K16	-9.5 ∗	-2.7 ∗	-2.2 ∘	-5.7 ∗	-11.6 ∗	-9.3 ∗	-8.7 ∗
LERP K1	-5.2 ∗	-1.9 ∗	-0.2	-3.9 ∗	-12.9 ∗	-8.5 ∗	-5.2 ∗
MinMax–Weight K32	-3.8 ∗	-1.6 ∗	-0.4	-3.5 ∗	-4.4 ∗	-2.2	-1.4
FWM–Weight K2	-5.4 ∗	-1.1 ∗	-1.0	-2.6 ∗	-3.6 ∗	-1.2	-1.4

Appendix

Table S16: Layer 3, all 14 conditions: accuracy minus softmax in percentage points, paired per item and computed from the per-item differences (not from rounded absolutes); the header row gives softmax’s absolute accuracy. BLiMP is the macro-average over 67 paradigms (200 pairs each); HellaSwag/PIQA/ARC-Easy use acc_norm ; LAMBADA and SciQ use acc ; SciQ-ns is without the support passage. Rows are ordered by 2.5B

\Delta\mathrm{NLL}

with rules between the large-gap, intermediate-gap and near-baseline groups. ∗ : BH

q<0.05

within the configuration

\times

metric family of 14 comparisons (the symbol of Figure 6 ); ∘ : 95% paired CI excludes zero but

q\geq 0.05

. Per-cell CIs: stageC/tables/paired_comparisons.csv ; absolute accuracies per condition: file tab_downstream_stageC_absolute.tex in compact/migrated/ .

Figure S8: Detached vs full calibration gradient, Nearest, seed 1337 (124M @ 100M, WikiText-103). MinMax–Weight

K{=}4

(solid colour) and

K{=}16

(faded), detached and full backward: (a) logged validation loss minus the seed-1337 softmax run and (b) logged layer-mean row span along training. Endpoints with CIs: Table S19 .

Diagnostic	Question	Result
A1 forward shift invariance	$F(S+c\mathbf{1})=F(S)$ , softmax, LERP, Nearest, $K{=}2\ldots 32$	Pass for all.
A2 gradient shift residual	$\rho=\|\sum_{j}g_{ij}\|/\sum_{j}\|g_{ij}\|$ per row	softmax $4.7{\times}10^{-8}$ ; LERP full $2.2{\times}10^{-7}$ ; LERP detach median 0.026–0.410, max near 1.
A3 all-ones directional finite difference (fp64)	autograd vs derivative along $s+c\mathbf{1}$	Detach $K{=}32$ : autograd $-2.0{\times}10^{-3}$ vs true 0. Full: $\approx 0$ .
A4 interior finite difference	local LERP slopes away from extrema	softmax, detach and full all match.
P1-0 projection sanity	removes exactly the common-mode gradient?	Forward bitwise equal to detach; projected $\rho_{\max}$ $1.7{\times}10^{-16}$ ; $\max\|g_{\mathrm{proj}}-g_{\mathrm{full}}\|=1.07{\times}10^{-3}$ .

Appendix

Table S20: Gradient-consistency diagnostics run when the defect was found.

Condition	Backward	Span	$h$	$\\|q\\|$	Entropy	Top-level mass	Eff. levels	TV
Nearest K2	detach	226.2	113.1	58.0	3.97	0.679	1.81	0.439
Nearest K4	full	19.8	4.94	10.9	3.77	0.637	2.31	0.355
Nearest K4	detach	903.8	225.9	88.4	0.66	0.990	1.06	0.873
Nearest K8	full	25.5	3.18	11.3	3.86	0.506	3.57	0.259
Nearest K8	detach	1544.7	193.1	72.8	1.03	0.927	1.47	0.702
Nearest K16	full	29.6	1.85	12.2	3.78	0.386	6.70	0.160

Appendix

Table S21: Final-checkpoint geometry (layer mean, 8 validation blocks), detached vs full.

Figure S9: Single-channel ablation trajectories , LERP

K{=}32

, seeds 7 (a, b) and 42 (c, d): new single-channel runs paired with the existing same-seed full and detach controls on one GPU (the seed-1337 record of Figure 2 is a different lineage). (a, c) Validation loss minus the same-seed full-gradient run at the five evaluated checkpoints; (b, d) layer-mean row span (log). The baseline is the same-seed full-gradient run, not softmax; the detach

-

full values at these checkpoints are tabulated in Table S23 . All series are read from the per-checkpoint evaluation files.

Explore similar work

From Attention Sensitivity to Layer Role: Revisiting Mixed-Precision Quantization of Transformers
Sep 28, 2026Nafiseh HosseinpourFardi, Negar Alihadi, Mahmoudreza Babaei +2Post-Training Quantization Mixed-Precision Quantization
Sharif University of Technology, Tehran, Iran · GISMA University of Applied Sciences, Potsdam, Germany · BRAINS, Brandenburg Research Center for Applied Intelligent Systems, Potsdam, Germany +3
Approximating Softmax in Pretrained LLMs: Model Sensitivity and Kernel Acceleration
Sep 27, 2026Shangzhen Zhu, Muyan Hu, Tomasz KozlowskiGumbel-Softmax Relaxation Transformer Architectures
University of Illinois Urbana-Champaign
EFQ-Softmax: Exp-Free Quantization for Softmax
Sep 9, 2026Haohui Han, Yuming Wan, Hongni Wang +4Gumbel-Softmax Relaxation Block Sparse Flash Attention
Xi'an Jiaotong University · Xi’an Jiaotong University · Huawei Technologies Co., Ltd +1

Sep 28, 2026, cs.LGJ/K move · Enter open · S save

From Attention Sensitivity to Layer Role: Revisiting Mixed-Precision Quantization of Transformers

Nafiseh HosseinpourFardi, Negar Alihadi, Mahmoudreza Babaei, Milad Hosseini+1

Sharif University of Technology, Tehran, Iran · GISMA University of Applied Sciences, Potsdam, Germany · BRAINS, Brandenburg Research Center for Applied Intelligent Systems, Potsdam, Germany+3

Post-Training Quantization Mixed-Precision Quantization Dynamic Attention Transformer Architectures Role

Open paper AskPDF arXiv

Most post-training quantization pipelines fit each weight matrix to its pretrained counterpart, one matrix at a time. Whether that proxy tracks what an attention block actually computes, or how errors in the Q, K and V projections compound inside the softmax, is rarely checked. We write the objective on the attention output instead, over all three projections at once, and reuse it throughout the pipeline. JAB defines one scalar loss over the joint Q, K, V weights of a block, evaluated against the block's real causally-masked attention output, and uses it twice: to fit the quantized weights (GPTQ warm start, then STE with learnable scales), and to score the block for a multiple-choice knapsack allocation. On attention-only quantization of Mistral-7B this works. At 3 bits JAB recovers 77-90% of the gap between uniform GPTQ and full precision, and its sensitivity estimate tracks an oracle costing 73 forward passes to within a fraction of a point. It stops working once MLP layers enter the allocation. A role-aware offset rule needing no sensitivity estimate at all beats JAB on GPT-2's MLP and on the full Mistral-7B model: with a 3-bit floor it quantizes 96.4% of the weights to 4.5 bits per parameter at 6.933 perplexity, within 4.4% of full precision (6.643) at 3.56x compression, against 7.158 for JAB at the same budget. Which matrix a weight sits in matters more than any sensitivity estimate we computed. Two things came out sideways. Block-local reconstruction is an unreliable proxy for end-to-end perplexity: one run improved a block's own objective 4.6x while perplexity rose 32x, which is why every allocation here is validated end-to-end. And on attention-only quantization, fine-tuning moved weights farther from their pretrained values while pulling attention outputs closer, with net gains. Post-training seems to recover attention behavior, not weights.

Abstract

Explore similar work

From Attention Sensitivity to Layer Role: Revisiting Mixed-Precision Quantization of Transformers

Approximating Softmax in Pretrained LLMs: Model Sensitivity and Kernel Acceleration

EFQ-Softmax: Exp-Free Quantization for Softmax