PaperlayerView exact revision on arXiv

Source-linked AI summary

CrossGMN: Graph Metanetworks for Cross-Architecture Weight-Space Transformations

Adir Dayan, Yam Eitan, Haggai Maron

arXiv:2610.01649v1cs.LG

TL;DR

Existing equivariant weight-space architectures have mainly addressed transformations within one architecture, while practical transformations often cross architectures. CrossGMN refines an initialized target using a trained source through symmetry-preserving graph message passing. It is universal under stated conditions and accelerates distillation across diverse settings, with speedups up to 8.89× on MNIST INRs.

  • Problem

    Existing equivariant weight-space architectures primarily address mappings within one weight space, whereas practical transformations map between architectures with distinct permutation symmetries.

    Source: §1 · p. 1

  • Method

    CrossGMN jointly processes a trained source and initialized target as parameter graphs, refining the target while remaining source-invariant and target-equivariant.

    Source: §1 · p. 2

  • Results

    Across INRs and classifiers, CrossGMN reduces subsequent distillation steps and supports transfer and heterogeneous-source compression, with up to 8.89× speedup on MNIST INRs.

    Source: §1 · p. 2, §7 · p. 9

  • Takeaways & Limitations

    The results demonstrate learned weight-space compression that can be reused across source architectures and datasets within the evaluated settings.

    Source: §1 · p. 2, §7 · p. 9

    The current setting requires access to source parameters and compatible input and output spaces, and is evaluated only at moderate scale, up to roughly 200k parameters.

    Source: §7 · p. 9

Abstract

from arXiv · show

Weight-space networks operate directly on parameters of other neural networks, enabling tasks such as predicting model properties, editing trained models, and generating weights. Weight-space symmetries such as neuron permutations make equivariance a key design principle. However, existing equivariant weight-space architectures have primarily been studied for transformations that preserve the network architecture. In contrast, many practical transformations, including model compression and upscaling, map a trained source network into a target network with a different architecture. In this setting, the source and target permutation symmetries act on different parameter spaces, making equivariance less straightforward to formulate. Our key idea for addressing this mismatch is to reformulate cross-architecture operators with two inputs: a trained source network and an initialization of the target network. This lets us define equivariant cross-architecture operators that refine the initialization of the target network using information from the source network, while being invariant to source-network permutations and equivariant to target-network permutations. Based on this formulation, we introduce CrossGMN, a graph metanetwork that jointly processes both networks through symmetry-preserving cross-network message passing. We prove CrossGMN is universal for continuous cross-architecture operators on compact sets under a general-position assumption. We evaluate CrossGMN for model compression, predicting a smaller network's parameters to accelerate subsequent knowledge distillation. Across 2-D and 3-D INRs and image classification with MLPs, CNNs, and Vision Transformers, CrossGMN speeds up distillation by up to 8.89x, transfers across datasets without retraining (3.78x), and a single model can accelerate compression from heterogeneous source architectures into a common target architecture.

1 Introduction

Practical transformations often map weights between different architectures, but existing equivariant weight-space models primarily handle mappings within one architecture. CrossGMN formulates these transformations using a trained source and initialized target, then learns symmetry-respecting refinements with theoretical and empirical support.

Motivation

  • Cross-architecture transformations include compression, knowledge transfer, and upscaling, but existing equivariant weight-space architectures primarily map within the same weight space.The source and target have distinct permutation symmetries acting on different parameter spaces, making equivariance unclear.Source: §1 · pp. 1–2

Formulation

  • Treating the trained source and initialized target as two inputs makes the operator invariant to source permutations and equivariant to target permutations.The target initialization specifies its architecture and provides parameters to refine.Source: §1 · p. 2

Approach and theory

  • CrossGMN jointly processes source and target parameter graphs with cross-network message passing and is universal for continuous operators on compact sets under a general-position assumption.The result also extends to finite collections of source architectures mapped into a common target weight space.Source: §1 · p. 2

Evaluation

  • Up to 8.89× speedup on MNIST INRs and 6.85× on ModelNet INRs, with up to 3.78× transfer to FashionMNIST without retraining.A single model also compresses heterogeneous CIFAR-10 MLP teachers into a common student architecture and generalizes to unseen teacher architectures.Source: §1 · p. 2

