PaperlayerView exact revision on arXiv

Source-linked AI summary

Transferable Graph Metanetworks

Yuxin Ma, Adir Dayan, Yam Eitan, Haggai Maron, Soledad Villar

arXiv:2610.00420v1stat.MLcs.LG

TL;DR

Metanetwork size generalization remains limited, despite the potential value of training on small networks and applying to larger ones. The paper modifies graph metanetworks to support transfer across widths, and reports improved performance, particularly for µP-trained inputs, with theoretical guarantees in that setting.

  • Problem

    Existing metanetwork size-generalization studies are limited, while training on small networks and applying to larger ones could address scarce large checkpoints and high training costs.

    Source: §1 · p. 3

  • Method

    The paper modifies graph metanetworks for invariance to function-preserving width changes and continuity across weights representing similar functions.

    Source: §1 · pp. 3–4

  • Results

    Both modification levels significantly improve size generalization across parameterizations, with matrix-product GMNs on µP-trained inputs performing well across the full test-width range.

    Source: §1 · p. 4

  • Takeaways & Limitations

    Strong size generalization is possible when the metanetwork architecture matches the distribution of input weights, including compatibility with function-preserving widenings and continuity under functional similarity.

    Source: §6 · p. 12

    The theory covers only MLP input networks and width generalization; depth generalization remains unaddressed.

    Source: §6 · pp. 12–13

Abstract

from arXiv · show

A weight space network (or metanetwork) takes the weights of another neural network as input and predicts properties of it. Most prior work trains such models on input networks of one or a few fixed sizes and evaluates them in-distribution. The few attempts at out-of-distribution size generalization remain limited in scope and have achieved only modest success. Consequently, the potential efficiency gains of training on small networks and evaluating on much larger ones remain largely unrealized. We propose Transferable Graph Metanetworks, which extend the graph metanetwork paradigm with a set of modifications that make performance transferable across input networks of different widths. The modifications follow two principles: invariance to the ways in which networks of different widths represent the same function, and continuity, such that weights representing similar functions receive similar predictions. We further study whether size generalization is possible for input networks trained independently from random initialization. Empirically, our modifications significantly improve size generalization on every task we consider. Performance is strongest on input networks trained under the maximal-update parameterization ($\mu$P), where it remains robust up to $42\times$ the training width. Theoretically, we explain these observations with infinite-width limit theory: we prove size-generalization guarantees for our model on $\mu$P-trained inputs, and explain why it can fail under other parameterizations.

1 Introduction

The introduction asks whether metanetworks can train on small input networks and generalize to larger widths, where existing approaches have limited success. It presents transferable graph metanetworks and reports improved size generalization, especially for inputs trained with µP.

Motivation

  • Training on small, plentiful networks and applying to unseen larger ones could address the scarcity of large checkpoints and the cost of training metanetworks on them.The introduction notes that prior GNN size-generalization results do not directly apply to metanetworks because their inputs have particular structure and symmetries.Source: §1 · p. 3

Research gap

  • Prior graph metanetworks accept variable-size networks, but size generalization was not a design goal and reported tests showed only limited success.Many equivariant-layer metanetworks instead target fixed weight spaces.Source: §1 · p. 3

Proposed framework

  • Transferable Graph Metanetworks modify graph metanetworks to transfer performance across input widths, using invariance to function-preserving changes and continuity for functionally similar weights.The authors leave depth generalization to future work.Source: §1 · pp. 3–4, Figure 1 · p. 4

Empirical and theoretical contributions

  • Both proposed modification levels significantly improve size generalization across parameterizations; matrix-product GMNs with µP-trained inputs perform well across the full test-width range.The paper also proves size-generalization guarantees for matrix-product GMNs on µP-trained inputs and identifies obstructions that can prevent generalization otherwise.Source: §1 · p. 4

2 Preliminaries and related work

The preliminaries review weight-space symmetries and graph metanetworks, then introduce compatibility and continuity as transferability requirements across sizes. They also explain how µP relates training dynamics across widths and supports the paper’s theoretical analysis.

