$\sigma$Transfer: Uncertainty Transfer from Small to Large Networks under $\mu\mathrm{P}$

2 RELATED WORK

• µP scales initialization and layer multipliers so feature updates remain stable as width grows, enabling training hyperparameters tuned on smaller models to transfer.µTransfer demonstrated transfer from a 40M-parameter proxy to a 6.7B-parameter language model.Source: §2 · pp. 2–3
• Even with methods that reduce Hessian costs, Laplace inference still requires evaluating posterior candidates in a precision sweep that can dominate inference for million- or billion-parameter models.The Hessian costs O(P^2) to store and O(P^3) to invert; last-layer restrictions, Kronecker-factored curvature, and online evidence maximisation reduce these costs.Source: §2 · p. 3

3 PRELIMINARIES

• Linearized Laplace places a Gaussian posterior around a trained network’s parameters and propagates it to output predictions using a first-order expansion.The prior is Gaussian with precision λ, and the parameter posterior is centered at the trained parameters.Source: §3 · p. 3
• The output Jacobian and likelihood curvature determine how the parameter posterior yields Gaussian output predictions.The exposition defines the Jacobian with respect to stored parameters and the output-space negative-log-likelihood curvature at trained predictions.Source: §3 · p. 3
• With trained parameters fixed, larger prior precision contracts posterior covariance and smaller precision permits greater uncertainty; validation NLL maps each λ to predictive NLL.Classification probabilities average softmax over the Gaussian logit posterior, while Gaussian regression adds observation noise.Source: §3 · p. 3
• µP scales initialization, forward multipliers, and updates with width so layers update without causing logits or activations to diverge.Parameter tensors can be represented in width-dependent units whose scale factors remain fixed during training, with initialization and updates adjusted consistently.Source: §3 · p. 3

4 METHOD: σTransfer

• σTransfer seeks the same limiting linearized output posterior across widths, making precision selected on a small network transferable to a larger one.A quantity is width-stable when it converges entrywise on a fixed finite input set.Source: §4 · p. 4
• For a general Gaussian prior, the Woodbury rewrite expresses posterior covariance through the prior kernel, so stable prior geometry can yield stable covariance.The kernel is Kn(X, X′) = Jn(X)SnJn(X′)ᵀ; an isotropic stored-coordinate prior is not guaranteed to be width-stable.Source: §4 · p. 4
• The method chooses a width-normalized prior so its induced prior kernel converges across widths, which in turn supports posterior-covariance stability.In width-normalized coordinates, the isotropic prior kernel is Jϕ,n(X)Jϕ,n(X′)ᵀ and converges to K∞ under Theorem 1’s conditions.Source: §4 · p. 4
• The coordinate change uses a block-diagonal scale map per parameter tensor, expressing an isotropic prior in normalized coordinates as a width-normalized prior in stored coordinates.The prior precision for a tensor with width factor t is implemented as the scalar λ/t^2, without constructing Sn as a dense matrix.Source: §4 · pp. 4–5
• The rescaling leaves the trained checkpoint, predictions, likelihood curvature, and λ as the only tuned prior hyperparameter unchanged.The method’s formal transfer conditions are stated in Section 5; fixed-noise regression has width-independent curvature, while softmax curvature converges when trained logits converge.Source: §4 · pp. 4–5

5 THEORETICAL GUARANTEES

• For fixed-depth, equal-width ReLU µP MLPs trained for finitely many width-matched steps, the σTransfer prior kernel and trained outputs converge jointly almost surely on finite input sets.The theorem assumes the regularity conditions stated in Appendix K.1 and extends to any fixed output dimension.Source: §5 · p. 5
• Given prior-kernel convergence and width-stable positive-semidefinite likelihood curvature, full-GGN output covariance converges uniformly over any compact positive precision interval.The interval is [λ−, λ+] with 0 < λ− ≤ λ+ < ∞.Source: §5 · p. 5
• Uniform convergence of posterior mean and covariance on a finite validation set implies uniform convergence of validation-NLL curves and convergence of their minimizers when the limit has a unique minimum.Even without a unique minimizer, the theorem bounds excess NLL from using a proxy-selected precision on the target.Source: §5 · p. 5
• For continuous posterior-moment-based scores over a finite action set, a positive limiting best-action gap yields matching choices at sufficiently large widths.Without that gap, proxy-action target-score loss is at most twice the largest proxy–target score difference.Source: §5 · p. 5
• Beyond the stated finite-step ReLU MLP setting, transfer applies when the normalized prior kernel, predictive mean, and likelihood curvature converge; precision transfer by fixed-checkpoint evidence needs additional conditions.Last-layer Laplace reduces the kernel condition to convergence of the normalized final-feature Gram matrix; evidence selection also requires convergence of the prior-weighted squared parameter norm and a unique limiting maximizer.Source: §5 · pp. 5–6

