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
Training compute removed
Analytic ledger, confirmed by autograd
Throughput once engaged
Machine-normalized, median of 12 runs
Backward passes into frozen layers
Down from 1,563 — counted per layer
Of the freeze's theoretical ceiling
Predicted +7.42% from geometry, measured +6.85%
From a twelve-run controlled study — four arms, three seeds each, with ablations that could have shown the freeze contributing nothing. Full method, and what these numbers do not establish, in the method section.
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.
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.
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.
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.
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.
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
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
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.
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.
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.
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.
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.
Twelve runs, four arms, reported in full.
Same model, same data, same seeds, same optimizer. Every arm runs identical code until the governor engages at the start of epoch 2 — which is why epoch 1 comes out identical to four decimal places across all four arms, and why it can be used as a control for the machine's own drift.
- model
- ResNet-18 · CIFAR stem — 3×3 stride-1 conv, no maxpool
- data
- CIFAR-10 · 50,000 images
- epochs
- 3
- hardware
- Apple M3
- arms
- 4 — baseline, freeze, batch, both
- runs
- 12 — 3 seeds per arm
Throughput per epoch
Images per second. The engine engages at the start of epoch 2.
Epoch 1 is identical by design — every arm runs the same code until the governor engages, which is what makes it usable as a control for the machine’s own drift. Across the twelve-run study, normalizing each run against its own epoch 1 puts the engaged epochs at +20.85%.
Where a training step spends compute
GMAC / image / stepA layer costs F going forward and up to 2F coming back — once for its weight gradient, once for its input gradient.
Forward — unchanged, because frozen layers still produce activations
Weight gradients — removed for every frozen layer
Input gradients — removed for the frozen prefix and the first trainable layer too
| Arm | Forward | Weight gradients | Input gradients | Total |
|---|---|---|---|---|
| Baseline | 0.555423 | 0.555423 | 0.553653 | 1.664499 |
| EnAi | 0.555423 | 0.515905 | 0.478156 | 1.549483 |
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
Validation accuracy
Shaded band is ±2 standard errors of measurement noise on the difference.
This pair closed -1.05pp apart. Repeating it across three seeds and pairing by seed — the correct test, since both arms share initial weights and data order — the gap shrinks to -0.33pp (t = -2, df = 2; significance at 5% would need |t| > 4.303). Not resolvable at this sample size — though all three seeds moved the same way, which we read as a small real cost rather than 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.
What we measured, and what we didn’t.
An efficiency claim is only worth what its methodology survives. Twelve runs across four arms, three seeds each — including the ablations that could have shown the novel part contributing nothing. Here is the part most pages leave out.
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.
The freeze reaches 92% of its theoretical ceiling
Removing that arithmetic caps the possible speed-up at +7.42%, a number derived from layer geometry with no reference to any clock. Measured, with machine state divided out: +6.85%. A result above the ceiling could not have been the freeze; this one lands just under it.
The freeze is not just the batch size
Ablation arms run each lever alone across three seeds. Freezing on its own is worth +6.9% and raising the batch alone +9.3%, so neither accounts for the other. Together they reach +20.9% — more than independent levers would predict, meaning they interact rather than stack.
Frozen weights and BatchNorm buffers do not move
Tensors are compared before and after real optimizer steps, rather than trusting the requires_grad flag.
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
Across three seeds, final accuracy sits −0.33pp against the baseline (t = −2.0, df = 2; |t| would need to exceed 4.30 to be significant at 5%). So the gap is not resolvable at this sample size — but all three seeds moved the same direction, and that consistency is weak evidence of a small real cost rather than proof of parity. Worth noting: an earlier single pair of runs showed −1.05pp, which did not survive replication.
Wall-clock reduction as a portable number
On this hardware the same baseline configuration varies by 20% between seeds from thermal throttling alone. Every timing figure here is therefore normalized against a within-run control rather than compared as raw seconds, and an absolute “X% faster” claim would not transfer to a machine with different cooling.
Generalization
One model, one dataset, one accelerator, three epochs. Twelve runs make the comparison internally sound; they say nothing about 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
- A cooled, interleaved re-run so more of the twelve survive the machine-state filter
- Longer schedules, where fewer optimizer steps stop being a handicap
- Larger models and multi-accelerator training, where the backward pass dominates by more
Built for people who will check the numbers.
Teams training on rented compute
Throughput is the bill. Fewer hours on the same instance is less money, and the change is a wrapper around the training loop rather than a rewrite.
+20.9% throughput once engaged
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.
Ablated across 4 arms, 3 seeds
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
Questions people actually ask
Including the one about whether the energy saving is real.
What is EnAi?
EnAi is a runtime optimizer for model training. It attaches PyTorch hooks to a running training job, measures how much each layer still contributes to learning, and freezes layers once their gradients collapse — removing backward computation mid-run without restarting the job or changing the model definition.
How does freezing layers reduce training cost?
A layer costs F operations going forward and up to 2F coming back: once to compute its weight gradient, once to compute the gradient it passes upstream. Freezing a layer at the head of the network removes both terms, and also removes the input-gradient term for the layer immediately after it, because autograd prunes the whole sub-graph. In the reference run this removed 6.91% of training compute while touching only 0.35% of the parameters.
Does freezing layers hurt model accuracy?
Slightly, by an amount we cannot resolve at our sample size. Across three seeds on CIFAR-10, final validation accuracy sits 0.33 percentage points below the baseline when paired by seed — t = −2.0 with 2 degrees of freedom, where significance at the 5% level would require |t| above 4.30. All three seeds moved the same direction, so we read this as a small real cost rather than parity. An earlier single pair of runs showed a 1.05-point gap, which did not survive replication; that is the reason we no longer quote single-pair results.
How do you verify the optimization actually does something?
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, which is PyTorch's autograd engine confirming the backward sub-graph was pruned. In the reference run this went from 1,563 backward passes per layer to exactly zero. The test suite also compares real weight tensors and BatchNorm buffers before and after optimizer steps rather than trusting the requires_grad flag.
How do you stop thermal throttling from faking the speed-up?
By not trusting elapsed seconds. On a fanless laptop the same baseline configuration varies by about 20% between runs from thermal state alone, which is several times the effect being measured — so a cross-run comparison of wall-clock largely measures the machine. The study is built so this can be corrected: both levers engage at the start of epoch 2, which makes epoch 1 bit-identical across every arm at a given seed, so epoch 1 is a direct probe of how that machine was behaving during that run. Every timing figure is the engaged epochs divided by that run's own epoch 1, then divided by the baseline's ratio to remove warm-up drift. Runs whose epoch-1 throughput is more than 8% off the median, or whose throughput never recovers, are flagged by a fixed rule stated in advance, and results are reported both with and without them.
Have you measured an actual energy reduction?
Not yet, and we say so rather than implying otherwise. The reference run used CodeCarbon, which on Apple Silicon without root access falls back to a constant-TDP estimate — it attributed zero watts to the GPU and reported near-identical CPU power in both arms, so its energy figure is the elapsed-time figure in different units. Direct power integration through powermetrics is implemented and is the next measurement. Until it runs, the defensible claims are compute reduction and wall-clock reduction.
Does EnAi work on NVIDIA GPUs?
The mechanism is framework-level rather than vendor-specific. It uses standard PyTorch forward and backward hooks and a batch sampler, none of which are tied to a particular accelerator, so nothing in the design prevents it. Being precise about what has actually been tested: the reference implementation was built and measured on Apple Silicon through the Metal Performance Shaders backend, deliberately without PyNVML or vendor power limiters. It has not yet been benchmarked on CUDA, and we would rather say that than imply numbers we do not have.
Do I have to change my model code?
No. EnAi attaches to modules that already exist through PyTorch's hook API and mutates a batch sampler in place. There is no checkpoint, no restart, and no edit to the model definition — the run continues from exactly where it was when the optimizer engaged.
How is this different from just using a bigger batch size?
Raising batch size is a known technique, and it is one of three levers here — so we ran the ablation that could have shown it accounts for everything. Across three seeds, with machine state normalized out: freezing alone is worth +6.9% throughput, raising the batch alone +9.3%, and the two together +20.9%. Neither lever explains the other, and the combination beats what two independent levers would produce, meaning they interact rather than stack. The third lever is a linear learning-rate rescale that keeps the optimization trajectory comparable so the accuracy comparison is not confounded.
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.