Compute in Transit 04 — The Fixed-Band Experiment

00

Topics We'll Cover

The Question
The Proposed System, End to End
The Setup
Bytes per Step: Payload and Wire (Interactive)
Loss: a Fixed Band Does Not Replace Top-k
Loss per Byte: Steps to Reach a Loss
Why: the Band Misses the Momentum's Energy
Emulating the Optics
The First Optics Runs Diverged
The Fix: Feed Back What Was Sent
ENOB and Bit Planes
What It Means for a Link Module
Limits, Next Steps, and What to Take Away

Where this deck sits

Deck 01 sets the budget and the five tests for a photonic link module (arXiv:2608.21536): a QSFP56 pluggable whose published operations are FFT, NTT, MAC, and modular and floating-point arithmetic, with no comparisons, rounding or top-k. Deck 02 found the strongest candidate in training, DeMo’s compressed momentum, and planned this experiment. This deck reports it. Code, logs and every number: DeMo_Fixed_Band_Experiment.

Concepts used here, and where they are explained. Each is defined, with links, in the series glossary.

DeMo Fixed band in place of top-k Error feedback Residual Captured energy Payload and wire bytes Ring all-reduce Bit planes ENOB Feedback of the values sent

01

The Question

DeMo (Peng et al., arXiv:2411.19870, ICLR 2026) lets each worker keep its own momentum. Every step it takes a chunked DCT of that momentum, sends the top-k coefficients of each chunk, and subtracts what it sent; the workers all-gather the values, and each applies the sign of the averaged update.

Top-k, as published

Keeps the largest coefficients of each chunk: the most energy for the bytes.
Needs comparisons to find them, and sends their positions as well as their values.
Not in the module’s published operation set.

A fixed band

Keeps the same positions in every chunk, every step: a gather, no comparisons.
Sends values only, so half the bytes of top-k at the same k.
Fits the module, but keeps the energy only if the momentum’s energy sits in that band.

Two questions

1. Does DeMo still converge, per byte sent, when top-k is replaced by a fixed band? 2. If the band is computed by an analogue optical transform, with its quantisation, clipping and readout noise, how much precision does it need?

01b

The Proposed System, End to End

Variant E, from a worker’s gradient on the TX GPU to the weight update on every receiving GPU. The delta leaves in 1-D chunks of 256; the module writes each chunk onto light as bit planes, the optics transforms it, and only the band’s 16 outputs are read and sent: values, no indices. The receivers place them at the fixed positions, average, invert the DCT and apply the sign, and each sender feeds back what it actually sent. Switch to E′ to see the band computed and summed at the switch instead (deck 05’s hypothetical variant).

Tags: each block says where it comes from: P the module’s paper (arXiv:2608.21536), X this experiment’s model, S the simulation (deck 05), A an assumption, such as reading only the band’s outputs.
Every frame is computed by a port of the experiment’s code (the delta update, the orthonormal DCT, the input quantisation and two’s-complement bit planes, the readout at an ENOB, the recombination, the scatter-mean, the inverse DCT and the sign), checked against the experiment’s own functions at b1f62d3. The code lives in the results site’s repository.
The chunk’s numbers are illustrative (a smooth shared signal plus each worker’s noise, and an out-of-band residual, from a seeded generator). The byte counts are the experiment’s recorded ones, for the whole model per worker per step (slide 03).

02

The Setup

A GPT of 15,833,472 parameters (8 layers, width 384, context 256), trained on TinyStories with a 4,096-token byte-level BPE. Eight simulated workers in one process, each with its own data shard, momentum and DeMo state; 8 sequences per worker per step; 3,000 steps. DeMo’s momentum decay is 0.999, and every compressed variant keeps 1/16 of the coefficients.

	Variant	What is sent	Seeds
A	AdamW on all-reduced gradients (the baseline)	dense BF16 gradients	3
B	DeMo as published: chunked DCT, top-k per 2-D 64×64 chunk	values and positions	3
C	DeMo with a fixed low-frequency band (2-D 64×64, 1-D 256 or 1-D 64 chunks)	values only	3 (1-D 64: 1)
D	DeMo with a calibrated band: chosen once from a short window’s mean energy map, then frozen	values only	3
E	C (1-D 256) with the band computed by an emulated optical transform	values only	1 per setting
F	Other fixed shapes of the same size: zig-zag, high band, random positions	values only	1 each