6.1 REGRESSION

• On FreeSolv, σTransfer’s evidence curves nearly overlap across widths; across the three regression datasets its selection gap is 0.4–0.9 bits versus 1.5–5.0 for SP + isotropic.Its predictive transfer gap |∆NLL| is orders of magnitude smaller across all three datasets.Source: Figure 2 · p. 6, Table 1 · p. 6, §6.1 · p. 7

6.2 IMAGE CLASSIFICATION

• Across MNIST and Fashion-MNIST, σTransfer has the smallest |∆NLL| in every Table 2 column; on MNIST, proxy acquisition ends 0.003 accuracy below target-reference, with ρ̄ = 0.75 versus 0.07 for SP + isotropic.The worst prior-transfer case is 0.6 × 10−2 for FMNIST LL-diagonal; decision transfer is also reported on Fashion-MNIST, though downstream utility is claimed on MNIST.Source: §6.2 · p. 7, Table 2 · p. 7, Figure 3 · p. 8

6.3 LANGUAGE MODELS

• Across ten language-model tasks, the σTransfer proxy choice is 0.65 ± 0.34 bits from the target optimum and raises target NLL by 0.0014 on average, versus 2.65 ± 0.63 bits and 0.013 under SP.On public 1B→7B pairs, seven of ten tasks select the target’s precision and eight have absolute target-NLL changes below 10−4.Source: §6.3 · p. 8, Figure 4 · p. 8
• On AG News, σTransfer has the highest mean proxy–target ranking correlation across tested widths and label budgets, and the smallest mean absolute test-NLL gap in nine of twelve settings.For 128→2048 transfer with 5,000 initial labels, proxy-selected labels leave target accuracy 0.0014 below target-selected labels; methods use λ = 1.Source: §6.3 · p. 8, Figure 5 · p. 9

6.4 COMPUTATIONAL SAVINGS AND LIMITATIONS

• Speedups reach ∼5000× on MNIST and a median ∼2.3× on the public 1B→7B pair, ranging from 1× to ∼330× across nine timed tasks.The target still fits one posterior after selection, so end-to-end savings are capped by the 61-point grid.Source: §6.4 · p. 9, Table 3 · p. 9

7 CONCLUSION

• Prior-precision transfer selects precision on the small network and reuses it zero-shot on the larger network, while decision transfer uses the small network’s uncertainty estimates for acquisition, OOD flags, or abstention.The paper reports theoretical results supporting accurate transfer across widths and empirical reductions in uncertainty-task costs of up to several orders of magnitude.Source: §7 · pp. 9–10

AI USE STATEMENT

• The authors report that generative AI tools assisted with theoretical literature navigation, proofs, implementation, experiments, analysis, literature search, feedback, and writing, and that they reviewed the assisted work.They state that numerical claims were checked against saved outputs and theoretical statements against proofs.Source: AI USE STATEMENT · p. 10

REPRODUCIBILITY STATEMENT

• The appendices provide parameterization conventions, implementation details, assumptions, and complete proofs, while experimental protocols specify model widths, seeds, posterior approximations, selection rules, metrics, and compute measurements.Code, configuration files, and result manifests are stated to be included in supplementary material.Source: REPRODUCIBILITY STATEMENT · p. 10

