Skip to content

The Forgiveness That Dies First

Published: at 01:42 PM
0 views

The learning rate is the most expensive hyperparameter in large-scale training, and it is also treated as the dullest one. Small models, big models, they all get swept the same way and tolerated the same way. Maximal update parametrization made it interesting. With muP you tune the rate once on a small proxy and reuse it unchanged on a model that is orders of magnitude wider. Tune wide, port wide, never retune. That is the whole promise, and it has shipped rates from millions of parameters to billions.

Everyone quietly assumed that trick dies the moment you train at low precision. I ran it and found the assumption is wrong in an instructive way. The transfer survives. What dies first is the thing that made transfer feel safe in the first place. The full study, with every number reproducible, is up on Hugging Face.

Why muP and quantization are even coupled

muP works by choosing the weight scaling so that activation magnitudes and per-step updates both stay order one as width grows. That is what makes the optimal learning rate approximately width-independent, and it is the license for zero-shot transfer. A practical side effect falls out of the same choice. Because activations and updates are both order one, a wide range of learning rates, more than a decade of them, stays near-optimal. That forgiveness is what makes proxy tuning robust. You can be wrong about the rate and the model still learns.

Quantization touches the same place. Low-precision rounding acts directly on the coordinate magnitudes muP is designed to keep bounded. A 4-bit grid resolves those order-one coordinates to only a handful of levels, so a single per-tensor scale for a whole layer has to serve every unit, and the per-step update is effectively a bound on how far a coordinate may move in one step, an effective cap on the learning rate. The two mechanisms are coupled by design, and the question that matters is which part of muP breaks first.

I did not modify the parametrization. I kept a standard muP and asked what simulated quantization does to it. That is the complement to u-muP, which changes muP to be precision-friendly by construction. I wanted to know whether the transfer property survives honest low precision on its own.

The setup, kept deliberately small

ReLU multilayer perceptrons on FashionMNIST, four width-to-width hidden layers, hidden widths from 64 to 1024 for the configs that carry the transfer claims, a factor of sixteen in width. Adam, batch 128, 1000 steps, three seeds per cell, six for the highest-variance ones. Quantization is simulated symmetric fake-quantization with straight-through gradients, at eight and four bits, applied to weights, activations, or both, in per-tensor, per-output-unit, and block-scaled sixteen-element shared-scale forms. The block form is MX-inspired but is not an exact OCP MXFP4 configuration, which uses thirty-two element blocks with E8M0 exponents. Every cell is an independent initialization. No transformer, no embedder, no billion-parameter claim. The point is the mechanism and the ordering of effects, not production numbers.

How you quantize matters as much as how much

The first result is the one with the most immediate practical weight. Weights-only 4-bit sits within about a hundredth of fp32. Per-unit and block-scaled 4-bit stay within roughly 0.05 to 0.08 nat of the fp32 best loss. Per-tensor 4-bit, applied to weights and activations, takes a hit of 0.3 to 0.4 nat. That is a four-to-eight times larger penalty for the same bit count, and the difference is entirely in how the quantization is laid out. A per-tensor scale for a full activation tensor is the pessimistic extreme. The finer-grained formats that real hardware, INT4, MXFP, per-channel, use, reclaim nearly all of it.

Best attainable test loss versus quantization format for each parametrization at widths 64, 128, 512, and 1024, mean with standard-deviation bars. Per-tensor 4-bit weight-plus-activation is the only damaging config; per-unit, block-scaled, and weights-only 4-bit sit close to fp32.

Figure 1. Best attainable test loss versus quantization format. Only per-tensor 4-bit is destructive; the finer-grained formats sit within noise of fp32.

The stability edge collapses first

The next result is where muP actually gets hurt. Per-tensor 4-bit collapses muP’s stable learning-rate edge by an order of magnitude. In fp32 the model tolerates a rate of 2e-1 at every width. Under per-tensor 4-bit that edge falls to between 1e-2 and 5e-2, depending on width. A rate that is comfortable in full precision becomes decisive at 4-bit. muP at width 64 reaches loss around 0.41 to 0.46 at 2e-1 in fp32, and 2.3 to 19 under per-tensor 4-bit. The robust check: raising the collision threshold from a terminal loss of 1.0 to 2.3 moves the edge only between 1e-2 and 5e-2, so this is not an artifact of the cutoff. Quantization converts muP’s wide flat optimum into a narrow channel.

The forgiveness, the decade of good rates, is the first casualty. That is the headline.

configSPmuP
fp325e-2 .. 1e-22e-1
8-bit1e-2 .. 5e-22e-1
4-bit per-tensor3e-3 .. 7e-41e-2 .. 5e-2
4-bit weights-only1e-22e-1
4-bit per-unit1e-2 .. 5e-22e-1

Largest stable learning rate (terminal loss under 2.0) across widths 64 to 512. muP’s edge holds through 2e-1 at every fine-grained format and only per-tensor 4-bit knocks it down, by an order of magnitude.

The transfer property survives even then

Here is the counterintuitive half. Under that same pessimistic per-tensor 4-bit stress, the width-independence of muP’s optimal rate remains. The refined grid places the optimum at 8e-3, 8e-3, and 1.3e-2 across widths 64, 128, and 512, all within the resolution of a factor-of-1.6 grid. Standard parametrization drifts the other way, from about 2e-3 down to 5e-4 across the same widths at fp32, a factor-of-four move in the wrong direction.

