Pretraining Transformers with Quantized Softmax in Attention

Motivation

• During pretraining, an approximate operator defines the learning rule as well as the forward computation, unlike an inference-time perturbation to a fixed function.For calibrated grids, row extrema also set the grid, so detaching them is not harmless as it is for the softmax row maximum.Source: §1 · p. 1

Approach

• K-interval attention approximates the exponential with K+1 grid values and varies row-wise calibration, interpolation versus hard rounding, and surrogate placement around normalization.The experiments use fp32 attention arithmetic with TF32 matmuls and do not measure kernels or claim speed-up.Source: §1 · p. 1

Main findings

• +0.89 nats above softmax at 2.5B tokens, MinMax–Weight-STE at K=4; changing either calibration or surrogate placement removes most of the deficit.The effects keep their direction at 1B parameters and across the tested seeds.Source: §1 · p. 2
• Detaching row extrema leaves the forward computation unchanged but omits calibration derivatives, and the resulting run diverges after initially tracking softmax.The possibility that failures reflect a gradient defect growing with learned score geometry is presented as a hypothesis, with counter-examples noted.Source: §1 · p. 2

Low-precision Transformers and the softmax path

• Low-precision Transformer systems quantize linear layers and attention matmuls, while the exponential and row normalization often remain high precision.EXAQ and BAPS state that training is untested, while IndexSoftmax is designed as a training-free replacement.Source: §2 · p. 2

Approximate and non-exponential softmax

• Prior approximate softmax methods include static-scale exponential approximations, removal of the row reduction, and alternatives whose Jacobian changes may or may not be absorbed into the learning rate.Under per-row calibrated reconstruction, the Jacobian terms are neither constant nor absorbable into the learning rate.Source: §2 · p. 2
• Non-exponential attention methods are trainable with stabilizers, but the paper makes no linear-complexity claim for its LERP at K=1.Source: §2 · p. 2

Surrogates and calibration gradients

• Unlike learned clipping or step-size parameters, dropping gradients through per-row score calibration changes ∂L/∂s itself.The paper reports no prior work that back-propagates through a per-row calibration statistic in attention pretraining.Source: §2 · pp. 2–3

Operator family

• The operator varies calibration, reconstruction, and surrogate placement; all weights are normalized to probabilities before producing the attention output.The three axes are MinMax versus FWM, LERP versus Nearest, and Weight-STE versus Prob-STE.Source: §3 · p. 3

Calibration

• MinMax sets interval width from the row span, while FWM fixes a window at the row maximum and assigns zero weight below it.Under MinMax, the learned negative tail sets resolution near the row maximum; FWM uses h = τ/K.Source: §3 · p. 3

Reconstruction

• LERP interpolates between adjacent exponential grid values, whereas Nearest rounds scores to grid centers and gives scores in one bin the same weight.The Nearest index is piecewise constant and has zero derivative almost everywhere.Source: §3 · pp. 3–4, Figure 1 · p. 3

Calibration gradients

• The full calibration gradient includes derivatives through row extrema and yields a zero-sum score gradient, as required by shift invariance.For FWM, the identity holds with the maximum derivative and no minimum term; detaching extrema drops terms and violates the identity.Source: §3 · p. 4
• For hard forwards, the index path has zero derivative almost everywhere, so the method uses a LERP surrogate and chooses which calibration derivatives it carries.The full calibration derivatives are exact for differentiable LERP, but for straight-through rules they are surrogate-carried derivatives rather than the hard-forward derivative.Source: §3 · p. 4

Surrogate placement

• Weight-STE places the straight-through replacement before normalization, while Prob-STE places it after normalization; both have the same hard forward but evaluate the normalization Jacobian differently.Their backward rules coincide when hard and interpolated weights agree, and their difference can be amplified by wider intervals.Source: Figure 1 · p. 3, §3 · p. 4

Training setup