• Theorem 1 covers fixed-depth ReLU MLPs under finite-step training and regularity assumptions; convergence conditions for Transformer settings and required proxy width remain open.Source: A LIMITATIONS AND FUTURE WORK · p. 14
• Extending σTransfer beyond width-only scaling to models whose width and depth grow together is identified as future work.The authors also identify medical prediction and generative models, including diffusion models, as possible future applications.Source: A LIMITATIONS AND FUTURE WORK · p. 14

B EVALUATION METRICS

• Precision selection is evaluated against target-side optima, while predictive transfer gap measures the resulting target test-NLL difference and cross-width KL measures predictive-distribution agreement.Target-side sweeps define reference optima; in zero-shot deployment the proxy-selected value is reused and the target fits one posterior at that value.Source: B EVALUATION METRICS · p. 14
• Decision-transfer metrics include acquisition-score rank correlation and top-K overlap, while AULC and OOD abstention metrics assess downstream utility and gating behavior.OOD abstention reports OOD prompts correctly rejected and in-distribution traffic served, with calibration on one OOD source and evaluation on another.Source: B EVALUATION METRICS · pp. 14–15

C EXPERIMENTAL DETAILS

• The main image-classification benchmark uses a two-hidden-layer ReLU MLP on MNIST across widths 128–4096, with 10 seeds and validation-NLL selection on a fixed 61-point precision grid.Targets reuse the proxy-selected numerical precision unchanged and construct one posterior; target-side sweeps provide reference optima.Source: C EXPERIMENTAL DETAILS · p. 15

C.1 CURVATURE APPROXIMATIONS

• Full GGN stores and inverts the full curvature matrix, diagonal GGN drops off-diagonal correlations, and KFAC captures within-layer structure at lower computational cost.Last-layer variants apply the approximation only to the final linear layer, with earlier layers fixed at their MAP values.Source: §C.1 · p. 15

C.2 THE GEOMETRIC MECHANISM

• On two moons, kernel drift falls from 20.4 ± 3.9 under standard practice to 0.64 ± 0.11 under σTransfer, while precision movement changes from 2.9 ± 0.1 bits to −0.5 ± 0.8 bits.Held-out uncertainty-rank agreement rises from 0.22 ± 0.05 to 0.69 ± 0.04.Source: §C.2 · p. 15
• On MNIST, linear CKA is about 0.999 for every method, indicating that the µP head scale stabilizes kernel magnitude rather than changing kernel shape.Kernel stability alone does not establish posterior stability after conditioning.Source: §C.2 · pp. 15–16

C.3 REGRESSION UNDER SCAFFOLD SHIFT

• The regression benchmark uses a tanh MLP on Morgan fingerprints with OGB scaffold splits, separating train, validation, and test molecules by scaffold.The inputs remain unstandardized, so features absent from training retain their 0/1 values when they appear in validation or test data.Source: §C.3 · p. 16

C.4 ADDITIONAL NORMALIZATION BASELINES

• σTransfer has the smallest reported mean NLL transfer gap across regression and classification settings, tying µP + empirical on MNIST LL-diagonal and FMNIST LL-full.These are comparisons of mean transfer gaps, not statistical significance or absolute predictive quality.Source: §C.4 · p. 16

C.5 ADDITIONAL CLASSIFICATION DATASETS

• The PenDigits and Letter experiments use full-GGN Laplace over all parameters and select precision by maximizing fixed-checkpoint Laplace log-evidence.The experiments extend beyond the convergence guarantee of Corollary 5.Source: §C.5 · p. 17
• Proxy-side evidence selection gives approximately 22× and 32× speedups on PenDigits and Letter, with target test-NLL changes of −0.007 and 0.078 relative to target-side evidence selection.The target precision is reused unchanged after selection on the proxy; timings include Laplace fitting and precision search but exclude network training.Source: §C.5 · p. 17

D EVIDENCE-SELECTED TRANSFER ON MNIST/FMNIST

• With plug-in evidence selection on MNIST/Fashion-MNIST, a transferred precision can have lower test NLL than the target’s own selection because the target evidence optimum can miss the test optimum.The paper therefore uses validation NLL as its primary criterion.Source: D EVIDENCE-SELECTED TRANSFER ON… · p. 17, Table 7 · p. 17

E ADDITIONAL OOD-TRANSFER RESULTS