Optimal learning rate (grid argmin) versus quantization format, median over three seeds with the seed-to-seed range. muP's optimum sits an order of magnitude above SP's and stays roughly width-invariant through width 1024; at per-tensor 4-bit it remains near 1e-2 across widths, while SP's drifts to low rates.

Figure 2. Optimal learning rate versus quantization format. muP’s optimum is higher and stays roughly width-invariant; SP’s drifts downward.

I did not stop at the width-invariance of optima, which is related to transfer but is not itself a transfer experiment. I ran the transfer directly. Train at width 64, select the rate on validation, freeze it, retrain widths 128 and 512 at exactly that rate, evaluate each on held-out test. For muP the transferred rate lands within the seed-to-seed variation of the per-width oracle at both precisions, which reads as no measurable loss from refusing to retune. Standard parametrization is the opposite. At fp32 it already gives up about 0.05 nat by refusing to retune at width 512. Under per-tensor 4-bit the same single width-64 rate drives the width-512 model to a test loss of 133.95, a divergence. Median SP transferred losses are still worse than the oracle. Core to the result: muP zero-shot transfers under aggressive activation quantization, and standard parametrization does not.

The qualitative picture holds on CIFAR-10. muP stays high and forgiving across width; SP droops and then breaks.

Zero-shot learning-rate transfer measured directly. Learning rate tuned on the validation split at width 64, frozen, then applied at widths 128 and 512 (blue) against the per-width validation-tuned oracle (orange), held-out test loss, log scale, mean over three seeds. muP's transferred rate is at or below the oracle at both precisions; SP's transferred rate loses at width 512 in fp32 and diverges under per-tensor 4-bit.

Figure 4. Zero-shot learning-rate transfer, measured directly. muP’s transferred rate matches the per-width oracle; SP’s does not survive.

The mechanism I measured

I instrumented the first width-to-width hidden pre-activation under per-tensor 4-bit on a width-512 muP model and recorded two quantities per step: how far a coordinate moves in units of the 4-bit grid step, and the relative rounding error.

The rounding error is large and roughly flat across all eight learning rates, sitting near 0.16 to 0.26. Because muP pins the pre-activation spread to order one, a single per-tensor scale resolves it into only a handful of grid levels regardless of rate. The learning rate is the knob that decides what that coarse grid does. At the low end a coordinate moves only 0.66 of a grid step per optimizer step, so it is effectively frozen and the run under-updates. Near the working rate the movement rises to just over one grid step per step, over 1.0, and the loss reaches its minimum. Push the rate up and the loss climbs sharply.

The candidate control variable is the update-to-grid ratio, how far a coordinate moves per step relative to the resolution the grid gives it, not the instantaneous rounding error, which is already large and flat. The finer-grained per-unit and block formats shrink the rounding error directly, which is why they also widen the stable channel back. Standard parametrization resolves the same puzzle in the other direction: under SP the deeper-layer activation magnitudes grow with width after an update, from about 0.10 at width 32 to past 1 at width 512 and past 4 at width 2048. The shared quantization grid must resolve a width-growing coordinate, which is exactly why SP’s quantization error and instability grow with width, and why muP’s do not.

Mechanism diagnostic under per-tensor 4-bit, muP, width 512, mean over two seeds. Left: per-step pre-activation movement in 4-bit grid units (blue) and relative 4-bit rounding error (orange) versus learning rate; movement rises to just over one grid step per step near the working rate, and the rounding error is large and flat. Right: final loss of the same runs versus learning rate.

Figure 3. Mechanism diagnostic: per-step movement in 4-bit grid units versus the rounding error, and the resulting terminal loss. The learning rate raises the update-to-grid ratio, not the instantaneous rounding error.

I state that mechanism at the level I measured it. The quantities move as the hypothesis predicts under one architecture, one optimizer, and simulated quantization. It is a candidate mechanism, not a proof.

What I take from this

Quantization attacks muP’s headroom before it attacks transfer. That is the sentence I want you to keep. The property everyone worried would break, the ability to tune once and port wide, holds up even at 4-bit. The property nobody was watching, the wide tolerant plateau that made tuning feel safe, is the part that gets squeezed down to a channel. If you train at low precision, you can still trust the transfer, but you can no longer be sloppy about the rate, because the room to be wrong is mostly gone.

The practical rules fall out cleanly. Choose a finer-grained format when you can, per-unit or block-scaled, because it reclaims most of the forgiveness for the same bit count. Do not extrapolate the stability of a full-precision rate to a 4-bit run, because the ceiling moves down an order of magnitude. And if you use muP for transfer, the transferred rate still lands, but recheck the channel it has to land in.

The limits are the honest part. Multilayer perceptrons over a factor of sixteen in width, not transformers, not the scales where muP is deployed, and simulated quantization rather than true integer paths with low-precision accumulators. The next steps are obvious and worth doing. Transformers, real per-channel and block formats with genuine integer paths, and a derivation of the exponent tying the stability edge to activation resolution from muP’s spectral condition. The result here is a clear answer to a bounded question, and it opens the interesting work rather than closing it.

I ran this because my own training runs live at low precision, and I wanted to know which of my assumptions about muP would survive contact with the hardware. The answer surprised me in exactly the way good experiments should. I lost the thing I did not know I was leaning on, and kept the thing everyone told me I would lose.

Code, data, figures, and the manuscript itself are public under CC BY 4.0 on Hugging Face, including every per-cell sweep so the tables here can be re-derived to the printed precision.


Previous Post
The Mirror Problem
Next Post
Making AI Agents Think Before They Retrieve