• GPT-2-style models have 124M or 1B parameters, with attention in fp32 and TF32 matmuls enabled during training.Other operations run under bf16 autocast; TF32 is disabled at evaluation.Source: §4 · p. 4

Experimental suites

• The suites test a balanced calibration-by-surrogate-by-K design, longer training, an early larger-model probe, and seed variation.Suite A trains 2.5B tokens, suite C probes 8× parameters at 0.1 tokens per parameter, and suite D alone measures seed spread.Source: §4 · p. 5

Evaluation

• Validation results are paired differences in NLL per token, with positive ΔNLL indicating worse performance than softmax.Evaluation uses each run's frozen training code because a shared implementation or larger batch can perturb single Nearest blocks.Source: §4 · p. 5

Statistics

• Single-seed intervals cover evaluation-block sampling, while suite D reports five-seed variation; single-seed differences below about 0.01 nats are not interpreted.The authors note no learning-rate sweep, so 2.5B-token persistence is specific to the tested schedule.Source: §4 · p. 5

Geometry probe

• The geometry probe records row span, native width, norms, entropy, and probability total variation on validation scores.The total variation measures operator distortion on the learned scores, not training loss or a mediator.Source: §4 · p. 5

5.1 CALIBRATION-GRADIENT INTERVENTION: SAME FORWARD, DELAYED FAILURE

• Detaching row extrema leaves the forward unchanged but delays divergence: the LERP K=32 run ends +0.737 nats above softmax, and restoring zero-sum alone does not repair the failure.Retaining only the maximum-dependent calibration gradient recovers almost the full full–detach gap on both seeds; retaining only the minimum-dependent gradient recovers almost none.Source: §5.1 · pp. 5–6

5.2 CALIBRATION × SURROGATE INTERACTION AND SEED ROBUSTNESS

• At K=4, MinMax–Weight ends +0.39 nats above softmax at 250M tokens and +0.89 at 2.5B; changing calibration or surrogate placement removes most of the deficit.Across five seeds, FWM reduces NLL by 0.067 nats relative to MinMax and Prob-STE by 0.021 relative to Weight-STE, with a +0.031 interaction; effects shrink with K.Source: §5.2 · pp. 6–7

5.3 K SCALING AND TRAINING HORIZON

• At 2.5B, MinMax–Weight falls from +0.89 nats at K=4 to +0.04 at K=64, while LERP K=4 improves from +0.086 at 125M to +0.005 at 2.5B.LERP's gap keeps closing while Nearest gaps plateau, so their relative NLL changes with training budget.Source: §5.3 · p. 7

5.4 LEARNED GEOMETRY: OBSERVATIONS, COUNTER-EXAMPLES AND A HYPOTHESIS

• MinMax–Weight alone among formal full-gradient conditions keeps growing its score span; at 2.5B its K=4 span is about seven times softmax's, with total variation 0.86.The authors describe a range–resolution feedback hypothesis but state that no intervention identifies the mechanism; across designed conditions, native width associates more strongly with ∆NLL than score span.Source: §5.4 · p. 8

5.5 DOWNSTREAM TRANSFER: GRADED, NOT BINARY

• All 35 large-gap comparisons across six benchmarks and seven configurations had negative point estimates, and MinMax–Weight K=4 lost 3.9–19.4 accuracy points, less at larger K.Of the near-baseline group's 35 comparisons, eight reached q < 0.05, but the authors caution that single-seed differences cannot be attributed to the operator rather than the run.Source: §5.5 · p. 8, Figure 6 · p. 9

6 DISCUSSION AND LIMITATIONS