1-D chunks matter because the module does 1×N power-of-2 transforms; DeMo’s own chunks are 2-D tiles.
Diagnostics every 200 steps, on worker 0: the share of energy that top-k and each band shape would keep at the same k, for a plain momentum, for the delta DeMo actually compresses, and for the gradient.
Runs on Colab GPUs (mostly A100; some T4 and L4). GPU runs are close but not bit-identical.

Source: results/tinystories-gpu/results.md (header and section 1) and docs/optics_investigation.md. DeMo is implemented from the paper and checked against the reference code in a test.

03

Bytes per Step: Payload and Wire (Interactive)

Two ways to count. Payload is what a worker originates each step. Wire is what it actually sends in its ring collective: an all-reduce for AdamW, an all-gather for DeMo, because each worker’s top-k positions differ.

Variant	Payload / step (bytes)	vs dense	Wire / step (bytes)	Wire vs dense
A AdamW, all-reduce	31,666,944	1.0×	55,417,152	1.0×
B DeMo, top-k	3,957,960	8.0×	27,705,720	2.0×
C–F any fixed band	1,979,184	16.0×	13,854,288	4.0×

Source: results.md, section 1. Values in BF16; top-k positions at the bytes needed for a position within a chunk (with int64 positions, as the reference code sends them, top-k is 9,895,920 bytes, 2.5× the compact count). Eight workers.

On a ring of N workers an all-reduce sends 2(N−1)/N times its buffer, close to twice whatever N is; an all-gather sends (N−1) times the payload. So with eight workers the payload ratios of 8× and 16× become 2.0× and 4.0× on the wire, and they shrink further as workers are added.

Workers on the ring8

04

Loss: a Fixed Band Does Not Replace Top-k

Final validation loss after 3,000 steps. Top-k DeMo ends 0.139 above AdamW. Every fixed band ends more than a nat above top-k.

Variant	Seeds	Final loss	Above top-k
A AdamW, dense all-reduce	3	1.646 ± 0.005	−0.139
B DeMo, top-k, 2-D 64×64	3	1.785 ± 0.007	–
C fixed low band, 1-D 256	3	3.124 ± 0.008	1.339
C fixed low band, 1-D 64	1	3.137	1.352
C fixed low band, 2-D 64×64	3	3.158 ± 0.009	1.373
D calibrated band, 2-D 64×64	3	3.087 ± 0.027	1.302
F zig-zag / high / random, 2-D 64×64	1 each	3.028 / 3.244 / 2.700	1.243 / 1.459 / 0.915

Source: results.md, section 1. ± is the sample standard deviation over seeds.

The low-frequency bands (low, zig-zag, calibrated) end 1.24 to 1.37 above top-k; the high band, 1.46.
A random fixed set of positions beats every structured band (2.700, one seed): the structure of a low band does not help.
The chunk shape barely matters for the low band: 3.124 (1-D 256), 3.137 (1-D 64), 3.158 (2-D).

05

Loss per Byte: Steps to Reach a Loss

Half the bytes per step would pay if the band reached a given loss in less than twice the steps. It does not come close. Steps to the first evaluation (every 200 steps) at or below each target, mean over seeds, shown only when every seed gets there; payload in brackets.

Variant	to 3.158 (C, 2-D)	to 3.124 (C, 1-D 256)	to 3.087 (D)	to 1.785 (B)
A AdamW	400 (12.67 GB)	400 (12.67 GB)	400 (12.67 GB)	2,000 (63.33 GB)
B top-k	400 (1.58 GB)	400 (1.58 GB)	400 (1.58 GB)	–
C low band, 1-D 256	2,667 (5.28 GB)	–	–	–
D calibrated band	2,533 (5.01 GB)	2,667 (5.28 GB)	–	–
F random fixed set	1,200 (2.38 GB)	1,200 (2.38 GB)	1,400 (2.77 GB)	–
F zig-zag band	2,200 (4.35 GB)	2,400 (4.75 GB)	2,600 (5.15 GB)	–

