Backward-State Policy Is Part of the Learning Algorithm

1 Introduction

• Low-precision implementations choose among the forward’s rounded value, the original, or a new random rounding for each gradient computation that rereads rounded state.The paper calls each such gradient computation a use and the collection of per-use choices the backward-state policy.Source: §1 · p. 2
• In 390M runs with an emulated FP8 attention backward, training fails when D reuses the forward’s rounding, but trains when D reads a new rounding from the same distribution.Source: §1 · p. 2
• A use must read the value its forward operation used, or a substitute giving the same gradient on average with the forward and incoming gradient held fixed.The reference is the gradient of the forward pass as it ran, with gradients passed through rounding unchanged.Source: §1 · pp. 2–3
• One low-precision normalization output can require different values: the gain gradient needs the original, while the next-layer weight gradient needs the rounded value.The paper also proves, under stated assumptions, that no single substitute is always right on average for both uses of rounded FP8 attention probabilities.Source: §1 · p. 3
• Final loss can stay within a prespecified margin even when reading the original for both gradients leaves an error relative to the reference.The discrepancy persists in 390M models trained with the rounded output feeding the next layer and the original read by both gradients.Source: §1 · p. 3
• Per-use requirements derived from the reference can be checked on single operators without training, and their predictions held in three PyTorch and Transformer Engine tests.The predicted effect of reusing the forward rounding depends on whether the use’s incoming gradient depends on that rounding.Source: §1 · p. 3

2.2 Changing only the backward removes a failure of shipped FP8 attention

• 5.097, 5.120, and 5.225 nats above BF16 were the FP8-attention runs’ validation losses; replacing only the attention backward with BF16 ended 0.0029, 0.0059, and 0.0052 nats above BF16.The comparison used the same FP8 forward; all three FP8-attention seeds failed, while the BF16-backward runs tracked BF16.Source: Figure 2 · p. 4, §2.2 · p. 4

2.1 Same rounding distribution, opposite outcomes

• In all three pairs, reuse failed while a new rounding trained usefully; the first pair’s final training losses were 7.314 and 3.256, respectively.The control also ended at 3.256, and the first pair was repeated on newer PyTorch and Transformer Engine versions with the same outcome.Source: §2.1 · p. 4

2.3 Replacing only the saved output also removes the failure

• With a replacement changing both the source output and its rounding, three seeds trained usefully, ending 0.0054, 0.0076, and 0.0067 nats above the BF16 runs in validation loss, while unchanged runs failed.Source: §2.3 · p. 5

3.1 The original matches the reference at only one use

• For a low-precision normalization store, the weight-gradient use needs the rounded value multiplied by the layer, while the gain-gradient use needs the original.Reading the original at the weight use adds −r⊤dy to dW; a new rounding is right on average at the gain use with the forward fixed.Source: §3.1 · p. 5
• Even averaged over the forward rounding, the original at the weight use can have error −E[r⊤dy], because the loss saw the rounding and dy can depend on it.A new rounding averaged with the forward fixed is right at the gain use, but retains the original’s error at the weight use.Source: §3.1 · p. 5
• In two BF16 models probed without training, the original’s average error at the weight use was detectable in all 32 layers, while its gain-use error was zero.The probe inserted a 6-bit store into all 16 layers of each model.Source: §3.1 · p. 5

3.2 No single substitute serves both uses

• In a two-parameter store model, reading the same rounded value at both uses leaves a stationary point at γ=1/3, w=1, whose expected loss exceeds the minimum by 1/18.Reading γ at the gain use and q at the weight use, or a new draw at the gain use, leaves the minimum stationary.Source: §3.2 · p. 5, Figure 3 · p. 6
• For rounded attention probabilities, no shared substitute matches both reference gradients on average with the forward held fixed for every incoming gradient, under the stated rounding and loss assumptions.The V-gradient needs the rounded probabilities, while softmax backward needs the probabilities produced by softmax.Source: §3.2 · pp. 5–6

3.3 Comparable final loss does not show that a policy follows the reference

• In 6-bit runs, the two policies whose gain use reads the rounded value end with lower root-mean-square gain than the other five in every layer on both seeds.The plotted gains are normalized against U/R at the same layer and seed.Source: §3.3 · p. 6, Figure 5 · p. 7
• At 6 bits, R/R ends 0.00425 nats below U/R; at 5 bits, it ends 0.00957 nats below, with only the 5-bit interval crossing the preset margin.The 6-bit U/U-versus-U/R interval also lies inside the margin.Source: §3.3 · p. 6, Figure 4 · p. 7