• Under MinMax, interval width grows with the negative tail, whereas FWM fixes the width independently of that tail.This is the proposed link between learned geometry and surrogate–forward mismatch.Source: §6 · p. 9
• The tail-policy comparison is single-seed and uses historical STE arithmetic rather than equal-cost runs; the channel ablation uses two seeds of one setting.Strict-forward retraining and reduced-precision baselines each use one seed at 100M tokens.Source: §6 · p. 9
• The largest frozen-checkpoint mean validation NLL difference between historical native and direct hard forwards is 3.4×10^-4 nats across 12 audited 124M K=4 endpoints.This is an evaluation-time quantity, not a per-block or training-time bound; bit-exact retraining reproduced contrast directions at one seed and horizon.Source: §6 · p. 9
• The 2.5B-token persistence finding is specific to the tested schedule, and no learning-rate sweep was run.The 1B-parameter suite is an early-regime probe on another corpus; no run exceeds 1B parameters or 2.5B tokens.Source: §6 · p. 9

7 CONCLUSION

• Detaching row extrema leaves the forward unchanged but causes delayed training failure, while zero-sum projection does not recover the full-calibration result.For hard rounding at K=4, MinMax with Weight-STE has a large loss gap; FWM or a probability-level surrogate reduces it.Source: §7 · p. 9
• Selected hard and interpolated configurations approach softmax validation loss, but small NLL gaps do not establish downstream-task equivalence.The remaining differences depend on operator, K, training horizon, and evaluation metric.Source: §7 · p. 9
• The study concerns numerical attention operators on public text corpora and benchmarks, with no human subjects, personal data, or deployment.The stated potential downstream use is more efficient training hardware; no specific additional risk is identified.Source: ETHICS STATEMENT · p. 10
• The authors report that closed-form gradients were checked against automatic differentiation, and that code, manifests, and per-block results will be released publicly.Stored result files regenerate the figures and tables.Source: REPRODUCIBILITY STATEMENT · p. 10

A OPERATOR DEFINITIONS AND EDGE CASES

• All operators normalize reconstructed weights as Pj=wj/W and produce the attention output O=PV from causal-row scores.Future positions receive zero weight; the row uses fp32 scores and subtracts its maximum before exponentiation.Source: A OPERATOR DEFINITIONS AND EDGE CASES · p. 13
• LERP interpolates between adjacent exponential grid values, whereas Nearest rounds each score to a grid center; attention probabilities are still normalized, not one-hot.Under FWM, the hard operator can additionally produce a zero tail.Source: A OPERATOR DEFINITIONS AND EDGE CASES · p. 13
• The stated formulas assume a nondegenerate row with h>0 and a locally fixed interval index; implementation branches separately handle spans above or below ε=10^-12.The edge-case table reports finite forward and backward values for the listed LERP, MinMax, and FWM settings.Source: A OPERATOR DEFINITIONS AND EDGE CASES · p. 13, Table S1 · p. 13
• MinMax sets grid width from the row span, while FWM anchors a fixed-width window at the row maximum; at K=4, FWM uses h=1.5 nats.At softmax-like spans of 25–32, MinMax uses a quarter of the row span, with much larger widths possible as the span expands.Source: A OPERATOR DEFINITIONS AND EDGE CASES · pp. 13–14

B GRADIENT DERIVATIONS

• For normalized weights P=w/W, the normalization Jacobian gives the upstream weight gradient used by the operator’s backward rule.The derivation evaluates upstream probability gradients at the forward probabilities.Source: B GRADIENT DERIVATIONS · p. 14
• Full MinMax calibration differentiates through shifted scores, grid values, and interval width; its per-key partials sum to zero when combined.This yields the row-gradient zero-sum identity for the full calibration gradient.Source: B GRADIENT DERIVATIONS · p. 14
• Detaching extrema drops calibration terms, leaving a common-mode component and incorrect relative gradients between keys.Projecting onto the zero-sum subspace removes only the common-mode component.Source: B GRADIENT DERIVATIONS · p. 14
• Under FWM, in-window weights depend on the row maximum, while tail keys have zero weight and zero gradient through the mask and clamp.The row minimum plays no role, and the resulting row gradient sums to zero.Source: B GRADIENT DERIVATIONS · p. 14
• Weight-STE and Prob-STE share the hard forward and LERP scalar derivatives but evaluate the normalization Jacobian at different points.Their centering distributions and normalization factors can therefore differ.Source: B GRADIENT DERIVATIONS · p. 14
• The hard MinMax map has extremum-mediated derivatives when interval indices are fixed, but STE backpropagation cuts the hard branch’s dependence on extrema.For FWM, the hard map is locally constant when the index and tail mask are fixed; reported formula checks agree with autograd to <10^-9.Source: B GRADIENT DERIVATIONS · p. 15

