Transforms in Training Workloads

Topics We'll Cover

• What a Training Step Adds to Inference
• The 6N Rule, Reproduced
• What the Inference Series Found
• The Gradient of a Convolution Is a Correlation
• Three 4f Passes: Forward, dX and dW
• Weight-Masked and Data-by-Data Products
• Block-Circulant Weights in Training
• The Training-Step Amdahl Table
• Why the Share Falls from Forward to Step
• Interactive: Training-Step Amdahl Explorer
• Masks Rewritten Every Optimizer Step
• Layer-Major Order
• What Does Not Map in Training
• What to Take Away

What a Training Step Adds to Inference

The inference series asked what share of a forward pass a Fourier-optical transform engine could take (FOptInf 02). Training runs the same forward pass and then two more kinds of work:

• The step is large. The analysis takes an optimizer step of 1,048,576 tokens (illustrative, about what Llama-scale runs use), so at 2,048 tokens a step is 512 sequences. Anything done once per step is amortised over all of them.
• The weights change after every step. An inference engine writes its Fourier-plane masks once; a training engine must rewrite them (slide 11).
• The question for this deck: which of the new work is a transform, which of it could sit on a mask, and does the share go up or down? Every number comes from analysis/train_flop_share.py, which extends the inference ledger with the backward pass and writes results.md. Same model shape (Llama-3-8B), same FLOP conventions (matmul 2mnk, real FFT 2.5 n log₂ n, complex multiply 6, add 2).

The 6N Rule, Reproduced