2 Preliminaries

The preliminaries define MLP parameters and their permutation symmetries, while focusing the theoretical analysis on MLPs and evaluating across a broader set of architectures.

Weight spaces

  • An MLP is specified by its layer widths and activation, with weight matrices and bias vectors forming its parameters.Its output is computed by composing affine layers with the activation function.Source: §2 · pp. 2–3

Scope

  • Source and target networks may differ in hidden widths and depths, but are assumed to share input and output dimensions.The theoretical analysis focuses on fully connected MLPs, while experiments also include CNNs and Transformers.Source: §2 · pp. 2–3

Permutation symmetries

  • Permuting hidden neurons changes an MLP's parameterization while preserving its realized function, motivating invariant and equivariant weight-space maps.Equivariance transforms outputs consistently with the input permutation; invariance leaves outputs unchanged.Source: §2 · p. 3

3 Cross-Architecture Weight-Space Transformations

Directly mapping source parameters to target parameters raises symmetry and continuity difficulties. The paper instead defines a two-input operator that refines a target initialization, a formulation that includes finite-step distillation under stated continuity conditions.

Direct mapping

  • A direct source-to-target map must be source-invariant while selecting among equivalent target parameterizations, and natural maps may be discontinuous.The passage gives model compression as a setting where mapping each larger network to an optimal smaller one can be discontinuous.Source: §3 · p. 3

Operator formulation

  • The alternative is a cross-architecture operator from source parameters and target initialization to refined target parameters.The initialization specifies a concrete target parameterization and supplies parameters for refinement.Source: §3 · p. 3
  • The operator is invariant to source permutations and equivariant to target permutations.Applying a source permutation and a target permutation to the inputs yields the correspondingly permuted target output.Source: §3 · p. 3

Knowledge distillation

  • For fixed hyperparameters and minibatches, k steps of SGD or Adam define a continuous equivariant cross-architecture operator when the student loss gradient is jointly continuous.The paper identifies knowledge distillation as an instance of this operator class.Source: §3 · pp. 3–4

4 CrossGMN

CrossGMN represents source and target networks jointly as parameter graphs and applies message passing while preserving their independent permutation symmetries. Its universality result covers continuous operators on compact general-position sets when connectivity includes the strided policy.

4.1 Architecture

  • The architecture represents source parameters and target initialization as parameter graphs, joins them into a joint graph, and uses a graph metanetwork to produce target-side updates.The target-side outputs are extracted and arranged as an update in the target weight space.Source: §4.1 · pp. 4–5, Figure 1 · p. 5
  • Cross edges must be invariant under independent source and target permutations; between hidden layers, this requires connecting every source neuron to every target neuron.Shared input and output coordinates instead use positionwise joins.Source: §4.1 · p. 4
  • The choice of which source and target hidden layers to join is a cross-edge policy; all-to-all connectivity uses nm joins for networks with n and m hidden layers.The strided policy uses n joins by partitioning the deeper network's layers into near-equal blocks.Source: §4.1 · p. 4
  • The strided policy retains all-to-all's benefits with fewer cross edges, is covered by the universality result, and performs comparably to all-to-all in experiments.The paper reports the empirical comparison in Table 6.Source: §4.1 · p. 5
  • Joint-graph message passing preserves both symmetry actions, and extracting only target-side outputs yields source invariance and target equivariance.A fixed positive residual-scale hyperparameter is also specified for the architecture.Source: §4.1 · p. 5

4.2 Expressivity

  • CrossGMN is universal for continuous cross-architecture operators on compact general-position sets when its cross-edge connectivity includes the strided policy.The theorem's general-position assumption requires that no two neurons within a hidden layer have the same bias value.Source: §4.2 · p. 5
  • The universality theorem provides an approximation guarantee for the knowledge-distillation operator introduced earlier in the paper.The result relies on the fixed source and target architectures and compact general-position weight sets specified in the theorem.Source: §4.2 · p. 5
  • The universality result extends to finite collections of source architectures sharing a target architecture, provided the source architecture is identifiable from the joint graph.The paper also evaluates this heterogeneous-source setting.Source: §4.2 · p. 6

5 Experiments