2.1 Weight-space symmetries, metanetworks, and graph metanetworks

  • Metanetwork designs include fixed-size equivariant layers, variable-size graph models, and methods based on probes or derived signals; the paper focuses on graph metanetworks.Probe-based methods require suitable probe data and may miss information not exposed by their chosen inputs or derived signals.Source: §2.1 · p. 5
  • A graph metanetwork applies message passing to an MLP computation graph, with one node per neuron and shared parameters independent of input width.Its graph-level prediction is produced by pooling final node features.Source: §2.1 · pp. 5–6

2.2 Any-dimensional learning, transferability, and size generalization

  • Compatibility preserves predictions on equivalent inputs across sizes, while continuity additionally transfers predictions to larger inputs that are close under a target-respecting metric.In weight spaces, equivalence and closeness correspond to networks representing the same and similar functions, respectively.Source: §2.2 · p. 6

2.3 µP, infinite-width limits, and tensor programs

  • µP and standard parameterization differ in how initialization and hyperparameters scale with width, shaping the relationship between trained weights across widths.Tensor-programs analysis characterizes limiting activations and outputs during training; under µP, networks at every width approximate the optimal infinite-width limit with error O(1/√n).Source: §2.3 · pp. 6–7

3 Transferable Graph Metanetworks

Transferable Graph Metanetworks exploit cross-width weight transformations that preserve a network’s function. Their first modification makes graph metanetworks invariant to duplication widening.

3.1 Duplication-compatible Graph Metanetworks

  • Weights at different widths can represent the same function, a cross-width structure that prior metanetwork designs had not exploited.The section focuses on widening hidden layers while preserving the represented function.Source: §3.1 · p. 7
  • Duplication widening duplicates hidden neurons with rescaled weights, while preserving the network’s function.It generalizes beyond simple MLPs to almost all standard architectures.Source: §3.1 · p. 7
  • A duplication-compatible GMN gives the same prediction for weights before and after any duplication widening.Fan-in-rescaled edge features cancel the widening’s scaling factor, so features duplicate consistently.Source: §3.1 · p. 8

2. Mean aggregation: replace the two sum-aggregations in (2) by mean-aggregations,

Mean aggregation and layer-wise mean readout make GMNs invariant to duplication widening. Matrix-product GMNs address broader function-preserving widenings and align continuity with a norm that tracks function similarity.

2. Mean aggregation: replace the two sum-aggregations in (2) by mean-aggregations,

  • Mean aggregation prevents wider neighborhoods from multiplying incoming messages, and layer-wise mean readout preserves duplication invariance.Together with fan-in-rescaled edge features, these changes make the GMN duplication-compatible without reducing fixed-width expressivity.Source: §2 · p. 8

3.2 Matrix-product Graph Metanetworks

  • Function similarity is controlled by the normalized spectral norm, but duplication-compatible GMNs need not be continuous in that norm for multilayer inputs.They are Lipschitz in the normalized entrywise ℓ1 norm, which may not track function similarity; this misalignment leaves predictions at unseen widths unconstrained.Source: §3.2 · p. 9
  • Block-wise widenings preserve the represented function by redistributing copied weights across rows; duplication widening is their constant-block special case.They are the most general widenings in which hidden neurons copy those of the narrower network.Source: §3.2 · p. 9
  • Matrix multiplication makes GMN weight processing both compatible with block-wise widenings and controlled by the spectral norm.The change addresses the duplication-compatible GMN’s limited widening invariance and continuity mismatch.Source: §3.2 · p. 9
  • The forward matrix-product GMN is invariant to every block-wise widening; the bidirectional version requires equal column sums in the widening blocks.With Lipschitz MLP components, the matrix-product GMN is Lipschitz on every bounded set.Source: §3.2 · p. 9
  • The matrix-product GMN is a strict restriction of the duplication-compatible GMN and is presumed less expressive at fixed width, though experiments show no performance loss.This is an expressivity trade-off, not an observed fixed-width performance drop.Source: §3.2 · p. 9