• Suite runs shared model, data, and optimizer protocols; the 124M and approximately 983M models used distinct layer, head, context, batch, and learning-rate settings.AdamW, gradient clipping, warmup, and cosine decay were specified; suite A included 124M runs through 2.5B tokens, while other suites varied model size and token budget.Source: §C.1 · p. 15
• Training used TF32 tensor-core products for QK⊤ and PV, an fp32 attention operator, and bf16 autocast outside attention; evaluation and diagnostics used TF32 off.Hardware and software records were incomplete for some cloud pods, and GPU-hour totals are approximate.Source: §C.1 · p. 15, Table S4 · p. 16
• In suite D, operator family tracked training stack, and same-GPU evaluation controlled evaluation-stack differences only, not training-stack differences.Some rented RTX 5090 pods were author-confirmed rather than identified in their logs; unlogged software versions remained unrecorded.Source: §C.1 · pp. 15–16, Table S2 · p. 16, Table S4 · p. 16
• Suites A–D varied model size, token budget, operator families, interval counts, and seeds; validation used FineWebEdu-3B and WikiText-103 blocks.Suite D contained 130 formal runs; the complete study reports 167 runs and 1,106 evaluated checkpoints.Source: §C.1 · p. 16

• Historical Prob-STE and Weight-STE expressions define the same hard forward mathematically, but differ at GPU rounding level and neither is bit-exact P_N at trained score scales.Batch sizes 1 and 8 differed by up to 8.1×10−3 on single blocks, motivating batch-1 evaluations.Source: §C.2 · p. 17
• On frozen checkpoints, Prob-STE minus shared Weight-STE evaluation differed in mean NLL by at most 1.1×10−4 nats on FineWebEdu and 3.6×10−4 on WikiText, with inconsistent signs.The audit compares historical expressions, not either expression against a direct hard forward.Source: §C.2 · p. 17
• Direct hard-forward replacement changed mean validation NLL by at most 3.4×10−4 nats across 12 K=4 endpoints, while single-block differences reached 0.026 nats.This is an evaluation-time result on frozen checkpoints and one RTX 3090 stack, not a per-block bound.Source: §C.2 · p. 17, Table S5 · p. 18
• Strict-forward retraining preserved both contrast directions and the MinMax magnitude; the larger FWM value mostly came from the new Prob-STE run (−0.012 against the historical one).The seed-7 supplementary comparison cannot separate rounding, environment, and training nondeterminism; it does not establish a general FWM–Prob improvement.Source: §C.2 · pp. 17–18, Table S6 · p. 18
• The audits do not bound accumulated training-time effects; the only training-time test covers suite D seed 7, not suites A–C or other seeds.Across stacks, the same checkpoint's mean NLL differed by up to about 5×10−4 nats, an observed range rather than a bound.Source: §C.2 · pp. 18–19

C.3 SELECTION OF τ, STATISTICS AND EVALUATION

• τ = 6 was selected before FWM training using an inference-time screen on frozen softmax-trained weights, not by establishing a training optimum.WikiText-103 validation participated in selection, and the wider candidates were eliminated on a 32-block subset.Source: §C.3 · p. 19
• The evaluation protocol distinguishes detectable differences from results inside a pre-declared margin, and reports downstream comparisons by task rather than as a combined score.For NLL-type metrics the margin is ±0.01 nats; no margin is declared for accuracy metrics.Source: §C.3 · pp. 19–20