4.1 Each use has a checkable requirement

• The checker reports whether each use’s requirement holds, fails, or cannot be decided, accounting for random draws that affect its incoming gradient.It certified the weight use reading the rounded value, but could not decide the gain use when its incoming gradient depended on the same draw.Source: §4.1 · p. 8

4.2 The requirements predict when reuse matters

• Reuse can change the backward’s average computation only when the incoming gradient depends on the forward’s rounding.Whether it depends on that rounding is determined by the computation around the use, including recomputation and the loss.Source: §4.2 · p. 8

4.3 Predictions about reuse held on three operators

• When checkpointed recomputation stopped at the rounding, mean error fell to 0.19 % of its no-checkpointing size; when it included the loss, it was 106.12 %.Both registered predictions held; at a rounding boundary, neither random-state setting serves both the dropout mask and the use needing the original.Source: §4.3 · p. 9
• Under negative log-likelihood, reuse left the average log-softmax gradient unchanged, but both copies were off by cosh ε−1 times each entry’s probability.Under the constructed loss, reuse shifted each entry by −εsinh ε times its probability; both predictions held for ε = 0.1 and 0.5.Source: §4.3 · p. 9
• In TE’s FP8 attention backward, all eight predicted query-gradient changes were reproduced exactly.The predictions distinguished incoming gradients equal to the sample’s deviation from its mean from fixed incoming gradients.Source: §4.3 · p. 9
• Every registered prediction about reuse held, including predictions that reuse changes nothing.The requirements can be checked use by use on an operator before training.Source: Figure 6 · p. 8, §4.3 · p. 9

5 Conclusion

• Copy accuracy and final loss cannot settle backward-state policy: one tensor can require different values at different uses, and every reuse prediction held in three operator tests.The requirements were checked on PyTorch and Transformer Engine operators without training.Source: §5 · p. 9

A Related work

• The paper distinguishes its per-use reference question from prior unbiased reconstructions and gradient estimators whose settings differ, including ActNN [8], GACT [24], and PRAC [23].Its setting has a forward that consumes the rounded value and one tensor feeding two gradients that require different values.Source: A Related work · p. 13

B Theory

• Policy codes list the value for the original-needing use first and the rounded-value-needing use second; U, R, and N denote original, forward-rounded, and new rounding.Source: B Theory · p. 13

• The expected loss has a unique minimum at γ = 1/2 and w = 1, where it equals 9/32.Source: §B.1 · p. 13
• Reading the forward-rounded value at the gain use adds w^2(1 − γ) to its average update; reading the original or a new rounding at the weight use adds −wγ(1 − γ).At the minimum, R/R has average update (1/2, 0), while U/U has (0, −1/4).Source: §B.1 · p. 14
• With three choices at each use, the nine policies fall into four classes; U/R and N/R match the reference update everywhere, while R/R is stationary at γ = 1/3 and w = 1.At that R/R stationary point, expected loss exceeds its minimum by 1/18.Source: §B.1 · p. 14
• At the minimum, no single value read at both uses keeps the minimum stationary, whether it is the original, rounded, new, or another value.The average weight update can vanish only when κ = 0, while the gain component is then 1/2.Source: §B.1 · p. 14

• The theorem assumes independent, unbiased, non-exact rounding of each probability to neighboring representable values and a quadratic loss in the attention output.The forward output is the value matrix multiplied by the rounded probabilities.Source: §B.2 · p. 14
• No shared substitute is right on average, with the forward held fixed, at both the value-gradient and softmax-backward uses for every incoming gradient.A shared substitute is read at both uses, as a vector at the value gradient and through its Jacobian at the softmax backward.Source: §B.2 · p. 14
• For a substitute right at the value gradient, the matrix supplied to the softmax backward falls short of J(p) by at least the rounding covariance Σ_r.The shortfall equals Σ_r for the forward-rounded probabilities; any randomness in the substitute adds to it.Source: §B.2 · p. 15

B.3 The errors of four policies, predicted and enumerated exactly