The experiments evaluate CrossGMN as a way to accelerate knowledge distillation across model families, datasets, and compression settings. They report broad but setting-dependent speedups, test out-of-distribution transfer, and examine the role of teacher–student communication.

5 Experiments

  • Experiments measure whether CrossGMN-predicted student initializations accelerate subsequent distillation, using trained-model populations that include INRs, MLPs, ViTs, and CNNs.The evaluation also includes width, depth, and combined compression settings.Source: §5 · p. 6, §5.1 · p. 7, §5.2 · p. 8

5.2 CrossGMN distillation speedup across model families and datasets

  • 4.99× wall-clock speedup versus 5.20× in distillation steps for MNIST INR width compression; joint width–depth compression achieves 8.36× versus 8.89×.The width setting compresses [2, 32, 32, 1] to [2, 8, 8, 1]; the joint setting compresses to [2, 16, 1].Source: §5.2 · p. 8
  • CrossGMN consistently accelerates INR distillation, improves MLP, CNN, and ViT depth-reduction settings, but not ViT attention-dimension or head-count reductions, and yields positive speedups across heterogeneous architectures.The results support learning and amortizing model compression directly in weight space.Source: §5.2 · p. 8

5.3 Out-of-Distribution Generalization

  • Cross-dataset and unseen-architecture evaluations support transfer across datasets and provide evidence of partial generalization to unseen architectures.The MNIST INR-trained models are applied unchanged to FashionMNIST INRs of the same architectures.Source: §5.3 · p. 8, Table 4 · p. 9

5.4 Ablations: the role of cross edges

  • Removing cross edges or replacing them with pooled-teacher FiLM conditioning reduces distillation speedup, while strided connectivity performs comparably to all-to-all with fewer cross edges.The communication ablation uses MNIST INR width reduction.Source: §5.4 · p. 9

6 Related Work

Prior symmetry-aware weight-space architectures operate within a single weight space, while related parameter-generation methods produce weights rather than transforming a particular trained network into another weight space.

6 Related Work

  • Table 6 examines teacher–student communication for MNIST INR width reduction from [2, 32, 32, 1] to [2, 8, 8, 1].Source: Table 6 · p. 9
  • Earlier symmetry-aware weight-space networks and graph-based models operate within a single weight space, unlike transformations between distinct source and target weight spaces.The related work also covers extensions to additional symmetries.Source: §6 · p. 9
  • Related parameter-generation methods condition on a target architecture or learn distributions over trained-model populations, generating new weights rather than transforming a particular trained network.The passage cites Zhang et al. (2019), Knyazev et al. (2021), Liu et al. (2019), Schürholt et al. (2022a; 2024), Erkoc et al. (2023), Wang et al. (2024), and Soro et al. (2025).Source: §6 · p. 9

7 Conclusion

CrossGMN learns equivariant transformations between different neural-network weight spaces, refining a target initialization using a trained source; experiments report reduced distillation steps across INRs and classifiers. The conclusion also identifies access, compatibility, and scale boundaries, while related work distinguishes this operator approach from existing weight-space, generation, growth, and compression methods.

7 Conclusion

  • Across INRs and classifiers, CrossGMN reduces subsequent distillation steps for width and depth compression and supports reuse across source architectures, dataset transfer, and partial generalization to unseen sources.The authors conclude that equivariant models can learn architecture-changing transformations directly in weight space.Source: §7 · p. 9
  • The current setting requires source parameters and compatible input/output spaces, and evaluation is limited to moderate scale, up to roughly 200k parameters.Scaling may require sparse or blockwise joint graphs and more efficient cross-network communication.Source: §7 · p. 9

A Extended Related Work

  • Prior weight-space models and expressivity results address a single weight space, whereas CrossGMN studies operators between distinct source and target spaces with independent permutation actions.Source: A Extended Related Work · p. 16
  • Parameter-generation methods generate weights for an architecture or model population rather than transforming a particular trained source network into another weight space.Source: A Extended Related Work · p. 16
  • Unlike analytic growth rules or source-specific growth-map fitting, CrossGMN learns a shared operator across source networks and reuses it on unseen sources without fitting a new operator.Source: A Extended Related Work · p. 16
  • CrossGMN refines a structured-pruning student initialization using the teacher, then leaves distillation to proceed toward that same teacher.The refinement is learned across a population of teachers to accelerate the subsequent optimization.Source: A Extended Related Work · p. 16