• 4.1 seconds versus 6.4 minutes (94×): proxy-side versus target-side OOD precision selection, with no measurable OOD-quality loss under σTransfer.Zero-shot transfer removes the target sweep but still requires one target posterior at the transferred precision.Source: E ADDITIONAL OOD-TRANSFER RESULTS · p. 17, Figure 6 · p. 18
• OOD precision is selected using MNIST-validation versus FMNIST and evaluated on disjoint MNIST-test versus EMNIST and KMNIST; signed regret can be slightly negative on the unseen distributions.In those cases, the proxy-selected precision generalizes marginally better to the evaluation distributions.Source: E ADDITIONAL OOD-TRANSFER RESULTS · pp. 17–18, Table 8 · p. 18
• On Fashion-MNIST, σTransfer’s mean acquisition-score rank correlation is 0.73 against 0.36–0.43, while the acquisition curves do not separate from random.The text claims decision transfer on both datasets, but downstream utility only on MNIST.Source: Figure 7 · p. 19
• On MNIST, σTransfer’s acquisition fidelity interval covers zero (−0.0001, [−0.005, +0.005]) and its proxy-minus-random AULC contrast is +0.018 ([+0.014, +0.022]).The SP + isotropic fidelity result is misleading because its target reference itself falls 0.050 below random.Source: Figure 8 · p. 19

F SEQUENTIAL ACTIVE LEARNING DETAILS

• Each run starts with 20 labelled examples and acquires 20 per round until 1,000 are labelled, scoring 500 sampled candidates per round with EPIG.The comparison uses widths 128 and 4096, 20 seeds, 200 reference inputs, and 256 posterior samples; acquisition uses a last-layer diagonal Laplace posterior at fixed λ = 1.Source: F SEQUENTIAL ACTIVE LEARNING DETAILS · p. 19
• On Fashion-MNIST, every method’s acquisition utility interval excludes zero because EPIG acquisition does not beat random; σTransfer is least affected at −0.011, versus −0.047 for SP + isotropic.The respective intervals are [−0.018, −0.003] and [−0.055, −0.039].Source: Figure 9 · p. 20
• On MNIST, mean acquisition-score rank correlation is 0.75 for σTransfer, versus 0.38 for µP + isotropic, 0.24 for SP + metric, and 0.07 for SP + isotropic.The proxy and target are widths 128 and 4096, respectively; curves use the same complete-trajectory seed set.Source: Figure 10 · p. 20

G DECISION TRANSFER FOR LANGUAGE MODELS

• The width-128 proxy supplies the uncertainty score and decides whether the width-2048 model should answer; the large model supplies the prediction.The proxy’s precision is selected using ID validation data and OOD source A, and its threshold is applied to proxy scores when evaluating on source B.Source: G DECISION TRANSFER FOR LANGUAGE MODELS · pp. 20, 22, 23
• 0.73 mean decision agreement under σTransfer on Fashion-MNIST, versus 0.36–0.43 for three baselines; SP + isotropic agreement is 0.36 on Fashion-MNIST versus 0.07 on MNIST.Source: Figure 11 · p. 21
• Compared with calibrated target-posterior decisions, using the proxy reduces OOD recall by approximately 10–12 percentage points, ID retention by 1–2 points, and answered accuracy by 7–8 points.Deployment requires only the proxy posterior and target MAP predictions.Source: G DECISION TRANSFER FOR LANGUAGE MODELS · p. 23
• Proxy OOD recall is 0.41–0.42 with 1 −MSP while retaining approximately half the ID prompts; applying it to the target gives OOD recall of 0.84–0.85 and answered accuracy of 0.64–0.66.The target’s mean NLL on answered ID prompts differs from ungated prediction by −0.004/−0.005 nats for SP/σTransfer.Source: G DECISION TRANSFER FOR LANGUAGE MODELS · p. 23

H PRIOR-PRECISION TRANSFER IN LANGUAGE MODELS

• The fixed-checkpoint plug-in evidence criterion selects precision on proxy ID training data, while target-side search is an evaluation reference; deployment builds one target posterior at the transferred precision.Source: H PRIOR-PRECISION TRANSFER IN LANGUAGE… · p. 23