• N/N’s average error at the softmax backward is −Σ_r Gp, so a new rounding alone does not satisfy the requirement there.A nonzero average given the forward rules out zero average conditioned on both the forward and its rounding.Source: §B.3 · p. 15
• All eleven preregistered statements resolved as predicted, including refutation of the false claim that N/N has no error at either use.The checks included zero-rounding and third-moment controls, per-draw and average zeros, and six planted code errors.Source: §B.3 · p. 15
• N/N reduces the softmax-backward error versus R/R in both cells and rounding laws, but introduces error at the gradient of V.Its softmax error is 0.4808 and 0.5598 of R/R’s in cell A, and 0.4222 and 0.4254 in cell B; its gradient-of-V error is −σ^2∥o∥^2.Source: §B.3 · p. 15

B.4 One new rounding suffices at the softmax’s backward

• Reading p_q at the gradient of V and K(a) at the softmax backward is right on average at both uses, under the stated independent-entry assumptions.The substitute requires a rounding not consumed by the forward.Source: §B.4 · p. 16
• Two independent new roundings in J2 average to J(p) and give zero error, whereas plugging one new rounding into J leaves an error.Using the forward’s rounding as one of J2’s two roundings leaves a nonzero error term because it is paired with the incoming gradient.Source: §B.4 · p. 16

B.5 What R/R adds to the reference’s update is not a gradient

• R/R’s average update is not a gradient in the two tested cells, with mixed-partial differences reaching 17/320; the null cell’s result depends on the symmetric rounding laws.Under stochastic rounding, exact enumeration also finds a nonzero derivative in the null cell.Source: §B.5 · p. 16
• U/U’s average update is the gradient of the unrounded loss but errs at the gradient of V, while U/R follows the reference at both uses.U/R’s average update can still fail to be a gradient until the correction c is added.Source: §B.5 · pp. 16–17

C.1 What the checker reads and reports

• The checker separates mean error from coupling, and tests a use against the reference with the forward fixed and over the substitute’s draw.A zero average over the forward’s rounding is weaker than the per-use requirement.Source: §C.1 · p. 17
• The checker distinguishes replay, independent reconstruction, restored random state, detachment, and impossible redraws; shared draws can receive an undecided verdict.Separate draws certify zero coupling but leave mean error undecided; replay certifies a zero gap.Source: §C.1 · p. 17

C.2 How the checker’s verdicts were tested

• All ten policy variants received their expected verdicts, and seven families of exact checks passed their declared mutation-test outcomes.The tests used rational arithmetic and enumerated every rounding outcome; one edit preserving the average was correctly accepted.Source: §C.2 · p. 17
• In a composed two-stage backward, separate draws give the reference on average, but one shared draw doubles the mean gradient despite each stage passing at fixed input.Conditioning each stage on its received input or checking the whole backward detects the error.Source: §C.2 · p. 17

C.3 How the three tests were registered

• The tests registered operators, setups, predictions, and reading rules before measurement; the reported tables include every registered statement and its outcome.The checkpoint and log-softmax tests were repeated in a separately registered configuration.Source: §C.3 · pp. 17–18

C.4 Checkpointing

• The test records the added mean-error term when a use reads rounded y instead of the original across checkpointing conditions with random state restored or not.The CPU model uses dropout, stochastic rounding, and checkpointed regions that end at the rounding or continue through the loss.Source: §C.4 · p. 18
• Recomputing the loss couples the new draw to its incoming gradient, producing the same expected error as the baseline.The registered rules distinguish removal from retention of the error.Source: §C.4 · p. 18
• PyTorch’s determinism check compares shapes, data types, and devices, so it does not detect changed tensor values..Source: §C.4 · p. 19
• S6’s no-restore condition changed the record because dropout altered the input to rounding; the dropout-free second configuration tested that control, and W6 held.This separates the dropout-related record change from drawing the rounding anew.Source: §C.4 · p. 19

C.5 Log-softmax

• The log-softmax test injects independent ±ε errors into the saved output and compares reuse with an independent copy against the exact-output gradient.The CPU backward was checked in double precision against its stated formula.Source: §C.5 · p. 19
• With fixed-target negative log-likelihood, reuse changes nothing on average, but neither copy is right; with the constructed loss, reuse shifts the average by −ε sinh ε.The constructed loss makes the incoming gradient equal the rounding error; losses whose incoming gradients depend on the saved output can therefore behave differently.Source: §C.5 · p. 19
• LS5b was predicted to miss because dividing by cosh ε corrects its mean but leaves the coupling scaled by 1/cosh ε.The second configuration also distinguished exact constants from their leading terms.Source: §C.5 · p. 19

C.6 FP8 attention