B Notation and Architecture-Specific Symmetries

This section defines parameter-graph representations and permutation actions for MLPs, CNNs, and Transformers, and explains how CrossGMN’s joint graph preserves independent source and target symmetries. It states the resulting source-permutation invariance and target-permutation equivariance, and establishes continuity and universality results for cross-architecture operators under stated conditions.

B.1 Notation

  • In the experiments, the target initialization is the anchor, and theoretical source/target terminology corresponds to experimental teacher/student terminology.The section writes the permutation action as g·v = ρA(g)v.Source: §B.1 · pp. 16–17, Table 7 · p. 17

B.2 Parameter graphs and induced permutation actions

  • Parameter graphs store each learnable scalar on a node or edge, with typed structural features and labels defining parameter locations and graph automorphisms.Matrix entries map to directed edges, convolutional kernels to parallel channel edges, and biases or normalization entries to nodes or auxiliary-node edges.Source: §B.2 · p. 17

B.2.1 Joint graphs and cross-edge policies

  • Joint graphs preserve independent source and target actions through layerwise joins of hidden-unit sets and positionwise joins of fixed, canonically indexed boundary nodes.Strided matching pairs target layers with nonempty source-layer blocks; equal-depth networks use shared-index matching.Source: §B.2 · p. 17, §B.2.1 · p. 18

B.2.2 Multilayer perceptrons

  • For MLPs, hidden-neuron permutations are precisely the standard parameter-graph automorphisms and preserve the network function.The MLP encoding uses one node per neuron and directed weight edges, with biases represented on nodes or through fixed auxiliary nodes.Source: §B.2 · p. 17, §B.2.2 · pp. 18–19

B.2.3 Convolutional networks

  • For CNNs, channel permutations preserve the realized function by relabeling channel nodes and incident kernel and bias edges.Each input/output channel pair is represented by parallel kernel edges, while spatial reuse does not add parameter edges.Source: §B.2.3 · pp. 19–20

B.2.4 Transformers

  • Transformer graph symmetries use a represented subgroup determined by structural labels, and the implementation does not claim equivariance to every function-preserving Transformer symmetry.The specified action permutes coordinates within heads, heads, and feed-forward units while fixing residual-stream coordinates; some other functional symmetries are not represented.Source: §B.2.4 · p. 21

B.2.5 Equivariance of the joint graph

  • For every metanetwork parameter choice, CrossGMN is equivariant to source and target permutation groups induced by neural graph automorphisms; for MLPs, these are the full hidden-neuron permutation groups.The proposition applies to the specified joint-graph cross-edge constructions for MLPs, CNNs, and Transformers.Source: §B.2.5 · p. 22
  • Disjoint source and target automorphisms preserve the joint graph, and GMN relabeling equivariance with target projection and residual addition yields source invariance and target equivariance.Graph metanetworks commute with relabelings (Lim et al., 2024, Appendix C.6).Source: §B.2.5 · pp. 22–23

C Theoretical Results

  • CrossGMN is universal for continuous source-invariant, target-equivariant operators on compact general-position sets, including finite collections of source architectures.The appendix also shows that mapping each larger network to an optimal smaller one can be discontinuous, while any fixed number of standard distillation steps defines a continuous cross-architecture operator.Source: C Theoretical Results · p. 23

C.1 Discontinuity of mapping a larger network to an optimal smaller one

For a width-2 ReLU teacher compressed to a width-1 student, an optimal student must switch between two components as the teacher varies. This makes a continuous parameter map that always selects an optimal student impossible, even in general position; refining a supplied initialization avoids requiring global optimality.