H.1 COMPARISON WITH THE DIAGONAL NEURAL G-PRIOR

• 5.5→10.2 bits mean optimum drift for the isotropic prior accompanies mean target-NLL degradation of 0.028, reaching 0.1 on TREC.On the same µP checkpoints, the isotropic prior drifts faster than SP; its mean optimum gap is 4.7 ± 1.3 bits at widths 128→2048.Source: Figure 13 · p. 24
• 3.95 bits for SP and 3.50 bits for µP are the mean absolute precision gaps, with mean target test-NLL increases of 0.0367 and 0.0229, respectively, under the diagonal neural g-prior.The experiment transfers α from width 128 to 2048 with fitted classification heads fixed throughout the precision search.Source: §H.1 · pp. 23–24

I EXTERNAL VALIDATION ON PUBLIC U-µP/SP PAIRS: DETAILS

• 0.56 bits is the mean precision gap for both σTransfer and SP with an isotropic prior under the corrected evidence criterion.The correction estimates the gain in log posterior from locally adjusting trained weights at each precision; all selected precisions are interior to the grid.Source: I EXTERNAL VALIDATION ON PUBLIC… · pp. 24–25
• Public-model proxy gates retain approximately 99.8% of in-distribution inputs and reject 96% of out-of-distribution inputs on average, with similar performance on a common target.Source: I EXTERNAL VALIDATION ON PUBLIC… · p. 25

J TECHNICAL FORMULATION OF THE σTransfer PRIOR

• For the MLP construction, the prior blocks scale as nI for weight and hidden-bias blocks and as 1 for the output bias.The resulting prior-map factors are √(n d_in) on input weights, √n on remaining weights and hidden biases, and 1 on output bias.Source: J TECHNICAL FORMULATION OF THE… · p. 26
• Under the stated assumptions, the kernel converges through joint almost-sure convergence of Gram entries on each fixed finite input set.The lemma’s result extends to any fixed output dimension.Source: J TECHNICAL FORMULATION OF THE… · p. 26
• Reported checks have relative Frobenius error below 2×10−8 for all-parameter regression, all-parameter classification, and last-layer classification; under blockwise rescaling, Proposition 1 function-covariance identities err below 2 × 10−15.Source: J TECHNICAL FORMULATION OF THE… · p. 27

K PROOFS OF THE THEORETICAL STATEMENTS

• Posterior and transfer results require stable data-space curvature, a compact precision search, and convergence of the selected subset’s prior-weighted squared norm for plug-in evidence-score transfer.Exact eventual decision agreement additionally assumes nonzero decision margins; analogous premises holding in probability yield agreement probability tending to one.Source: K PROOFS OF THE THEORETICAL STATEMENTS · p. 27

K.1 FORMAL STATEMENT AND PROOF OF THEOREM 1: PRIOR-KERNEL STABILITY

• With fixed-depth, finite-step μP training, shared training paths, and no asymptotic preactivation atom at the ReLU kink, smoothing ReLU leaves outputs and forward and backward Gram entries asymptotically unchanged.The finite-horizon stability lemma also establishes vanishing discrepancies in effective hidden matrices, wide vector parameters, and scalar parameters.Source: §K.1 · pp. 27–28
• Under Assumptions 1 and 2, the fixed-depth MLP prior-kernel matrix on a finite input set converges entrywise almost surely as width grows, and trained outputs converge jointly.The theorem follows from convergence of the kernel’s Gram factors and continuity of their finite sums and products.Source: §K.1 · p. 29
• For fixed output dimension, the output and output-pair kernel entries converge jointly almost surely under the same assumptions, including the prior-kernel blocks.Separate readout and output-bias blocks contribute the Kronecker-delta factor, while shared hidden blocks contribute the remaining kernel terms.Source: §K.1 · pp. 29–30

K.2 COORDINATE-INVARIANT INTERPRETATION OF THEOREM 1

• The σTransfer prior is equivalent to an isotropic local prior in width-normalized perturbation coordinates, where stability means the normalized-coordinate Jacobian kernel has a large-width limit.The stored-coordinate and normalized-coordinate Jacobians are related by the coordinate transformation described in the subsection.Source: §K.2 · p. 30