D.2 FIVE-SEED SUITE

• Across five seeds, FWM keeps the K=4 gap small, while approximation costs emerge between 40M and 60M tokens for conditions that degrade.Under FWM the K=4 gaps are +0.032 nats for Weight-STE and +0.013 for Prob-STE; no condition is degraded at 20M tokens.Source: §D.2 · p. 21, Figure S1 · p. 21, Table S8 · p. 21

D.3 HORIZON AND K DETAILS AT 2.5B

• MinMax–Weight K=4 reaches a +3.66-nat collapse at 375M tokens before partially recovering to a +0.89 plateau; FWM–Weight K=4 instead holds flat after an early cost.The MinMax–Weight collapse coincides with single-step gradient norms of 104 and 98% of steps clipped over the last 1.5B tokens; FWM–Weight K=4 is +0.043 at 25M tokens.Source: §D.3 · p. 21, Figure S2 · p. 21

D.4 TAIL POLICY IN THE MINMAX →FWM CHANGE

• The untruncated fixed-step operator recovers most of the MinMax–FWM gap, so zeroing the tail is not required for the observed improvement in this setting.It ends −0.012 nats versus FWM under Weight-STE and +0.001 under Prob-STE; the resolved comparison uses more distinct nonzero grid levels, so it is not at equal deployment cost.Source: §D.4 · p. 22, Table S10 · p. 22

D.5 REDUCED-PRECISION EXPONENTIAL BASELINES

• −0.0034 nats for BF16-exp and +0.0028 nats for FP8-rounded exp are the 100M validation ∆NLLs versus the same-seed softmax run.Both differences were detected by evaluation intervals and were below 0.01 nats in magnitude; each endpoint is from one run evaluated on one RTX 3090.Source: §D.5 · p. 23, Table S11 · p. 23
• The reduced-precision exponential runs stay within ±0.011 nats of softmax at the five evaluated checkpoints, while LERP K=4 opens its gap between 40M and 60M.The checkpoint comparisons do not bound behavior continuously between evaluations.Source: §D.5 · p. 23, Figure S3 · p. 24
• The exponential-baseline evidence is limited to one seed and 100M tokens, without a coarser format or a cost-matched comparison to five-level reconstruction.It establishes neither equivalence nor a cross-seed ranking against training randomness.Source: §D.5 · p. 24

D.6 SCORE GEOMETRY ALONG TRAINING AT 2.5B

• MinMax–Weight K=4 reaches row span 287, while softmax’s reported span of 25–32 is in the range of the largest dispersion reported for Gemma 7B.The softmax mass in MinMax–Weight K=4’s top rounding bucket reaches 99.5%; this reflects the operator mapping mass to one grid value, not the model concentrating probability on one key.Source: §D.6 · p. 24, Figure S4 · p. 25
• In the 250M MinMax–Weight run, row span and NLL gap have a within-run Spearman correlation of about 0.9 across ten checkpoints.Across five-seed conditions, endpoint gap tracks interval width h and span only as a cross-condition association.Source: §D.6 · p. 24, Figure S6 · p. 25
• The proposed range–resolution feedback remains a hypothesis: the reported associations are descriptive, and several isolating interventions were not run.The authors identify freezing h mid-training, swapping the surrogate from a saved checkpoint, QK-normalization, and a learning-rate sweep as possible ways to separate remaining factors.Source: §D.6 · pp. 24–25
• Geometry diagnostics move before the loss in the detach failure but after it in formal MinMax–Weight runs, so they are not a universal early warning.The cited timing comparison is shown for the 250M MinMax–Weight run and the detach failure.Source: §D.6 · p. 25, Figure S5 · p. 25

E.1 DOWNSTREAM EVALUATION OF THE 2.5B SUITE