• In TE’s FP8 attention backward, the independent-minus-shared query-gradient contrast averages to a^2 when the incoming gradient is the shared deviation, and to zero for fixed h=a.The test calls the captured backward directly with inputs arranged so the values deciding the result are exactly representable.Source: §C.6 · pp. 20–21
• Checks confirmed the FP8 fused backend, intended inputs, and unit scales on tensors and quantizers; decoded outputs and incoming gradients matched intended values, repeated gradients were byte-identical, and an emulated backward ignoring saved output was rejected.Source: §C.6 · p. 21

D Setup

• The appendix describes training-run and probe specifications, with operator tests in Appendix C and additional runs in Appendix E.Source: D Setup · p. 21

D.1 Model, data and software

• Experiments use the FOG 390M Llama 3 setup, 10B FineWeb-Edu tokens, FP8 Transformer Engine attention, and BF16 controls on eight H20 GPUs.Source: §D.1 · p. 21

D.2 The forks of Section 2.1

• Forks resume at step 999 and alter only D’s output copy, comparing the forward’s rounding with an independent stochastic rounding and a reference that computes D from P.The independently rounded output is recomputed from dequantized FP8 inputs; runs are paired across parent and rounding seeds.Source: §D.2 · p. 22
• Pre-registered fork rules classify non-finite losses and specified excursions as failures, and require late-logit and loss bounds for a useful run.Source: §D.2 · pp. 22–23

D.3 TE attention trained from scratch

• From-scratch 390M TE runs use a 20k-step schedule across three seeds, with BF16 references and pre-specified endpoint loss and logit criteria.The 1.5B comparison uses a separate rule with a 0.05-nat endpoint gap and no logit condition.Source: §D.3 · p. 23

D.4 The normalization store

• On both BF16 models, U/U’s weight-use error statistic exceeds the largest zero-mean control statistic, while its gain-use error is exactly zero.The preregistered comparison required the weight-use statistic to exceed the largest control in at least 13 of 16 layers.Source: §D.4 · p. 24
• Normalization-store policies are compared at 6- and 5-bit widths after 8,000 BF16 training steps on two seeds, using a pre-fixed ±0.010-nat equivalence margin.Source: §D.4 · p. 24
• The continuation probe compares original-reading errors across U/R and U/U against a no-store control, using paired window estimates and pre-registered criteria.The rules also test persistence and gain drift; incomplete or void states leave the result undecided.Source: §D.4 · p. 24

D.5 Compute

• The 112 training runs used 5,856 GPU-hours; a 390M 20k run took a median 79 GPU-hours, with a 66–95 range.The accounting excludes startup work and work after each segment, as well as probes and operator tests that do not train the model.Source: §D.5 · p. 24

E.1 Runs from scratch and on newer software

• The further-run table reports loss gaps for useful runs and excursion-opening steps for failures, including from-scratch runs and a newer-software fork.Source: §E.1 · p. 25, Table 8 · p. 25
• At 1.5B parameters, the FP8 run’s failing excursion opens at step 4900 with loss 7.40 versus a 2.96 running minimum, and ends at 3.47 at step 4999.Source: §E.1 · p. 25
• Repeating seed 0 BF16 on cu129 shifts late training loss by −0.0010 nats and validation loss by −0.0014; widening seed-2 gaps by those amounts keeps both within 0.02 nats.Source: §E.1 · p. 25

E.2 The fork at step 1999

• Both stochastic copies train usefully in construction (a), while original and nearest-rounded copies fail; in construction (b), stochastic copies and BF16-stored PV train usefully.Source: §E.2 · pp. 25–26
• In construction (a), stochastic-copy RMS error is 1.35 times the original TE error with correlation 0.0009, while nearest rounding has ratio 0.96 and correlation 0.7536.Source: §E.2 · p. 26

E.3 Repairing selected attention rows

• Repairing D from P only on saturated rows trains usefully, whereas repairing only other rows fails; saturated rows occupy at most 25.2% of any layer at the fork.The authors identify this comparison as a sufficient subset for repair in this fork.Source: §E.3 · p. 26

E.4 The same dependence in a normalization backward

• Moving the normalization store from BF16 to E4M3 multiplies median paired statistic T by 256.71, matching the 256.50 prediction; round-to-nearest gives T exactly zero.The BF16 paired statistic detects dependence, while the registered per-window positive-sign test is unresolved for BF16 and positive in 24 of 32 E4M3 windows.Source: §E.4 · p. 27