Before trusting the ledger on new layers, check it on the old one. Training a dense model costs about 6N FLOPs per token (2 forward, 2 for dX, 2 for dW), plus an attention term that grows with the sequence (Kaplan et al., arXiv:2001.08361; the same accounting appears in the Chinchilla paper's appendix, Hoffmann et al., arXiv:2203.15556). Here N counts the parameters that multiply activations: the 8,029,995,008 of the untied model less the 128,256 × 4,096 input embedding, which is a lookup.

Source: results.md, section “1. Cross-check: the 6N rule”; the pytest suite checks the ratio to 1e-12.

• The match is exact by construction: the ledger charges every matmul's backward at twice its forward, which is what the rule says. Its value is as a check on the bookkeeping (no layer counted twice or missed).
• Not counted, as in the rule: softmax, norms, activations, the AdamW update and activation recomputation. FlashAttention's backward recomputes the attention scores (Dao et al., arXiv:2205.14135), and activation checkpointing re-runs parts of the forward pass (Chen et al., arXiv:1604.06174); both add digital work, which would lower every share in this deck a little further.

What the Inference Series Found

FOptInf 02 counted the forward pass. The short version:

This script's forward column agrees: Hyena-2 at 2,048 tokens, 0.20% of the forward pass (slide 08). The physics of the engine (the lens as a Fourier transformer, the 4f system, masks and SLMs, ENOB) is in FOptInf 01 and is not repeated here.

The Gradient of a Convolution Is a Correlation

One channel of a causal long convolution, length L, with the output gradient g = ∂ℒ/∂y arriving from above:

forward      y[i]   = Σk h[k] u[i−k]                    Y  = H · U
input grad   ∂ℒ/∂u[j] = Σi g[i] ∂y[i]/∂u[j] = Σi g[i] h[i−j]
                                         dU = G · conj(H)
filter grad  ∂ℒ/∂h[k] = Σi g[i] u[i−k]          dH = G · conj(U)

• Each gradient is a correlation: the output gradient slid against the other operand. By the correlation theorem a correlation of real signals is a product with a conjugate spectrum, so dU = IFFT(G·conj(H)) and dH = IFFT(G·conj(U)).
• Causal padding carries over. The forward pass is causal (y[i] uses inputs at or before i); its gradient is anti-causal (du[j] uses gradients at or after j), the mirror image. Both need lags from −(L−1) to L−1 free of wrap-around, so the same zero-padding to n ≥ 2L that made the forward FFT linear makes the backward one exact; keep the first L outputs.
• What it costs per channel: one FFT of G and one inverse FFT for dU (the same two transforms as the forward), the product with conj(H), the product with conj(U) (U's spectrum is kept from the forward pass), and an accumulation. dH is summed over the step's sequences in the Fourier domain and inverse-transformed once per step.

Source: results.md, section “6. Numeric checks of the backward formulas”. The spectral formulas agree with direct sums to rounding; the finite-difference column is limited by the difference step, not the formula. Length 1,000 is not a power of two: it is padded to 2,048.

Three 4f Passes: Forward, dX and dW

Mapped onto a 4f system (lens, Fourier-plane mask, lens; FOptInf 01, “The 4f System”), one channel of one layer needs three passes per training sequence. What matters is what sits on the mask:

• Forward and dX have the shape an optical engine wants: a data stream through a mask that holds a weight. The backward-dX mask is the forward mask's conjugate.
• dW is different in kind. Both operands of G·conj(U) belong to the sequence being processed. Whichever goes on the mask, the mask must be rewritten for every sequence of every layer.

Weight-Masked and Data-by-Data Products

The ledger splits every Fourier-domain product by its operands, because that decides whether a mask written once per optimizer step can hold it:

Static share = transforms plus weight-masked products, over all step FLOPs: what an engine whose masks change only at optimizer steps could take. The optical share adds the data-by-data products as well.

Block-Circulant Weights in Training

The inference series' one large-share case made every projection and MLP matrix block-circulant (CirCNN, Ding et al., arXiv:1708.08917; block 256 here, illustrative): each block is defined by one vector, so a block product is an FFT, a pointwise product and an inverse FFT. Training applies the same split per block:

• The weight gradient is now a data-by-data product per token, not per sequence, because every token's input blocks form a new conj(X). At 4,096 tokens, 20.73% of this model's optical work is data-by-data.
• The LM head stays dense and, with every other matrix circulant, is most of what remains digital; the explorer's bars (slide 10) show it.

The Training-Step Amdahl Table

Share of FLOPs that is transform work in the forward pass, the backward pass and the whole step; the static share; the data-by-data part of the optical work; and the Amdahl bounds 1/(1−share). The transformer is 0.00% throughout and omitted.

Source: results.md, section “2. A training step”. The hybrid's ratio and the circulant block size are illustrative.

• Hyena-2 at 4,096 tokens: 0.21% of a forward pass, 0.15% of a training step, 0.14% with once-per-step masks. If the optics made all of it free, a step would be at most 1.001x faster.
• Block-circulant at 4,096: 12.19% of the step (Amdahl 1.139x), but only 9.66% with static masks (1.107x), because a fifth of the optical work is the weight gradient.

Why the Share Falls from Forward to Step

Every row of the table has the same shape: backward share well below forward share, step share in between. The ledger shows why, per Hyena channel and per matmul:

• The backward pass adds no long FFTs beyond the forward's count, because U's spectrum is kept from the forward pass and dH is accumulated in the Fourier domain. For Hyena-2, where FFTs are most of the optical work and matmuls most of the total, the backward share is therefore about half the forward: 0.21% against 0.11% at 4,096 tokens.
• The step is a weighted mean, about one part forward to two parts backward (for Hyena-2 the backward pass is 1.998x the forward in FLOPs, results.json), so it lands nearer the backward share: 0.15%.
• Block-circulant falls less (14.18% to 11.16%), because there the Fourier-domain products, which do double, are a large part of the optical work.
• The per-step terms are negligible. Filter FFTs and gradient IFFTs once per step are shared by 256 sequences at 4,096 tokens, so they are a tiny correction to the per-sequence work.

Interactive: Training-Step Amdahl Explorer

The FLOP ledger of train_flop_share.py (step_per_sequence for the four variants, its op classes and the section-4 mask budget), ported to JavaScript and tested to give identical numbers (js/train_flop_model.js, checked by analysis/test_js_parity.py). The default state reproduces the results.md section 2 row “Hyena-2 + block-circulant weights, 4,096”: 14.18%, 11.16%, 12.19%, 9.66%, 20.73%, 1.139x, 1.107x.

Things to try: move the step size and watch the shares stay put while the digital step time grows and the weight-mask frames do not; then switch to the transformer and see the bars turn entirely digital. Sequences beyond the model's trained context are shape-only extrapolations, and the 64-GPU cluster and mask device are illustrative.

Masks Rewritten Every Optimizer Step

An inference engine writes its filter spectra once and keeps them. A training engine must rewrite every weight spectrum after every optimizer step, for the forward pass (H) and the backward pass (conj(H)). On the inference series' illustrative device, a 2-megapixel micromirror device at 1,031 Hz with 8-bit depth, holding one complex value per pixel (optimistic):

Source: results.md, section “4. Holding the masks”. Digital step time = step FLOPs on an illustrative 64 GPUs at 400 TFLOP/s sustained each.

• Weight masks alone rival the step. Rewriting the Hyena-2 filter spectra takes 0.52 s per step at 2,048 tokens and 2.08 s at 8,192 (longer sequences have longer filter spectra), against a 2.05 s digital step; for the block-circulant model, whose digital step is only 0.15 s, the rewrite alone is several times the step.
• Data-by-data masks are out of reach: 133 s of frames per step for Hyena-2, 385 s for block-circulant, against seconds of digital work. The weight gradient stays digital, and only the static share of slide 08 is a candidate.
• Bigger steps help the weight masks only: their frame count does not depend on how many tokens the step holds, while the digital step time grows with it. The explorer on slide 10 computes this live.

Layer-Major Order

The frame counts above assume the most favourable schedule: layer-major order, in which every sequence of the step passes through one layer while that layer's mask is loaded, before the mask is rewritten for the next layer.

Whether a cluster can run layer-major for the layers that touch the engine, and what that does to memory, parallelism and step time, is the subject of deck 03, “Optical Engines in a Training Cluster”.

What Does Not Map in Training

And the hidden costs that apply even where it maps, now per step instead of once: mask rewrites, conversions per value, and the precision a gradient format asks for (deck 02).

What to Take Away

• The ledger is right where it can be checked: it reproduces the 6N rule plus attention exactly for the transformer, and the spectral backward formulas match direct sums to rounding.
• The gradient of a convolution is a correlation. dU = IFFT(G·conj(H)) uses the weight's own spectrum, mirrored: a mask that can be written once per step. dH = IFFT(G·conj(U)) multiplies two data spectra: a mask per sequence, so it stays digital.
• Training does not rescue the transform share. The backward pass has as many long FFTs as the forward while its matmuls double, so the share falls: Hyena-2 at 4,096 tokens is 0.21% of a forward pass, 0.15% of a step and 0.14% with static masks.
• Block-circulant weights (speculative at LLM scale) raise the step share to 12.19% (Amdahl 1.139x), but a fifth of that optical work is the weight gradient, so once-per-step masks hold 9.66% (1.107x).
• The masks change every step. Even in layer-major order, rewriting the weight spectra takes from a quarter of a digital step (2,048 tokens) to a whole one (8,192) for Hyena-2 on the illustrative device, and several steps' worth for block-circulant; data-by-data masks would take minutes.
• Next: deck 02, “Optical Backpropagation and In-Situ Training”, turns from training LLMs on optics to training the optics: in-silico diffractive networks, in-situ backpropagation by the adjoint method, physics-aware training, direct feedback alignment, and the ENOB a gradient format asks for.