C.1 Discontinuity of mapping a larger network to an optimal smaller one

  • In the width-2-to-width-1 ReLU example, the optimal student switches between the teacher’s two spatially separated components as their coefficients cross.For t < 0 it realizes (1 − t)u; for t > 0 it realizes (1 + t)v; at t = 0 both are optimal.Source: §C.1 · pp. 23–25
  • No continuous map can select an optimal width-1 student for every teacher along the continuous path, because the one-sided optimal limits differ at t = 0.The discontinuity also rules out arbitrarily small uniform approximation of an optimal selection by a continuous map.Source: Figure 3 · p. 24, §C.1 · p. 25
  • The score Γ_t(h) measures squared-error reduction from fitting the best scalar multiple of shape h, so minimizing student error is equivalent to maximizing that score.Source: §C.1 · p. 24
  • The counterexample holds in general position: teacher hidden-layer biases remain distinct along the path, and the width-1 student satisfies the condition vacuously.The discontinuity therefore does not rely on coincident hidden biases.Source: §C.1 · p. 25
  • A cross-architecture operator can instead refine a supplied student initialization without requiring a globally optimal student.The paper states that a fixed number of distillation steps gives a continuous equivariant operator under its assumptions.Source: §C.1 · p. 25

C.2 Finite-step distillation is a cross-architecture operator

Finite-step distillation defines a cross-architecture operator when its updates use a shared continuous coordinatewise rule and the student-loss gradients satisfy the stated continuity assumptions. For every fixed step count, the resulting map is continuous and equivariant, and functionally equivalent teachers produce identical trajectories from the same initialization.

C.2 Finite-step distillation is a cross-architecture operator

  • Gradient descent, SGD, momentum, RMSProp, and Adam fit the shared coordinatewise update form under fixed schedules and the stated optimizer conditions.Source: §C.2 · pp. 25–26
  • For every fixed k, the student parameters after k distillation steps define a continuous cross-architecture operator that is equivariant to source and target permutations.The stated assumptions include jointly continuous student-loss gradients and a shared continuous coordinatewise update rule.Source: §C.2 · pp. 25–26
  • Functionally equivalent teachers give identical student gradients and therefore identical finite-step trajectories from the same initialization.Source: §C.2 · p. 26
  • The continuity argument proceeds by induction from the initialization, since each update composes jointly continuous gradients with a continuous coordinatewise rule.Source: §C.2 · p. 26

C.3 Universal approximation of cross-architecture operators

CrossGMN uniformly approximates continuous equivariant cross-architecture operators on compact general-position sets for fixed MLP source and target architectures. The proof factors an operator into source-invariant weights and target-equivariant maps, then realizes and combines those factors through joint-graph message passing.

C.3 Universal approximation of cross-architecture operators

  • The proof approximates any continuous source-invariant, target-equivariant operator by a finite mixture of source-invariant scalar functions and target-equivariant maps.Source: §C.3 · p. 27
  • CrossGMN realizes these mixtures by computing the factors on their respective parameter graphs, communicating source scalars across an edge, and multiplying them on the target side.The construction approximates the scalar and target-map factors uniformly before combining them.Source: §C.3 · pp. 27–29
  • General position excludes parameter sets with equal biases among two neurons in a hidden layer.Source: §C.3 · p. 27
  • The same approximation conclusion holds when further symmetry-preserving cross edges are added.Source: §C.3 · p. 28
  • On compact general-position sets for fixed source and target MLP architectures, CrossGMN uniformly approximates continuous cross-architecture operators to any prescribed error.The theorem uses joint graphs with positionwise boundary joins and strided hidden-layer joins.Source: §C.3 · p. 29

C.4 One CrossGMN for finitely many source architectures

One CrossGMN can approximate continuous operators for a finite family of source architectures with one parameter set, provided the family is architecture-identifiable. Its extension to CNN and Transformer weight spaces depends on fixed-architecture universality conditions.

C.4 One CrossGMN for finitely many source architectures

  • One CrossGMN with a single parameter set uniformly approximates each continuous source-invariant, target-equivariant operator in an identifiable finite source family on compact sets.For finite MLP families, graph depth, layer types, and layer cardinalities can identify the source architecture.Source: §C.4 · p. 30

C.5 Scope for CNN and Transformer weight spaces

  • For CNNs and Transformers, equivariance follows from their implemented parameter graphs, while universality requires fixed-architecture invariant and equivariant GMN universality to hold.The extension uses the neural graph automorphism groups defined in Appendix B.2.Source: §C.5 · p. 31