K.3 FORMAL STATEMENT AND PROOF OF THEOREM 2: POSTERIOR-COVARIANCE STABILITY

• When prior-kernel blocks converge and likelihood curvature converges in operator norm, the GGN posterior covariance converges uniformly over every compact interval of positive precisions.The result assumes finite data and evaluation sets and a positive-definite prior covariance at each width.Source: §K.3 · p. 30
• If the final-layer feature Gram matrix converges, the last-layer prior kernel is 1 + ΣL(x, x′) and its Laplace posterior covariance is stable; a vanishing exact-Hessian residual gives the GGN covariance limit.The exact-Hessian convergence is uniform in precision on compact positive intervals.Source: §K.3 · pp. 32–33

K.4 STRUCTURED-CURVATURE EXTENSIONS: DIAGONAL AND KFAC

• For diagonal and block-Kronecker curvature, finite-width computations retain the width-normalized coordinate interpretation, while width stability and transfer conclusions hold conditionally on projected spectral-measure convergence and the other stated assumptions.The spectral-measure convergence condition is not directly tested in the structured-curvature experiments; their reported finite-width transfer is empirical evidence, not verification of the asymptotic condition.Source: §K.4 · pp. 33–34
• Prior-kernel convergence alone does not ensure diagonal posterior-covariance convergence: an alternating width sequence can keep the kernel constant while covariance fails to converge, and KFAC also depends on its retained block-factor spectrum.The example motivates the additional spectral condition for structured-curvature stability.Source: §K.4 · pp. 34–35
• Proposition 2 extends to positive-semidefinite curvature approximations whose projected spectral measures satisfy its assumption, with implementation checks covering full, diagonal, and KFAC curvature for Linear/MLP models.Source: §K.4 · p. 35

K.5 FORMAL STATEMENT AND PROOF OF THEOREM 3: INFERENCE-HYPERPARAMETER TRANSFER

• Uniform validation-NLL convergence yields consistent precision minimizers when the limiting minimizer is unique, with transferred excess NLL bounded by twice the proxy–target discrepancy.The result covers scalar Gaussian regression and softmax classification with Gaussian logit posteriors.Source: §K.5 · pp. 35, 37
• Fixed-checkpoint plug-in evidence maximizers and finite-grid optima transfer under additional q_n convergence and full-GGN assumptions, but refitted evidence is not covered.The criterion applies only to the chosen Laplace parameter subset; likelihood-null directions can change its optimum without changing the predictive posterior.Source: §K.5 · pp. 37–38
• Approximate precision minimizers converge to the unique limiting optimum, and finite-grid minimizers eventually agree with the limiting grid optimum.The approximate-optimization result requires optimization error below the limiting margin; the grid result requires a positive gap between its best and second-best limiting points.Source: §K.5 · pp. 38–39

K.6 FORMAL STATEMENT AND PROOF OF THEOREM 4: DECISION AND ACTIVE-LEARNING TRANSFER

• For a finite action set, posterior convergence and a unique limiting optimum with positive margin make proxy-selected actions asymptotically optimal and decisions eventually agree across widths.The proxy precision may be selected by validation NLL or the fixed-checkpoint plug-in evidence score.Source: §K.6 · pp. 39–40

K.7 EXAMPLES OF ACQUISITION FUNCTIONS COVERED

• Marginal posterior variance, Gaussian information gain, integrated variance reduction, and predictive information gain converge over finite candidate sets or batches under positive observation noise.Convergence of the covariance blocks and positive-definite noise-augmented matrices supports these acquisition-score limits.Source: §K.7 · p. 41
• BALD-style scores converge, while top-K rankings and selective or detection flags transfer when score or threshold margins are nonzero.The BALD result also covers a fixed common-random-number Monte Carlo rule, but independent fixed-sample draws across widths require a separate error bound.Source: §K.7 · pp. 41–42
• Under conditional convergence and nonzero-margin conditions, the selected action agrees at each round, yielding agreement of the entire finite acquisition trajectory.The result follows by induction along the common limiting acquisition history.Source: §K.7 · p. 42