4 Experiments

The experiments compare GMN variants on two prediction tasks while training at one small width and evaluating at larger widths. The modifications improve size generalization across tasks, datasets, and models, with the strongest robustness when the matrix-product GMN receives µP-trained inputs.

4 Experiments

  • The GMN is trained at one small width and evaluated at larger widths on INR classification and CNN generalization prediction tasks.The reported setups include bidirectional GMN and forward ScaleGMN architectures; other architectures and experiment details are deferred to the appendix.Source: §4 · pp. 10–11

Size generalization on hand-crafted widenings.

  • Hand-crafted duplication and block-wise widenings preserve the represented function, and the modified GMNs’ results match the theory in this controlled sanity check.These widenings are applied to training-set networks.Source: Size generalization on hand-crafted… · p. 11

Size generalization on independently trained weights.

  • Across every dataset, performance improves from the baseline GMN to the duplication-compatible GMN and then the matrix-product GMN; the matrix-product GMN on µP inputs stays nearly flat across larger widths.Removing either the matrix-product model or µP inputs visibly degrades the near-constant scores.Source: Size generalization on independently… · p. 11
  • For every model, size generalization is much better on µP-trained inputs than on SP-trained inputs, despite matching training distributions at the training width.The authors attribute this pattern intuitively to more similar training dynamics across widths under µP and greater weight variation across widths under SP.Source: Size generalization on independently… · p. 11

5 Theory of size generalization on trained weights

The theory connects size generalization to how trained-weight distributions change with width and to the distance the GMN respects. Under µP, it gives vanishing-gap guarantees, while identifying an obstruction under SP.

5 Theory of size generalization on trained weights

  • Both guarantees use infinite-width training-dynamics theory to control the size-generalization gap, defined as the worst test-error discrepancy between training width and larger widths.The authors expect their empirical improvements under SP to extend beyond the regime covered by the theory.Source: §5 · pp. 11–12
  • The first guarantee applies to one-layer inputs because the matrix-product GMN is continuous in a distance that measures function similarity after quotienting hidden-unit permutations.Under µP, hidden-neuron weights behave like width-independent i.i.d. draws during training, making trained weights at different widths close in that distance with high probability.Source: §5 · pp. 11–12
  • For one-hidden-layer MLPs under µP, the trained-weight laws converge uniformly across larger widths in the distance respected by the matrix-product GMN, yielding a vanishing size-generalization gap for Lipschitz targets and losses.This applies to duplication-compatible GMNs with Lipschitz components, including matrix-product GMNs.Source: §5 · p. 12
  • Under SP, initialization makes the relevant trained-weight distance diverge across widths, so the weights do not settle to a common limit and unseen-width predictions remain unconstrained.This explains why the one-layer guarantee does not extend to SP.Source: §5 · p. 12
  • At fixed depth, a second µP guarantee gives a uniformly vanishing size-generalization gap when the loss is Lipschitz and the target is representable by the matrix-product GMN.As stated, it assumes exact target representability and measures distribution differences using an integral probability metric generated by GMN losses.Source: §5 · p. 12

6 Conclusion

Strong size generalization is possible when the metanetwork architecture matches the distribution of input weights. The paper develops width-focused structures for comparing networks that represent the same function and leaves depth generalization open.

6 Conclusion

  • Matching the metanetwork architecture to the input-weight distribution is required for strong size generalization, which the paper establishes theoretically and empirically.Under μP, wider networks trained with matched settings have similar dynamics and lie near the same nontrivial limit, with O(n^-1/2) finite approximation error.Source: §6 · p. 12
  • The analysis addresses input-network width, while depth generalization remains open; the theory also covers only MLP input networks.The authors say it is unclear whether comparably simple GMN modifications can provide invariance to depth-changing transformations.Source: §6 · pp. 12–13

