Fourier Optics for Training

Presentations in This Series

1. Transforms in Training Workloads →

   What the backward pass does to a Fourier-domain layer: the gradient of a convolution is a correlation, so dX reuses the filter's spectrum (conjugated) but the weight gradient multiplies two data-dependent spectra. An Amdahl analysis of training-step FLOP shares for transformer, Hyena-2, hybrid and block-circulant variants, which share a once-per-step mask could hold, and the mask-rewrite bill when weights change every step.

   liveFOptTrn 01
2. Optical Backpropagation and In-Situ Training →

   Training the optics itself: in-silico training of diffractive networks, in-situ backpropagation by the adjoint method in photonic meshes, physics-aware training with a digital twin in the loop, direct feedback alignment on a photonic co-processor, and what analogue noise and the reality gap do to convergence. The precision a gradient format (BF16, FP8 E5M2) asks of the ADC.

   liveFOptTrn 02
3. Optical Engines in a Training Cluster →

   Where an optical transform engine could sit in data-, tensor- and pipeline-parallel training: step time against mask-rewrite time, layer-major scheduling, the optimizer step and frequency-domain gradient compression (DeMo), and why the bytes saved in the link matter more than the FLOPs saved in the engine.

   liveFOptTrn 03

Analysis and Code

1. </>
   analysis/ →

   train_flop_share.py: the training-step FLOP-share (Amdahl) analysis for a Llama-3-8B-shaped model in four variants, forward and backward, split into weight-masked and data-by-data Fourier-domain work; a cross-check against the 6N rule; DeMo's transform cost; SLM frames per optimizer step; the ENOB a training format needs; NumPy checks of the backward formulas. js/train_flop_model.js is the JavaScript port behind deck 01’s explorer, checked against the Python by test_js_parity.py. results.md holds every number.

   Python · pytest
2. </>
   Inference analysis →

   The forward-pass ledger this series extends: same model shape, same variants, same FLOP conventions.

   sister series

Glossary: Concepts and Where They Are Explained

Every concept the series relies on, with a short explanation. Each entry links to the slides that explain it and, where the numbers exist, the section of analysis/results.md. Optics and inference concepts are explained in the Fourier Optics for Inference glossary and linked, not repeated.

The backward pass

The 6N rule

Training a dense model costs about 6 FLOPs per parameter per token: 2 for the forward pass and 4 for the backward (the activation gradient and the weight gradient), plus an attention term that grows with sequence length. The analysis reproduces it exactly for the transformer variant.Explained in: FOptTrn 01 · the 6N rule, reproduced · · results.md §1

The gradient of a convolution is a correlation

For y = h ∗ u, the input gradient is G correlated with h, IFFT(G·conj(H)), and the filter gradient is G correlated with u, IFFT(G·conj(U)). A backward long convolution is one more FFT and two more Fourier-domain products than the forward; the analysis checks the formulas against direct sums and finite differences.Explained in: FOptTrn 01 · the gradient of a convolution is a correlation · FOptTrn 01 · three 4f passes · · results.md §6

Weight-masked and data-by-data products

A Fourier-domain product whose second operand is a weight spectrum (forward H, backward conj(H)) can be a mask written once per optimizer step. One whose operands both depend on the sample (conj(U)·G for the weight gradient) needs a new mask for every sequence, and in practice stays digital.Explained in: FOptTrn 01 · weight-masked and data-by-data products · FOptTrn 03 · data-by-data products stay digital · · results.md §2

Static share and the training Amdahl bound

The fraction of a training step's FLOPs that are transforms plus weight-masked products: what an engine with once-per-step masks could take. 1/(1−share) bounds the step speed-up. For Hyena-2 at 4,096 tokens it is 0.14%; for block-circulant weights 9.7%.Explained in: FOptTrn 01 · the training-step Amdahl table · FOptTrn 01 · interactive Amdahl explorer · FOptTrn 03 · the static Amdahl bound · · results.md §2 · compare the inference bound