Source: results.md, section 1, Steps to reach a loss. Payload = steps × bytes each worker originates per step. A dash: not every seed reached the target. The 2-D low band (3 seeds) does not reach even its own mean on every seed.

Per byte, top-k still wins

Top-k reaches the 1-D band’s final loss (3.124) by step 400, after 1.58 GB per worker. The band itself needs 2,667 steps and 5.28 GB just to reach 3.158, and not every seed reaches 3.124. On the wire every DeMo figure is seven times its payload and AdamW’s 1.75 times (slide 03), so the comparison between DeMo variants is unchanged, and top-k’s lead over AdamW shrinks from 8× to 2× per step.

06

Why: the Band Misses the Momentum's Energy

A fixed band can only work if the energy of what DeMo compresses sits in it. The diagnostics measure the share of energy each selector would keep at the same k, on worker 0, averaged over every diagnostic after step 200, all tensors together. A random set keeps about the keep fraction (1/16, about 0.06).

Run, chunks	Tensor	Top-k	Low band	Random
C, 2-D 64×64	momentum	0.643	0.013	0.065
C, 2-D 64×64	gradient	0.619	0.305	0.059
C, 2-D 64×64	delta	0.649	0.000	0.065
C, 1-D 256	momentum	0.429	0.146	0.049
C, 1-D 256	gradient	0.485	0.228	0.053
C, 1-D 256	delta	0.417	0.000	0.048
D, 2-D 64×64	momentum	0.578	0.109 (calibrated band 0.190)	0.063

Source: results.md, section 2. Momentum = a plain momentum of the same decay, never compressed; delta = what DeMo actually compresses (momentum plus the residual error feedback keeps); gradient = the step’s gradient.

In DeMo’s 2-D chunks the low band keeps 0.013 of the momentum’s energy: about a fifth of a random set (0.065), against 0.643 for top-k at the same k.
The gradient’s low band is much fuller (0.305 in 2-D). But DeMo compresses a momentum (decay 0.999), not a gradient, and in 2-D chunks the momentum’s low band is nearly empty.
1-D chunks help the band, but not enough: 0.146 of the momentum, about a third of top-k’s 0.429.
The delta holds nothing in the band: error feedback removes the band every step, so what remains to compress is all out of band. That residual becomes the optics’ problem (slide 08).

07

Emulating the Optics

Variant E computes the 1-D 256 band as an analogue optical transform would, so that its errors enter training. The emulator follows the module’s constraints: power-of-2 transform sizes, binary (NRZ) input, and only multiply-accumulates for anything data-dependent, so no per-chunk maximum.

delta (BF16) → scale to full scale, saturate → 8-bit two’s complement → one optical pass per bit plane, noisy readout → recombine with weights 2^b

Input full scale: kappa = 4 times the delta’s RMS, measured over a calibration window (steps 200 to 299): a sum of squares, no comparisons. Inputs beyond it saturate, and the clip rate is logged.
Bit planes: each binary digit of the input passes through the transform separately; the planes are recombined with weights 2^b, the top one negative. The number of planes is the number of optical passes (4, 8 or 12 tested).
Readout noise per plane: RMS = FS × 2^−ENOB × 2/√12, with FS = 16, the largest row L1 norm of a 256-point orthonormal DCT. ENOB 6, 8, 10 and 12.
The emulator is accurate on white input (where the band holds 1/16 of the energy): 8 planes without readout noise reproduce the exact band to 0.9% RMS error, 12 planes to 0.2%, and 8 planes at ENOB 8 to 16%.
A perfect-optics control (no quantisation, no noise) is C with E’s power-of-2 chunking: at width 384 the 17 LayerNorm vectors are cut into runs of 128 rather than DeMo’s 192.

Source: docs/optics_investigation.md, sections 1, 2 and 5; results.md, section 3. The noise model is the series’ own, not a measured characteristic of the module.

08

The First Optics Runs Diverged

As first designed, error feedback removed the exact band from each sender’s delta, while every receiver applied the optics’ version. All six runs diverged; the same band computed exactly trained normally.

