Skip to content
Reference implementation is running

Your model stops learning long before your GPU stops paying for it.

EnAi watches gradients while a training run is in flight and retires layers the moment they converge. Backward compute drops mid-run — with no restart, no rewrite, and no change to your model code.

ResNet-18 · CIFAR-10 · Apple M3 · PyTorch 2.13 · Metal Performance Shaders

forward · all layersbackward · trainable layers onlyfrozen
0.00%

Backward compute removed

Confirmed by autograd, not estimated

+00.0%

Throughput across the full run

462 → 545 img/s on active epochs

0

Backward passes into frozen layers

Down from 1,563 — measured per layer

−0.0%

Wall-clock training time

354.1s → 320.1s, same images seen

Measured on a single controlled pair of runs. Full method, and what these numbers do not establish, in the method section.

The problem

Training bills you for work that finished hours ago.

Long after the early layers have stopped changing in any meaningful way, every training step still computes their weight gradients, still propagates activations back through them, and still pays the memory and bandwidth to do it.

01

Convergence is uneven

The first convolutional layers settle into edge and texture detectors in a fraction of the time the deeper layers need. A standard training loop treats every layer as equally unfinished, from the first step to the last.

02

Backward is the expensive half

A layer costs F going forward and up to 2F coming back — once for its weight gradient, once for the gradient it passes upstream. Freeze a layer at the head of the graph and both terms disappear, along with the gradient for the layer just behind it.

03

Static configs cannot react

Batch size, learning rate and the set of trainable parameters are chosen before step one and never revisited. A run that could safely speed up at the halfway mark has no mechanism to notice, let alone act.

The waste is not in the model and not in the hardware. It is in the fact that nothing is watching the run while it happens and nothing is allowed to change while it does.

How it works

A control loop that runs inside training, not around it.

EnAi attaches to the layers it is allowed to touch, measures what they are still contributing, and acts at epoch boundaries. Everything it does is reversible and observable.

1

Observe

Forward and full-backward hooks sample each watched layer every N steps: activation RMS, gradient RMS at the output, gradient RMS at the weights, and that layer’s share of the network’s total gradient energy. On the steps in between, both hooks return on their first line — the steady-state cost is a branch, not a synchronization.

register_forward_hook · register_full_backward_hook

2

Decide

A governor reads that telemetry at each epoch boundary. When a layer’s share of gradient energy collapses, it has stopped contributing to learning and becomes a candidate. In the reference run the two watched layers fell from 35.7% of total gradient energy to effectively nothing.

early_grad_energy_share: 0.357 → 0.000

3

Adapt

Freeze the layer and its BatchNorm, raise the batch size into the headroom the freeze just released, and rescale the learning rate linearly so the optimization trajectory stays comparable. The run continues from where it was — no checkpoint, no restart, no change to the model definition.

batch 32 → 64 · lr ×2 · 4 modules frozen

What the governor actually did

⚡ ENGINE ENGAGED — start of epoch 2
   trigger  : early-layer gradient share 0.357 → below threshold
   lever 1  : froze conv1, bn1, layer1.0.conv1, layer1.0.bn1
              38,848 params non-trainable · BatchNorms pinned to eval
   lever 2  : batch size 32 → 64
   lever 3  : base LR ×2 → 0.05000  (linear with batch size)
   compute  : 1.6645 → 1.5495 GMAC/image/step  (6.91% cut)

✅ VERIFIED — after the freeze, autograd never entered the watched layers.

The learning-rate rescale is not cosmetic. Without it the optimized run trains at an effectively halved rate and the accuracy comparison stops measuring the thing it claims to measure. Disable it with --no-lr-rescale and the confound is visible immediately.

Why this is harder than it sounds

Three ways to build this that look like they work.

Each of these produces a system that reports a saving while quietly doing something else. All three were found by measurement rather than by reasoning, and each one now has a test that fails if it comes back.

Freezing BatchNorm does not freeze BatchNorm

Setting requires_grad = False stops the weight and bias updates. It does nothing about the running mean and variance, which are updated as a side effect of the forward pass. And model.train() silently puts the module back into training mode at the start of every epoch.

How it is handled

Frozen BatchNorms are held in eval mode and re-asserted after every model.train() call. A test asserts they revert without that call, so the fix cannot be quietly deleted.

Stale gradients keep the optimizer moving

SGD skips a parameter only when its gradient is None. A gradient tensor left over from the previous step keeps updating a parameter that has supposedly been frozen — silently, and with no error anywhere.

How it is handled

Freezing explicitly clears the gradient to None. A test compares the actual weight tensors before and after real optimizer steps rather than trusting the flag.

Changing batch size can cost more than it saves

Rebuilding the DataLoader to change batch size tears down and respawns the worker pool. On macOS those workers are spawned, not forked, so each one re-imports the framework. Measured: about 20 seconds of dead time on a run whose baseline was 32 seconds.

How it is handled

The batch size is changed by mutating a sampler in place. The loader and its workers are never touched, so the transition costs milliseconds instead of restarting the input pipeline.

Results

One controlled pair of runs, reported in full.

Same model, same data, same seed, same optimizer. The two arms run identical code until the governor engages at the start of epoch 2 — which is why epoch 1 comes out the same to the decimal in both.