A Structure of the weight space

  • The weight-space structure combines invariance to function-preserving widenings with continuity under a distance reflecting functional similarity.These are presented as natural requirements for transferring predictions across widths.Source: A Structure of the weight space · p. 17

A.1 Widening by duplicating neurons

  • Block-wise widening preserves the represented function exactly when its layerwise weight blocks satisfy the stated if-and-only-if conditions.This characterizes the most general function-preserving widening that keeps wide hidden neurons as copies of narrow ones.Source: §A.1 · p. 18

A.2 Widening by zero-padding neurons

  • Zero-padding adds switched-off neurons, but the authors do not expect invariance to it to help size generalization for independently initialized trained networks.Random initialization and permutation-equivariant training do not single out neurons expected to remain zero.Source: §A.2 · p. 20

A.3 The weight space as a consistent sequence

The appendix formalizes duplication widening as a consistent sequence of weight spaces, with compatible norms and permutation symmetry. It also proves properties of the modified GMNs, including duplication compatibility and the matrix-product GMN’s broader widening compatibility.

A.3 The weight space as a consistent sequence

  • Duplication widening gives a consistent sequence whose limit space identifies weights related by widening and supports distances comparing widths, following Levin et al. [2026].The embeddings compose and respect hidden-neuron permutations; compatible normalized norms induce norms on the limit space.Source: §A.3 · pp. 20–22
  • Fan-in rescaling is motivated by the rescaled coordinates, where duplication widening repeats features.The limit-space construction also symmetrizes distances over hidden-neuron permutations.Source: §A.3 · pp. 21–22
  • The normalized spectral and ℓ1 norms are dominated by the normalized Frobenius norm, but width-independent norm domination depends on depth.For L ≥ 2, neither the normalized Frobenius nor ℓ1 norm is dominated by the normalized spectral norm with a width-independent constant.Source: §A.3 · pp. 23–24

B Theoretical properties of transferable GMNs

  • The appendix supplies proofs for duplication compatibility, block-wise widening compatibility, task-model misalignment, and Lipschitz continuity of the matrix-product GMN.It also establishes that the modified GMNs cannot additionally be made compatible with zero-padding widening.Source: B Theoretical properties of… · p. 24

B.1 Duplication-compatibility

  • After the three modifications, the simple GMN gives the same prediction on a network and its duplication widening.The proof tracks copied node and edge features through message passing and the layer-wise mean readout.Source: §B.1 · pp. 24–25

B.2 Compatibility with block-wise widening

  • The duplication-compatible GMN can disagree on block-wise-widened weights representing the same function, while the forward matrix-product GMN is compatible with block-wise widenings.The bidirectional matrix-product variant is compatible with the subclass satisfying the column-balance condition.Source: §B.2 · pp. 25–27

B.3 Task-model misalignment

This section analyzes continuity, duplication compatibility, and parameter sensitivity for task targets and graph metanetworks, including width-uniform bounds for matrix-product GMNs and bounds on how constants depend on input-network norms.

B.3 Task-model misalignment

  • The prediction target and duplication-compatible GMN are Lipschitz in different norms, creating a misalignment that can obstruct size generalization.The target is Lipschitz on bounded normalized-spectral-norm balls but need not be continuous in normalized ℓ1; the GMN is Lipschitz in normalized ℓ1, while its class can be discontinuous in normalized spectral norm.Source: §B.3 · pp. 29–30

B.4 Spectral continuity of matrix-product GMNs

  • Matrix-product GMNs are Lipschitz in the normalized spectral norm with constants independent of hidden widths.The bounds depend on the GMN parameters, input and output dimensions, depth, and the bounded weight-space radius.Source: §B.4 · pp. 31–32

B.5 Zero-padding compatibility

  • Forward matrix-product aggregation can preserve zero-padding compatibility during message passing, but bidirectional aggregation supports only a restricted zero-padding subclass.For the forward model, preprocessing and node updates must also preserve zeros for the relation to propagate.Source: §B.5 · p. 33
  • No nontrivial pooling rule is invariant to both zero-padding and duplication; mean pooling targets duplication compatibility at the cost of zero-padding compatibility.Zero-padding invariance forces sum pooling, while duplication invariance forces mean pooling.Source: §B.5 · p. 33