Mask rewrites per optimizer step

An inference engine writes its filter spectra once; a training engine must rewrite them after every optimizer step, for the forward and the backward pass. Processing the whole step layer by layer amortises each write over every sequence, but the frame budget still rivals the digital step time.Explained in: FOptTrn 01 · masks rewritten every step · FOptTrn 03 · masks change every step · FOptTrn 03 · interactive: step time against mask time · · results.md §4 · compare mask capacity

Training the optics

Diffractive deep neural networks (D²NN)

Stacks of passive diffractive layers whose transmission is trained in simulation and then fabricated (Lin et al., Science 2018). Training is entirely digital; the optics only runs inference, so model mismatch after fabrication is the central problem.Explained in: FOptTrn 02 · in-silico training and the reality gap

In-situ backpropagation (the adjoint method)

In a reciprocal photonic circuit, the gradient with respect to every phase shifter can be measured by sending the error signal backwards and interfering it with the stored forward field (Hughes et al., Optica 2018; demonstrated in a programmable photonic chip by Pai et al., Science 2023). The hardware computes its own gradient.Explained in: FOptTrn 02 · the adjoint method · FOptTrn 02 · in-situ backpropagation, demonstrated

Physics-aware training

Run the forward pass on the physical system and the backward pass through a differentiable digital model of it, so the update uses the real output but an approximate gradient (Wright et al., Nature 2022). It tolerates mismatch that pure in-silico training does not.Explained in: FOptTrn 02 · physics-aware training · FOptTrn 02 · interactive: three trainers

Direct feedback alignment (DFA)

Replace each layer's backpropagated error with a fixed random projection of the output error. The random projection is a large matrix multiply an optical scattering medium does for free; a photonic co-processor has trained networks this way (Launay et al., 2020). It scales poorly to transformers' accuracy.Explained in: FOptTrn 02 · direct feedback alignment

The reality gap

The difference between the model a network was trained on and the hardware it runs on: fabrication error, drift, crosstalk, temperature. In-situ and hardware-in-the-loop training close it; in-silico training does not.Explained in: FOptTrn 02 · in-silico training and the reality gap · FOptTrn 02 · interactive: three trainers

Precision

Gradient precision and FP8 E5M2

Training keeps activations and weights in BF16 or FP8 E4M3 and gradients in FP8 E5M2, which trades mantissa for range. Applying the inference series' noise rule, matching E5M2 needs about 5.3–5.8 ENOB and BF16 about 10.3–10.7, higher for spiky gradients whose crest factor is larger.Explained in: FOptTrn 02 · precision for training formats · · results.md §5 · the noise rule

Bit planes

An analogue engine with few bits of input resolution carries each integer as P binary planes, one optical pass per plane, recombined digitally with weights 2^b. The read-out noise per plane is multiplied by √Σ4^b (147.8 for 8 planes). Measured in a training loop, 4 planes cost about 0.12 in loss against 8 and 12 gained nothing.Explained in: FOptTrn 02 · precision for training formats · CiT results · ENOB and planes

Analogue noise in a training loop

Zero-mean read-out noise in the forward pass acts like injected noise, which training can tolerate or even exploit; noise in the gradient adds to minibatch noise. Biased errors (offset, nonlinearity, drift) do not average out and must be calibrated.Explained in: FOptTrn 02 · analogue noise in a training loop · FOptTrn 02 · interactive: three trainers

The training cluster

Layer-major scheduling

Passing every sequence of a step through one layer before moving to the next, so one mask write serves the whole step. It conflicts with pipeline parallelism and activation memory, which prefer sequence-major order.Explained in: FOptTrn 01 · layer-major order · FOptTrn 03 · layer-major scheduling against the pipeline

Error feedback