• The pre-declared groups classify endpoints by their 2.5B ∆NLL versus softmax, from large-gap conditions (+0.13 to +0.89) to near-baseline conditions below +0.01.The intermediate group spans +0.02 to +0.06 nats; the ranges are descriptive and three members sit at rounded edges.Source: §E.1 · pp. 25–26
• On WikiText-103, PTB, and C4, condition ordering correlates strongly with the in-domain ordering, while MinMax–Weight K=4 has larger gaps on WikiText-103 and PTB.The reported Spearman correlations are 0.996, 0.952, and 0.996; the MinMax–Weight K=4 gap is unchanged on C4.Source: §E.1 · p. 27
• The intermediate-gap group separates from softmax on BLiMP and LAMBADA in 6 of 28 comparisons, while no condition loses more of the SciQ support-passage benefit after correction.FWM–Weight K=16 does not separate on BLiMP (−0.5 pp, q = 0.10); the mean near-baseline difference of +0.3 pp is descriptive only.Source: §E.1 · p. 27, Figure S7 · p. 27, Table S16 · p. 28
• Probe A and B outcomes are difficult to interpret as evidence about the targeted abilities because the softmax reference itself barely solves either probe.For example, MinMax–Prob K=4 has −1.28-nat target NLL and +1.47-nat induction gain on probe A, but the authors do not interpret this as long-range copying.Source: §E.1 · p. 28, Table S17 · p. 29
• On copying probe C, the slope of target-NLL difference against log2 L is steeper under Prob-STE than Weight-STE at matched calibration and K=4.The slopes are 0.94 versus 0.11 for MinMax and 0.88 versus 0.23 for FWM.Source: §E.1 · p. 29, Table S18 · p. 29
• The large-gap prediction is supported on benchmark and probe layers, but near-baseline differences are detectable on held-out LM NLL and in some benchmark comparisons.P3 is supported on copying probe C only; P4 cannot be tested because the reference model’s induction gain falls to 17% by distance 896, and P5 fails as measured.Source: §E.1 · p. 29

F THE DETACHED-GRADIENT RECORD AND ITS REPLICATION

• At 40M tokens, detach is +0.118 and +0.074 nats above matched full-gradient runs for seeds 7 and 42; most of the gap is open by 60M.Both seeds reproduce the delayed degradation; the older seed-1337 result is consistent but comes from a different code lineage.Source: F THE DETACHED-GRADIENT RECORD AND ITS… · p. 31, Table S23 · p. 31
• Maximum-only recovers 0.98–0.99 of the detach–full gap, unlike minimum-only (0.003 and 0.072) across the two seeds.Maximum-only ends within 0.015 nats of full with a softmax-like span and query norm; minimum-only has a severely expanded span and a gradient-norm rise.Source: F THE DETACHED-GRADIENT RECORD AND ITS… · p. 32, Table S25 · p. 33
• The omitted maximum-dependent term has a norm ratio of 0.36–0.52 versus 0.008–0.009 for the minimum-dependent term; preserving it leaves local relative gradient error below 1%.This is a fixed-upstream local diagnostic, not a per-layer guarantee or a bound on whole-model score-gradient error.Source: F THE DETACHED-GRADIENT RECORD AND ITS… · pp. 32–33, Table S24 · p. 33

G ASSETS, CODE AND REPRODUCTION

• Formal K-interval runs use the full calibration gradient, while softmax runs use softmax’s own backward; their legacy range_grad field is not read by softmax.The audit records 243 run directories and 940 step checkpoints; manifests distinguish complete trainings, checkpoint files, and evaluation jobs.Source: Table S26 · p. 34, G ASSETS, CODE AND REPRODUCTION · p. 34

H EXTENDED RELATED WORK

• Prior approximate-softmax methods use static scales, downstream fine-tuning, or learned per-head parameters; this work focuses on gradients through per-row calibration statistics.The related-work discussion distinguishes those approaches from propagating derivatives of the per-row extrema used inside the attention operator.Source: H EXTENDED RELATED WORK · p. 35