B.5.1 Lipschitz sensitivity of GMNs to their parameters.

  • MLPs with fixed architecture and Lipschitz activations are Lipschitz in their parameters and inputs uniformly over compact parameter sets.The stated bound applies to parameter and input pairs within the respective compact parameter set and input space.Source: §B.5.1 · pp. 33–34

B.5.1 Lipschitz sensitivity of GMNs to their parameters

  • For fixed-architecture Lipschitz MLP components, Proposition 12 gives a matrix-product GMN parameter-sensitivity bound with c_A,T independent of hidden widths.The bound applies for every width n and θ_n in V_n.Source: §B.5.1 · pp. 35, 37
  • K_t and C_t grow at most as c_t(1+R_n)^(t+1), while the final sensitivity constant B_n is at most c''(1+R_n)^(T+1).The constants c_t depend only on A, t, L, d_in, d_out, and d_h.Source: §B.5.1 · p. 37

C Implementations of transferable GMNs

This section describes adapting duplication-compatible and matrix-product approaches to existing graph metanetworks, including NG-GNN and ScaleGMN, to preserve width-transfer properties. It also describes a limitation of matrix-product models under positive-scaling symmetry.

C Implementations of transferable GMNs

  • Width-growing aggregations must be layer- and direction-wise and invariant to repeating each element; matrix-product variants restrict weight influence to forward or backward matrix products.For duplication compatibility, weights are initialized with fan-in rescaling; shared message and update maps can remain unchanged when their inputs already satisfy the required invariance.Source: C Implementations of transferable GMNs · p. 38

C.1 Modifications on the GMN of Lim et al. [2024]

  • For Lim et al. [2024], the duplication-compatible GMN separates forward and backward messages, uses mean aggregation, and replaces global sums with layer-wise means.Its matrix-product variant removes edge updates and restricts messages to be linear in the weight scalar and source state alone.Source: §C.1 · pp. 39–41

C.2 Modifications on the NG-GNN of Kofinas et al. [2024]

  • For Kofinas et al. [2024], duplication compatibility standardizes fan-in-rescaled weights, preserves positional and probe features, and restricts PNA to repeat-invariant aggregators with identity degree scaling.Dropout can break compatibility during training, but compatibility holds at evaluation when dropout is off.Source: §C.2 · p. 43
  • The NG-GNN matrix-product variant drops edge features and updates, uses raw fixed weight scalars, and can retain probe features computed by ordinary forward passes.The general recipe requires other components that read weight entries to be modified so weights enter only through matrix products.Source: C Implementations of transferable GMNs · p. 38, §C.2 · p. 43
  • NG-GNN is restricted to mean aggregation and uses layer normalization over fixed hidden dimension dh, or omits width-dependent normalization, so weights retain matrix-product form [2024].The original normalization uses width-growing global statistics that introduce nonlinear weight dependence.Source: §C.2 · pp. 43–44

C.3 Modifications on the ScaleGMN of Kalogeropoulos et al. [2024]

  • ScaleGMN [2024] enforces scaling equivariance through typed positional encodings and specialized message and node-update maps; its positive-scaling bidirectional form uses separate backward edge features.The implementation also uses residual connections and dropout.Source: §C.3 · p. 45
  • Duplication-compatible ScaleGMN [2024] uses fan-in-rescaled edge features, mean aggregation, and layer-wise mean readout; dropout preserves compatibility only at evaluation.The last-layer readout is already duplication-compatible, while the full graph readout must be modified.Source: §C.3 · p. 46
  • Matrix-product forward ScaleGMN [2024] is a less expressive specialization of duplication-compatible ScaleGMN, while positive-scaling bidirectional matrix-product ScaleGMN is impossible because reciprocal weights cannot be expressed as matrix products.This obstruction does not arise for sign symmetries.Source: §C.3 · p. 46