Run (1-D 256)	Loss at 400	Loss at 3,000	Clip rate 300–400	Clip rate 400–600
C, exact band (seed 0)	4.07	3.12	–	–
E, ENOB 6, 8 planes	4.56	11.15	0.8%	38.6%
E, ENOB 8, 8 planes	4.59	214.65	1.0%	41.5%
E, ENOB 10, 8 planes	4.47	42.18	0.8%	34.9%
E, ENOB 12, 8 planes	4.51	336.98	0.8%	42.9%
E, ENOB 8, 4 planes	5.20	344.39	0.9%	50.8%
E, ENOB 8, 12 planes	4.49	558.83	0.8%	37.0%

Not analogue noise: ENOB 12 ended worse than ENOB 6, and 12 planes worse than 8. The cause is the residual. Error feedback keeps everything not sent, decaying only by 0.999 per step; with a fixed band, the band is removed every step and the out-of-band residual piles up.

Step	200	400	1,000	2,000	3,000
Band share of the delta’s energy	1.4–2.3e-3	2.3–2.9e-4	1.1–1.5e-4	3.9–4.2e-5	1.6–1.8e-5

The full scale goes stale: set at step 300 by a residual that keeps growing, it saturates 30 to 50% of inputs.
The band sinks below the error: per band coefficient the RMS is about √(16 × share) of the input RMS, 0.016 by the end. Quantisation alone adds 0.009; readout noise at ENOB 8 adds 0.17 after recombining 8 planes (0.011 at ENOB 12).
The error is a bias, not noise: the residual changes slowly, so quantisation and clipping errors repeat almost identically each step, are never corrected, and DeMo’s sign update turns a consistent error into a full step the wrong way, every step.

Source: docs/optics_investigation.md, sections 1 and 2 (diagnostics of the three C 1-D 256 seeds); final losses in results.md, section 3.

09

The Fix: Feed Back What Was Sent

Each idea was first tried on a small model on the CPU (TinyShakespeare, 4 layers, width 128, 600 steps, 1-D 256 chunks; about 12 minutes a run) before spending GPU time. Final validation loss at step 600, with the clip rate over the last interval:

Configuration (8 planes, ENOB 8 unless stated)	Final loss
C, exact band (reference)	2.464
E as first designed: full scale fixed, feedback of the exact band	8.809 (41.8%)
Full scale re-set every 50 steps, feedback of the exact band	3.404 (0.6%)
… at ENOB 12	3.069 (0.7%)
… with no readout noise at all	2.967 (0.7%)
Full scale re-set every 50 steps, feedback of the band sent	2.535 (0.5%)
Full scale fixed, feedback of the band sent	2.520 (5.7%)

The failure reproduces at small scale.
Re-setting the full scale removes most clipping, yet the loss still drifts: saturation was not the whole problem.
Even with no readout noise the loss stays 0.5 above the exact band: quantisation and clipping errors, never fed back, are enough.
Feeding back what was actually sent fixes it, with or without re-setting the full scale.

Why the fix is physically allowed

DeMo exchanges its compressed values with an all-gather: every worker receives every worker’s values, its own included. So a sender sees what the optics produced for its own data, noise and all, and can subtract that from its delta instead of the exact band. Nothing extra crosses the link. The optional full-scale re-set uses the delta’s mean square over an interval: multiply-accumulates only.

Source: docs/optics_investigation.md, sections 3 and 4; results/shakespeare-optics-check/results.md.

10

ENOB and Bit Planes

The optics runs were repeated on TinyStories with feedback of the band sent and the full scale re-set every 100 steps, plus one with the full scale fixed, and the perfect-optics control. One seed each.

Configuration (1-D 256)	Final loss	Clip rate (whole run)
C, exact band (3 seeds)	3.124 ± 0.008	–
E, perfect optics (control)	3.141	0.00%
E, ENOB 6 / 8 / 10 / 12, 8 planes	3.423 / 3.017 / 2.878 / 2.905	0.28% / 0.29% / 0.30% / 0.30%
E, ENOB 8, 4 / 12 planes	3.137 / 3.024	0.22% / 0.30%
E, ENOB 8, 8 planes, full scale fixed	3.021	1.34%

Source: results.md, sections 1 and 3; docs/optics_investigation.md, section 6.