model
ResNet-18 · CIFAR stem — 3×3 stride-1 conv, no maxpool
data
CIFAR-10 · 50,000 images
epochs
3
hardware
Apple M3
runs per arm
1

Throughput per epoch

BaselineEnAi

Images per second. The engine engages at the start of epoch 2.

0150300450600429434epoch 1batch 32462544epoch 2batch 32 → 64463547epoch 3batch 32 → 64

Epoch 1 is near-identical by design — both arms run the same code until the governor engages. Epochs 2 and 3 average +18%.

Where a training step spends compute

GMAC / image / step

A layer costs F going forward and up to 2F coming back — once for its weight gradient, once for its input gradient.

Baseline1.6645
0.55540.55540.5537
EnAi1.5495
0.55540.51590.4782

Forwardunchanged, because frozen layers still produce activations

Weight gradientsremoved for every frozen layer

Input gradientsremoved for the frozen prefix and the first trainable layer too

Compute ledger in GMAC per image per step
ArmForwardWeight gradientsInput gradientsTotal
Baseline0.5554230.5554230.5536531.664499
EnAi0.5554230.5159050.4781561.549483
Verified by autograd

The frozen layers receive exactly zero backward passes.

Backward hooks stay attached to the frozen layers for the entire run. After the freeze they stop firing completely, while the forward hooks keep firing — the layers still compute activations, they just cost nothing to train. That silence is PyTorch reporting that the backward sub-graph was pruned. It is not inferred from a stopwatch.

Before freeze

1,563

backward passes, per layer

After freeze

0

forward passes continue: 1,644

conv1bn1layer1.0.conv1layer1.0.bn1frozen after epoch 1

Validation accuracy

BaselineEnAi

Shaded band is ±2 standard errors of measurement noise on the difference.

50%60%70%80%79.73%78.68%identical — 59.39%epoch 1epoch 2epoch 3

Final gap is -1.05pp — about 1.8 standard errors, so it is not statistically significant, but it is not comfortably inside the noise either. The optimized run also took 33% fewer optimizer steps, which is a plausible mechanism for a small real regression over only three epochs. We call that unresolved, not parity.

Parameters are the wrong unit

Freezing those four modules made 38,848 parameters non-trainable — 0.35% of the model. It removed 6.91% of the compute.

Early convolutional layers are tiny in parameters and expensive in operations, because they run at full spatial resolution. Anyone quoting parameter count here would understate the result by roughly twenty times.

Wall-clock354.1s → 320.1s
Optimizer steps4,689 → 3,127
Images seen150,000 — identical
Method

What we measured, and what we didn’t.

An efficiency claim is only worth what its methodology survives. Here is the part most pages leave out.

Established
  • 6.91% of training compute removed

    Computed from the layers’ real output shapes captured by temporary hooks, not hand-derived, so the ledger stays correct if the architecture changes.

  • Zero backward passes into the frozen layers

    Counted per layer by hooks that stay attached for the whole run: 1,563 before the freeze, 0 after, while forward hooks keep firing.

  • +10.6% throughput, −9.6% wall-clock

    Both arms processed exactly 150,000 images. Timings exclude one-off kernel recompilation but still charge it to the total.

  • Frozen weights and BatchNorm buffers do not move

    Tensors are compared before and after real optimizer steps, rather than trusting the requires_grad flag.

Not established
  • Energy reduction

    Our meter was CodeCarbon, which on Apple Silicon without root falls back to a constant-TDP estimate. It attributed 0 W to the GPU — the device doing nearly all the work — and reported near-identical CPU power in both runs. The resulting energy delta is the time delta in different units, so we do not present it as an energy measurement.

  • Accuracy parity

    Final accuracy moved −1.05pp, roughly 1.8 standard errors. Not statistically significant, but not comfortably inside the noise either. The optimized run also took 33% fewer optimizer steps, which is a plausible mechanism for a small real regression over three epochs.

  • Generalization

    One model, one dataset, one accelerator, three epochs, one run per arm. Nothing here establishes behavior at larger scale, on other architectures, or across multiple devices.

What closes the gaps

  • Direct power instrumentation with root access, so the energy figure is measured rather than inferred
  • Repeats with alternating order, reporting medians instead of single runs
  • Longer schedules, where fewer optimizer steps stop being a handicap
  • Larger models and multi-accelerator training, where the backward pass dominates by more
Who this is for

Built for people who will check the numbers.

Teams training on rented compute

Throughput is the bill. A run that finishes in 90% of the wall-clock time on the same instance costs 10% less, and the change is a wrapper around the training loop rather than a rewrite.

−9.6% wall-clock, same images seen

Researchers running many short experiments

Sweeps and ablations spend most of their compute on early epochs, which is exactly where the early layers converge and stop earning their gradients. Shorter iterations mean more experiments per day on the same hardware.

+18% throughput once engaged

Anyone reporting efficiency numbers

The instrumentation is the product as much as the optimization is. Hook-level counters, an analytic compute ledger and a per-phase power sampler produce numbers that survive someone checking them.

Per-layer counters, not estimates

Get early access

We're opening the engine to a small group of teams running real training workloads. Tell us where to reach you.

No newsletter. We'll only email you about access.