D Experiments

The experiments compare four base GMN architectures and test whether their modifications satisfy compatibility and continuity properties. Widening tests examine duplication and more general function-preserving transformations.

D Experiments

  • The experiments cross permutation or scale symmetry with forward or bidirectional direction, then add duplication-compatible and matrix-product variants where defined.The matrix-product bidirectional ScaleGMN is undefined, and row 4 is not run on the CNN-accuracy task because the bidirectional ScaleGMN has a numerical-instability problem.Source: D Experiments · pp. 46–47, Table 2 · p. 47

D.1 Verification of GMN compatibility and continuity

  • Numerical checks assess whether the GMN modifications satisfy compatibility and continuity on MLP and CNN weights.The checks use the maximum absolute difference across sampled networks, width multipliers, and output coordinates, with 10^-6 as the passing threshold.Source: §D.1 · p. 47

Duplication compatibility.

  • For MLP weights, the reported duplication-widening results align with the theoretical prediction of Proposition 1.The tests use width-32 MLPs widened with randomly selected hidden-layer multipliers.Source: Duplication compatibility · p. 47
  • Removing any one of the three modifications breaks duplication compatibility in the modified GMN.Table 3 reports maximum output differences between an input network and its duplication widening.Source: Table 3 · p. 48
  • Tests cover duplication and four other widening families; the reported GMN and ScaleGMN results align with the prediction for general block-wise widening.The families include two Kronecker widenings, balanced block-wise widening, and general block-wise widening.Source: Duplication compatibility · pp. 47–48

D.2 Size generalization on hand-crafted widenings

Hand-crafted, function-preserving widenings test whether GMNs trained at the smallest width retain their performance on wider inputs. Duplication-compatible models succeed under duplication widening, while only forward matrix-product models succeed under block-wise widening.

D.2 Size generalization on hand-crafted widenings

  • GMNs are trained only at the smallest input width and evaluated on wider networks generated by duplication or block-wise widening, both of which preserve the represented function.Widened test copies retain their source networks’ labels.Source: §D.2 · p. 48
  • Across test widths, duplication-compatible and matrix-product GMNs reproduce in-distribution performance under duplication widening.Baseline GMNs degrade sharply on wider inputs, whereas the modified models have flat error curves.Source: Figure 3 · p. 49, §D.2 · p. 49
  • Across test widths, only forward matrix-product GMNs reproduce in-distribution performance under block-wise widening.All other models, which are not compatible with this widening, degrade sharply on wider inputs.Source: §D.2 · p. 49, Figure 4 · p. 50

D.3 Size generalization on independently trained weights

  • The next experiments test size generalization on independently trained input networks with random initializations.Source: §D.3 · p. 50

D.3.1 Dataset generation

The datasets comprise independently trained INR and CNN weights at multiple widths, evaluated for classification or accuracy prediction. The section also tracks how input quality and prediction targets vary with width under SP and µP.

INR datasets

  • MNIST and FMNIST INRs are tested up to width 1024, or 42.7× and 32× their respective training widths.Image identities are partitioned across widths before fitting, so no image contributes weights to more than one width.Source: §D.3.1 · p. 50
  • Under SP, the fitting learning rate is retuned at each width; under µP, a value selected at base width is held fixed across widths.The µP and SP fits coincide in distribution at training width and differ only out of distribution.Source: §D.3.1 · p. 51
  • Under SP, INR quality peaks a few widths past training and then degrades; under µP it stays essentially flat with a low failure rate.Some of the accuracy decrease under SP may reflect declining input quality, but the considerable decrease cannot be explained by input quality alone.Source: §D.3.1 · p. 51