D.1 CrossGMN implementation

  • The experiments treat the teacher and target initialization as a joint graph, with a target-side readout predicting the residual refinement.Classifier and INR node and edge features encode parameter values, layer information, and source/target roles; classifier features may also include family and activation embeddings.Source: §D.1 · p. 31
  • The CrossGMN backbone is bidirectional SCALEGMN, chosen for better empirical performance than plain GMN in most experiments, although universality is established for plain GMN.The implementation uses skip connections, edge updates, dropout 0.1, and family-specific scale- or sign-equivariant message passing.Source: §D.1 · p. 31

D.2 Evaluation protocol

The protocol compares CrossGMN-refined anchors with the same teachers distilled under matched optimizers, minibatches, data order, and evaluation steps. It also specifies family-specific compression cohorts and reach–budget criteria.

D.2 Evaluation protocol

  • Matched teachers, optimizers, minibatches, data order, and evaluation steps isolate the benefit of refining an anchor with CrossGMN; speedup statistics include failures.A failed CrossGMN run receives speedup zero, and reported metrics include median speedup, acceleration and failure rates, plus wall-clock speedup including initialization.Source: §D.2 · pp. 32, 34, 35

D.3 MNIST and FashionMNIST INR compression

  • MNIST INR operators are evaluated on 1,000 held-out teachers, while MNIST-trained operators are applied to a 1,000-teacher FashionMNIST cohort without retraining.Both populations use SIREN teachers with architecture [2, 32, 32, 1].Source: §D.3 · p. 35

D.4 ModelNet40 3-D INR compression

  • ModelNet40 evaluates pre-fit 3-D SIREN distance-field teachers on an object-disjoint test split, with all settings assessed on 4,096 points.The teacher architecture is [3, 48, 48, 48, 48, 1].Source: §D.4 · pp. 35–36

D.5 CIFAR-10 MLP compression

  • The heterogeneous CIFAR-10 MLP study trains on four source architectures mapped to [3072, 8, 10] and tests generalization on two held-out source architectures.The unseen-architecture cohort contains 500 teachers.Source: §D.5 · p. 36

D.6 CIFAR-10 Vision Transformer compression

  • The CIFAR-10 ViT population contains 6,100 teachers, and its reduction anchors prune attention dimensions or heads or select blocks before combined head-and-depth reduction.The teacher architecture has four blocks and two heads; depth reduction retains the first and third blocks.Source: §D.6 · pp. 36–37

D.7 MNIST CNN compression

  • The MNIST CNN experiment reduces convolutional widths from [8, 6, 4] to [4, 3, 2], leaving the classifier head unchanged and retaining 41.5% of parameters.Evaluation uses a fixed cohort of 400 teachers.Source: §D.7 · p. 37

D.8 Ablations

  • Direct neuron-level cross edges outperform removing communication or replacing it with pooled FiLM, while strided and all-to-all cross-edge policies yield similar speedups.Other ablations indicate stronger dependence on training population, message-passing depth, and width than on ScaleGMN versus a parameter-matched GMN; these are single runs, not seed-variance estimates.Source: §D.8 · p. 37

D.9 Reach–budget plots

  • Reach–budget curves show the fraction of teachers meeting fixed teacher-relative criteria by distillation budget; thirteen of forty panels show both CrossGMN and its anchor eventually reaching the threshold for nearly all teachers.For ModelNet40, only fidelity reach is reported; an asterisk marks a fallback threshold when no threshold meets the selection conditions.Source: §D.9 · pp. 37–38, Figure 5 · p. 39

D.10 Complete results

Complete results show that settings with median speedup above one also have nonnegative median quality differences at the budget, with a stated exception for two ViT reductions. Classifier teacher populations vary in prediction agreement and error overlap.

D.10 Complete results

  • Every reported setting with median speedup above one has a nonnegative median quality difference at budget, but CrossGMN does not accelerate ViT attention-dimension or head-count reductions.For those two reductions, the structured anchors already provide strong initializations.Source: §D.10 · pp. 42–43, Table 17 · p. 42

D.11 Classifier teacher-population diversity

  • CIFAR-10 MLP teachers have lower agreement than ViTs (0.715 versus 0.792), while MNIST CNNs show high agreement (0.891) but low error overlap (0.210).Heterogeneous MLP and unseen-architecture populations have κ values of 0.602 and 0.574, respectively.Source: §D.11 · p. 43, Table 18 · p. 43

Related papers

Loading 2610.01649v1…