A compressor that keeps everything it did not send in a residual and adds it to the next step's input, so nothing is dropped, only delayed. In DeMo the delta compressed each step is the momentum plus this residual, which decays by 0.999 per step. What is subtracted must be what was actually sent: in the Compute in Transit experiment, feeding back the exact band instead of the band the emulated optics produced made the optics' errors a bias, and all six first-design runs diverged.Explained in: FOptTrn 03 · the optimizer step and DeMo · FOptTrn 02 · a measured case: optics inside DeMo’s feedback loop · CiT results · the feedback loop

Fixed frequency band

Sending the same transform coefficients (for example the lowest 1/16 of each chunk's DCT) every step instead of the largest ones (top-k): no comparisons, so a link module could compute it, and every worker's band lines up for summing. On a 15.8M-parameter GPT it ended at a validation loss of 3.124–3.158 against 1.785 for DeMo's top-k and 1.646 for AdamW.Explained in: FOptTrn 03 · a fixed band · CiT results · loss curves

Captured energy

The share of a tensor's energy (the sum of its squared transform coefficients) that a selector keeps. On the experiment's momentum with 2-D 64×64 chunks, top-k kept 0.643 and the fixed low band 0.013; a random set keeps about the keep fraction. It explains why the fixed band trains so much worse.Explained in: FOptTrn 03 · a fixed band · CiT results · captured energy

All-gather against all-reduce wire bytes

A ring all-reduce of B bytes sends 2(N−1)/N·B per worker whatever N is; a ring all-gather of P bytes per worker sends (N−1)·P, which grows with N. On 8 workers DeMo's 8× and 16× smaller payloads became only 2× and 4× fewer wire bytes. Summing a fixed band in the transform domain (it is linear) restores the all-reduce pattern.Explained in: FOptTrn 03 · the optimizer step and DeMo · FOptTrn 03 · where the band is computed · CiT results · bytes

Time and energy to a target loss

The measured steps to reach a loss times the step time (compute plus, for each sync round, latency and bytes over bandwidth), and the same steps times the energy per step. It weighs a compressor's saved bytes against the extra steps it needs: to 3.158 on 2 DCs × 4 over 100 GbE, AdamW takes 1.63 s and the fixed band 3.85 s.Explained in: FOptTrn 03 · training across data centres · CiT results · transit explorer

Cross-data-centre training and break-even bandwidth

Data-parallel workers split across data centres sync over a WAN far slower than the fabric inside one. The break-even bandwidth is the WAN rate below which one way of syncing reaches the target loss sooner than another, found by bisection: a band in a pluggable link module never breaks even with AdamW from 0.1 to 1,000 GB/s; summing the band at the switch wins below 1.57 GB/s.Explained in: FOptTrn 03 · where the band is computed · FOptTrn 03 · training across data centres · CiT results · break-even

DiLoCo-style sync

Workers train locally and sync only every H steps, dividing communication by H; DiLoCo (Douillard et al., 2023) also changes the optimizer. The Compute in Transit simulation's DiLoCo-style rows assume the steps to a loss are unchanged: an extrapolation, not a measurement. At H = 100 the bytes a band saves matter little.Explained in: FOptTrn 03 · training across data centres · CiT results · extrapolations

DeMo (decoupled momentum)

An optimizer that DCT-transforms the momentum in small chunks and sends only the largest coefficients between workers (Peng et al., 2024). Its transforms are about 5×10⁻⁶ of a step's FLOPs: what it saves is link bytes, not compute. Measured on a 15.8M-parameter GPT with 8 workers (Compute in Transit), top-k DeMo ended at 1.785 ± 0.007 against AdamW's 1.646 ± 0.005, with 8× less payload and 2× fewer wire bytes.Explained in: FOptTrn 03 · the optimizer step and DeMo · · results.md §3 · Compute in Transit

Explained in the sister series

The lens as a Fourier transformer · The convolution theorem · Spatial light modulators · The phase problem · MZI meshes · The noise rule for float workloads · Conversion-bound optics · Long convolution by FFT · Hyena · Block-circulant weights · Hybrid architectures · Optical share and the Amdahl bound · Fourier-plane mask capacity · Amdahl's law for a port.