CNN datasets

  • CNN accuracy-prediction datasets use grayscale CIFAR-10 and SVHN networks, with test widths from c = 16 to c = 512 (32×).Networks have three convolutional layers, global average pooling, and a linear classifier; µP and SP networks coincide at training width c = 16.Source: §D.3.1 · p. 51
  • Prediction-target drift is measured by training-range coverage and median-accuracy shift; it is clearer under SP than under µP.Wider networks are expected to fit training data better, so accuracy may increase with width.Source: §D.3.1 · p. 51

D.3.2 Evaluation and results

On CIFAR-10 accuracy prediction, µP matrix-product variants perform reasonably well across all four criteria, while the forward baselines fail under SP. The theory’s guarantees are framed around a fixed training rule across widths, with further practical constraints on hyperparameters and training duration.

D.3.2 Evaluation and results

  • Because accuracy targets drift with width, four criteria separate value error (MAE, R2), affine-recalibrated fit (R2_recal), and rank ordering (Kendall’s τb).Source: §D.3.2 · p. 52
  • At c=512 under µP, forward ScaleGMN shifts from R2 = −1.450 (R2_recal = +0.682; τb = 0.761) to duplication-compatible R2 = +0.714 (τb = 0.756), chiefly tracking accuracy drift.Source: §D.3.2 · p. 52
  • Under µP, matrix-product GMN and ScaleGMN perform reasonably well on all four criteria, while the forward baseline and duplication-compatible GMN perform badly under SP.Source: §D.3.2 · p. 52, Figure 6 · p. 54

E Theory of size generalization on trained weights

  • The theory fixes a training rule across widths, including dataset, loss, architecture, optimizer, and parameterization, while initialization is independent across widths and hyperparameter and auxiliary-randomness laws are width-invariant.Source: E Theory of size generalization on… · p. 56
  • Width-independent hyperparameter laws exclude width-dependent learning rates, though optimal rates are observed to stay close across widths under µP [Yang et al., 2021].Source: E Theory of size generalization on… · p. 56
  • The shared hyperparameter law also prevents training epochs from growing with width, despite larger models typically taking longer to converge.Source: E Theory of size generalization on… · p. 56
  • The independence assumptions also exclude validation-selected stopping or restarting after a diverged run, and the authors do not expect guarantees when optimizer or data pipeline changes across widths.Source: E Theory of size generalization on… · p. 56
  • The analysis uses Tensor Programs to study limiting trained-weight distributions and states two size-generalization guarantees, in Theorems 4 and 5.Source: E Theory of size generalization on… · p. 56

E.1 Preliminary: Tensor Programs and µP

Tensor Programs provide a framework for characterizing infinite-width limits of neural-network computations. For one-layer MLPs, the appendix explains why µP can make trained networks at different widths agree, while standard parameterization (SP) does not yield the same initialization behavior.

E.1 Preliminary: Tensor Programs and µP

  • Tensor Programs express neural computations as sequences of random initialization, matrix multiplication, width averaging, and entrywise nonlinearities, enabling analysis of their infinite-width limits.The framework also covers forward and backward passes and a fixed finite number of weight updates when the architecture can be represented as a tensor program.Source: §E.1 · p. 57
  • The master theorem characterizes scalar convergence, and its corollary gives convergence of empirical vector measures to their limiting laws in Wasserstein distance.The stated framework uses a fixed dataset and training horizon, with randomness entering through initialization and propagating through downstream quantities.Source: §E.1 · pp. 56, 58, 59

E.2 Proof of Theorem 4

  • For one-layer MLPs, µP makes hidden-neuron vectors draws from a width-independent law at initialization and preserves a width-independent limiting law at each fixed training step.The appendix notes that this one-layer specialization is the classical mean-field scaling.Source: §E.2.1 · p. 61, §E.2 · p. 61

E.2.1 A high-level explanation

  • Under µP, convergence of hidden-neuron feature distributions and output biases to common limits yields closeness of trained one-layer networks across widths.The proof relates Wasserstein-2 convergence of feature distributions and convergence of output biases to closeness in the normalized spectral distance.Source: §E.2.1 · p. 62
  • Under SP, initialized networks of different widths are not expected to be close in normalized spectral distance; the appendix reports this initialization obstruction for every depth.The text says training is unlikely to remedy the discrepancy, based on numerical checks.Source: §E.2.1 · p. 62