Every repaired run trains, with clipping near 0.3%. Re-setting the full scale adds nothing once errors are fed back: 3.017 with it, 3.021 without.
ENOB 8 is enough for the band itself. ENOB 6 costs 0.3 against the exact band (3.423). Four planes cost 0.12 against eight; twelve gain nothing.
With 8 or more planes, from ENOB 8 up, the optics band ends below the exact band, by up to 0.25 (ENOB 10). The control (3.141, 0.017 above C’s mean) rules out the power-of-2 chunking as the cause; it ran on a T4, the others on an A100.

A hypothesis, not a result

Why would an inexact band beat the exact one? Possibly because quantisation and clipping are non-linear: the values the optics produces carry a little of the large out-of-band residual into the 16 positions sent, and with feedback of what was sent that part is then removed from the delta, so some out-of-band information reaches the model. It would fit the random band beating the low band. But every optics run here is a single seed, and on TinyShakespeare the repaired optics ended slightly above the exact band (2.520 to 2.535 against 2.464). Testing it needs several seeds and a measurement of the out-of-band energy in the values sent.

11

What It Means for a Link Module

The experiment answers deck 02’s question for this model and scale. Against the five tests of deck 01:

Test	Fixed band computed in transit
Operations	pass A transform and a gather; full-scale tracking by multiply-accumulate. Top-k would need comparisons.
Precision	pass ENOB 8 and 8 planes match or beat the exact band, once the values sent are fed back.
Bottleneck (convergence per byte)	fail Every fixed band ends more than a nat above top-k; top-k reaches the band’s final loss in 400 steps.
Placement	tension See below.

What a fixed band could gain

Every worker sends the same positions, and DeMo averages before taking the sign, so a fixed band could be summed in the network rather than gathered: an all-reduce of 1,979,184 bytes sends 3.46 MB per worker with eight workers, 16× under dense on the wire too (derived from results.md’s payload with the ring formula; not a run). That is the strongest case for computing in transit.

What the fix depends on

Feedback of the values sent relies on the all-gather returning each worker’s own optical output. An all-reduce returns only the sum. Keeping the fix would need the module to hand each sender back its own result (a loopback), or the band to be computed at the sender. And summing streams needs them to meet: a switch, not a pluggable.

Verdict

A photonic link module (arXiv:2608.21536) cannot do DeMo’s top-k, and a fixed band computed in transit is not a substitute for it at this scale: the optics is precise enough, but the band is the wrong selection. The candidate deck 02 ranked first survives only if a selection with no comparisons can be found that keeps the momentum’s energy.

12

Limits, Next Steps, and What to Take Away

Limits

One small model (15.8M parameters), one data set, 3,000 steps, eight simulated workers.
One seed for every E and F run; the optics-beats-exact finding rests on single seeds.
The learning rate was tuned on top-k and AdamW only, and reused for every fixed-band variant.
The optics is emulated with the series’ own noise model, not measured on the module.
Wire bytes assume ring collectives; no latency or link model yet.

A discrete-event model of training transit across data centres, fed by these bytes per sync (wire, not payload) and the steps-to-loss ladder.
More seeds for E, and a measurement of out-of-band energy in the values sent, to test the hypothesis of slide 10.
Selections without comparisons that track the momentum: a band re-calibrated periodically, or fixed random sets that change by a known schedule.

A fixed band does not replace top-k. Top-k ends at 1.785, AdamW at 1.646, every low-frequency band 1.24 to 1.37 above top-k, a random fixed set 0.92 above.
Half the bytes does not pay: top-k reaches 3.124 in 400 steps (1.58 GB); the band needs 2,667 steps and 5.28 GB to reach 3.158.
Why: in 2-D chunks the low band keeps 0.013 of the momentum’s energy, a fifth of a random set; top-k keeps 0.643.
Count wire bytes: with eight workers DeMo’s all-gather turns 8× and 16× payload savings into 2.0× and 4.0×.
The optics needs its errors fed back. Feeding back the exact band diverged in all six runs; feeding back the values sent, which an all-gather returns, trains, and ENOB 8 suffices.
The optics improvement over the exact band is a hypothesis, from one seed per run.