E.2.2 Proofs of the 1-layer results

For one-layer MLPs trained under µP, the appendix proves convergence of trained weight distributions across widths and derives size generalization for duplication-compatible GMNs under stated regularity assumptions.

E.2.2 Proofs of the 1-layer results

  • As min{n, m} grows, trained weight distributions converge in the normalized spectral distance under the proposition’s assumptions.The proof first establishes conditional convergence and then removes the conditioning on the shared training randomness and hyperparameters.Source: §E.2.2 · pp. 63–64
  • The proof uses Tensor Programs to obtain limiting laws for hidden-neuron features and output biases, then relates these limits to distance between networks.Its setup assumes a finite training horizon and regularity conditions inherited from the Tensor Programs framework.Source: §E.2.2 · pp. 62–63
  • Under the corollary’s hypotheses, one-layer µP-trained inputs support size generalization for Lipschitz duplication-compatible GMNs, target functionals, and losses on compact input domains.The hypotheses include q>2 and the assumptions of Proposition 17.Source: §E.2.2 · p. 65

E.2.3 Obstruction on MLP with larger depths

The appendix identifies an initialization obstruction to extending its width-closeness result to deeper MLPs: hidden-to-hidden matrices prevent width-n and width-2n networks from being close in the stated distance.

E.2.3 Obstruction on MLP with larger depths

  • At initialization, µP-trained MLPs with at least two layers do not have weight distributions converging in W1 between widths n and 2n under the normalized spectral distance.The proof compares a rank-at-most-n duplicated matrix with a randomly initialized width-2n matrix whose (n+1)-st singular value is bounded below with high probability.Source: §E.2.3 · p. 66
  • Whether the same obstruction persists for trained weights at fixed positive time is unknown; existing µP theory provides no guarantee of closeness.The text notes that observed low-rank structure in some trained weights may remove the obstruction.Source: §E.2.3 · p. 66

E.3 Proof of Theorem 5

The proof uses tensor-program limits and a distance generated by the GMN class to control the size-generalization gap. Under the stated assumptions, it establishes uniform convergence on compact A and a size-generalization guarantee for matrix-product GMNs on μP-trained inputs.

E.3 Proof of Theorem 5

  • The d∞ route fails for MLPs with hidden-to-hidden layers because different-width weights remain separated even at initialization, although this does not imply size generalization fails.d∞ closeness is sufficient, but not necessary, for the size-generalization gap to vanish.Source: §E.3 · p. 67
  • dA(µn, µm) dominates the size-generalization gap, and the proof aims to show this GMN-generated distance vanishes as n, m →∞.The earlier d∞ route is sufficient but not necessary for the gap to vanish.Source: §E.3 · p. 67
  • Appending GMN, target, and loss evaluation to the training tensor program yields a deterministic infinite-width limit for each fixed Θ.Uniform integrability and Vitali convergence then give convergence of expected risks.Source: §E.3 · pp. 67, 70, 71
  • Exact uniform representability by the matrix-product GMN satisfies the target condition, but arbitrary functionals of the realized function are excluded.The excluded class may depend on infinitely many inputs or features not reducible to finitely many evaluations and width-independent nonlinearities; GMN expressivity across widths remains an open question.Source: §E.3 · pp. 69–70
  • Compactness of A and the common Lipschitz bound upgrade pointwise convergence a_n(Θ) → a∞(Θ) to uniform convergence.The limiting function inherits the Lipschitz bound by pointwise convergence.Source: §E.3 · p. 71
  • Corollary 21 gives a size-generalization guarantee for matrix-product GMNs under μP, assuming Assumption 3 and the remaining hypotheses of Proposition 20.The proof takes a supremum over m ≥ n and applies Proposition 20.Source: §E.3 · p. 72
Loading 2610.00420v1…