Source-linked AI summary
Graph Metanetworks for Processing Diverse Neural Architectures
Derek Lim, Haggai Maron, Marc T. Law, Jonathan Lorraine, James Lucas
TL;DR
Metanetworks process neural networks themselves, but existing equivariant designs are specialized and difficult to extend to diverse architectures. Graph Metanetworks encode networks as parameter graphs for graph neural network processing; the paper reports theoretical guarantees and better performance than existing baselines across evaluated tasks.
Problem
Existing equivariant metanets require specialized, architecture-specific layers and do not directly handle many standard modules or varying network architectures.
Source: §1 · pp. 1–2Method
Graph Metanetworks represent input networks as parameter graphs and process those graphs with graph neural networks.
Source: §1 · p. 2, §2.1.2 · p. 4Results
GMNs outperform existing metanetwork baselines across all evaluated tasks and process diverse neural architectures.
Source: §6 · p. 10Takeaways & Limitations
The approach extends metanetwork processing to architectures including Transformers, residual networks, normalization layers, and group-equivariant networks.
Source: §1 · p. 2The authors have not tested GMNs at the billion-parameter scale, and they provide no formal constraints on parameter-graph design.
Source: §6 · p. 10
Abstract
from arXiv · showhide
Neural networks efficiently encode learned information within their parameters. Consequently, many tasks can be unified by treating neural networks themselves as input data. When doing so, recent studies demonstrated the importance of accounting for the symmetries and geometry of parameter spaces. However, those works developed architectures tailored to specific networks such as MLPs and CNNs without normalization layers, and generalizing such architectures to other types of networks can be challenging. In this work, we overcome these challenges by building new metanetworks - neural networks that take weights from other neural networks as input. Put simply, we carefully build graphs representing the input neural networks and process the graphs using graph neural networks. Our approach, Graph Metanetworks (GMNs), generalizes to neural architectures where competing methods struggle, such as multi-head attention layers, normalization layers, convolutional layers, ResNet blocks, and group-equivariant linear layers. We prove that GMNs are expressive and equivariant to parameter permutation symmetries that leave the input neural network functions unchanged. We validate the effectiveness of our method on several metanetwork tasks over diverse neural network architectures.
1 INTRODUCTION
Graph Metanetworks represent neural networks as parameter graphs and process them with graph neural networks, aiming to respect parameter symmetries across diverse architectures. The paper reports theoretical guarantees and strong performance across evaluated metanet tasks.
1 INTRODUCTION
- Metanetworks process neural networks themselves to predict their properties, generate or optimize them, and transform them.The paper uses “metanets” as a shorthand for metanetworks.Source: §1 · p. 1
- Flattening network parameters discards input structure, while ignoring hidden-unit permutation symmetries can greatly degrade metanet performance.Permuting hidden-layer neurons in an MLP can leave its output unchanged.Source: §1 · p. 1
- Earlier equivariant metanets rely on specialized layers and struggle with normalization, residual connections, attention blocks, and varying input architectures.Designing such layers requires analyzing each architecture’s symmetries; existing methods were limited to simpler networks.Source: §1 · pp. 1–2
- GMNs encode each neural-network parameter as an edge in a compact graph and use standard graph-learning methods to process the representation.The parameter graph is designed to handle parameter-sharing layers such as convolutions and attention without scaling with activation count.Source: §1 · p. 2
- GMNs are theoretically equivariant to input-network permutation symmetries and can express the forward pass of any input feedforward neural network.The paper formulates these symmetries through neural graph automorphisms and shows that computation-graph metanets are at least as expressive as prior methods.Source: §1 · p. 2
- GMNs solve a variety of metanet tasks with diverse neural architectures, achieving strong quantitative performance across all tasks explored; they are at least as expressive as past equivariant metanets and easier to implement.Evaluated image classifiers included 2D CNNs, 1D CNNs, DeepSets, ResNets, and Vision Transformers.Source: §1 · p. 2
2 GRAPH AUTOMORPHISM-BASED METANETS
Graph metanets represent neural networks as graphs and process them with GNNs equivariant to parameter symmetries. The section develops parameter-graph constructions for diverse layers and explains the automorphisms and GNN operations that support this approach.
2.1.1 COMPUTATION GRAPHS
- Computation graphs represent neurons as nodes and weights as edges, but weight sharing can create many graph edges for each parameter.Figure 2’s small convolution example has 16 edges for 4 parameters; a 32 × 32 input can use the same four parameters across 1024 activations.Source: §2.1.1 · p. 3, Figure 2 · p. 3
2.1.2 PARAMETER GRAPHS
- Parameter graphs assign each network parameter a single edge, aiming for efficient representations that still capture the correct parameter permutation symmetries.The graphs are built from modular layer subgraphs that can be stitched together.Source: §2.1.2 · p. 4
- For convolution layers, parameter graphs use channel nodes and parallel edges for kernel positions, with positional encodings identifying each weight’s spatial location.This representation contains one edge per parameter and captures parameter permutation symmetries as graph automorphisms.Source: §2.1.2 · p. 4
- Attention-layer graphs use feature-dimension nodes and separate edges for query, key, value, and output maps; head identity can be encoded as a node feature.The construction is designed to represent each shared parameter with a single edge.Source: §2.1.2 · pp. 4–5
- Residual connections are represented by parameter-free edges, while bias nodes connect to neurons through bias-parameter edges.The bias-node design supports message-passing-like operations among bias representations.Source: §2.1.2 · pp. 4–5
2.2 NEURAL DAG AUTOMORPHISMS
- Neural DAG automorphisms preserve adjacency, node types, and weight-sharing constraints, and induce parameter permutations that leave the network function unchanged.Graph metanets are equivariant to these induced permutations because GNNs are equivariant to graph permutation symmetries.Source: §2.2 · p. 5
2.3 FORMULATING METANETS AS GNNS
- GMNs use GNNs that update edge representations, predicting per-parameter outputs from edges or network-level outputs by pooling edge features.The empirical setup does not use global features, though the expressive-power results rely on them.Source: §2.3 · p. 6
3 EXPRESSIVE POWER OF GRAPH METANETS (GMNS)
The section establishes the expressive power of graph metanets through simulation results. GMNs can express prior equivariant metanets on MLPs and the forward pass of any feedforward network represented by a computation graph, with the latter result limited to computation-graph inputs.
4 EXPRESSIVE POWER OF GRAPH METANETS (GMNS)
- On MLP inputs, GMNs can express StatNN and NP-NFN, making them at least as expressive as these existing equivariant metanets.The paper attributes StatNN’s operations to permutation-invariant statistics and says global features can capture NP-NFN’s global operations.Source: §3 · p. 6
- On computation-graph inputs, GMNs can express the forward pass of any input feedforward neural network.The formal result is limited to GMNs operating on computation graphs; the authors leave a proof for more general parameter graphs to future work.Source: §3 · pp. 6–7
- The forward-pass expressivity result applies only to GMNs operating on computation graphs, and its extension to more general parameter graphs remains unproved.The authors leave a formal proof for those graphs to future work.Source: §3 · p. 7
4 RELATED WORK
Prior metanets use a range of designs, including specialized equivariant layers, statistics, and sequence models. The related-work discussion contrasts those approaches with graph metanets’ parameter-level graph representation and theoretical guarantees.
4 RELATED WORK
- DWSNets and NFN derive equivariant linear layers for simple MLPs, while NFTs use specialized parameter-equivariant attention and layer-number features.NFN also handles CNNs.Source: §4 · p. 7
- Other metanets use statistics, flattened-weight convolutions, hierarchical LSTMs, per-parameter MLPs, or set functions over chunks.Methods for implicit neural representations may require special training or access to input-network forward passes, unlike graph metanets’ stated ability to process independently trained diverse networks without those requirements.Source: §4 · p. 7
- Compared with Zhang et al.’s computation-graph approach for MLPs, graph metanets use parameter graphs that scale efficiently with parameter sharing and offer theoretical guarantees.Other graph methods often use coarser representations, such as treating a whole convolutional layer as one node, rather than associating each parameter with an edge.Source: §4 · p. 7
- The framework can use GNN variants that learn edge representations, including message-passing, graph-equivariant, and graph-transformer architectures.These are relevant because the neural-network parameters are represented as graph edges.Source: §4 · p. 7
5 EXPERIMENTS
Across three metanet evaluations, GMNs predict CIFAR-10 classifier accuracy across varied architectures and settings, outperforming baselines; they also perform strongly on INR editing and representation tasks.
5.1 PREDICTING ACCURACY FOR VARYING ARCHITECTURES
- The task predicts each CIFAR-10 classifier’s test accuracy across six settings, including reduced training data and OOD tests on larger-hidden-dimension networks.GMNs use message passing and mean-pool edge representations for invariant predictions.Source: §5.1 · pp. 7–9
- GMNs outperform both baselines in all six CIFAR-10 accuracy-prediction settings, with larger benefits in low-data and OOD regimes.The table covers 15,000-network training, 10% of that set, and tests on higher-hidden-dimension networks after training on lower-dimensional ones.Source: Table 1 · p. 8, §5.1 · p. 9
- Each dataset contains about 30,000 CIFAR-10 classifiers: varying 2D CNNs, or diverse 1D CNNs, DeepSets, ResNets, and Vision Transformers.The varying CNNs differ in hidden dimension, convolution and fully connected layer counts, and BatchNorm or GroupNorm use.Source: §5.1 · p. 8
- Test accuracies span 10%–77.5% for Varying CNNs and 7.5%–87.8% for Diverse Architectures, whose parameter counts reach 87,658.Varying CNN parameter counts range from 970 to 21,165.Source: §5.1 · p. 8
5.2 EDITING 2D INRS
- For 2D INR editing, GMNs outperform all methods on Contrast and are beaten only by NFT on Dilate.The comparison reports test MSE, where lower is better.Source: §5.2 · p. 9, Table 2 · p. 9
5.3 SELF-SUPERVISED LEARNING WITH INRS
- On self-supervised neural-network representations, GMNs outperform all baselines in the linear-regression evaluation task.The table reports test MSE for a linear regressor on learned representations.Source: Table 3 · p. 9, §5.3 · p. 10
6 CONCLUSION
The paper presents GMNs as a graph-based approach with theoretical and empirical benefits, while identifying scale and symmetry coverage as open boundaries.
6 CONCLUSION
- GMNs are permutation-equivariant and provably expressive, and the paper reports outperforming metanetwork baselines across all evaluated tasks.The approach processes diverse architectures, including layers used in state-of-the-art models.Source: §6 · p. 10
- The billion-parameter regime remains untested, and GMNs account for permutation-based parameter symmetries but not scaling-induced ReLU symmetries.The authors note that scalable GNN methods may help with scale, but have not tried GMNs at that scale.Source: §6 · p. 10
- Future work could improve the basic message-passing GNN, extend theory to parameter graphs, and apply GMNs to pruning, learned optimization, or finetuning.The paper also identifies state-of-the-art networks as potential objects for processing and analysis.Source: §6 · p. 10
A MORE GRAPH CONSTRUCTION DETAILS
The section builds parameter graphs from layer-specific subgraphs, with features chosen to represent parameters and architecture while preserving neural DAG symmetries. It also describes computation DAGs and the automorphisms used to formalize those symmetries.
A.1 CODE FOR CONSTRUCTING PARAMETER GRAPHS
- The construction procedure converts supported PyTorch layers into parameter subgraphs and stitches them together; it can also reconstruct a module from its architecture and graph edge weights.The section also illustrates an S_n-equivariant linear layer and additional isolated linear, normalization, and spatial-grid subgraphs.Source: A MORE GRAPH CONSTRUCTION DETAILS · p. 15, Figure 6 · p. 16, §A.1 · pp. 17–18
A MORE GRAPH CONSTRUCTION DETAILS
- Parameter subgraphs encode linear, convolutional, and group-equivariant layers with layer/type features and parameter values; convolutional edges additionally encode kernel position.For group-equivariant layers, each input-output channel pair is connected by d_B edges, where d_B is the dimension of the equivariant-map space.Source: A MORE GRAPH CONSTRUCTION DETAILS · pp. 15–16
- Normalization layers use learned mean and variance nodes linked to normalized neurons, while residual edges carry a constant rather than parameter values.Node and edge features identify the normalization type and layer; residual edges are labeled as residual connections.Source: A MORE GRAPH CONSTRUCTION DETAILS · p. 17
B.1 DAGS AND COMPUTATION GRAPHS
- A computation DAG defines activations recursively from input values, bias nodes, weighted incoming activations, and a nonlinearity; its output-node activations define the network function.The DAG is assumed connected, and parameter sharing can assign one parameter to multiple edges.Source: §B.1 · p. 18
B.2 NEURAL DAG AUTOMORPHISMS
- A neural DAG automorphism preserves edges, fixes input, output, and bias nodes, and preserves parameter-sharing constraints.These conditions define the graph permutations used to represent parameter symmetries.Source: §B.2 · p. 18
B.3 NEURAL DAG AUTOMORPHISMS GENERALIZE KNOWN PARAMETER SYMMETRIES
Neural DAG automorphisms capture familiar hidden-unit and channel permutations, and these transformations leave the represented network function unchanged. Extending the function-preservation proof from computation graphs to parameter graphs is presented as a conjecture.
B.3 NEURAL DAG AUTOMORPHISMS GENERALIZE KNOWN PARAMETER SYMMETRIES
- For MLPs, neural DAG automorphisms are exactly permutations of hidden neurons within each layer; for convolution-only CNNs, layer-preserving hidden-channel permutations are automorphisms.The CNN result establishes this direction but does not establish the converse.Source: §B.3 · pp. 18–19
B.4 NEURAL DAG AUTOMORPHISMS PRESERVE NETWORK FUNCTIONS
- Every neural DAG automorphism of a computation graph leaves the neural network function unchanged.The proof establishes that activations under permuted parameters are permuted versions of the original activations, while output nodes remain fixed.Source: §B.4 · pp. 19–20
B.5 FROM COMPUTATION GRAPHS TO PARAMETER GRAPHS
- Extending the function-preservation result to parameter-graph automorphisms is conjectured, with a locally-injective homomorphism proposed as a route to lifting automorphisms.The authors identify parameter-graph multigraphs as a difficulty for this proof extension.Source: §B.5 · p. 20
B.6 EQUIVARIANCE OF GRAPH METANETWORKS TO NEURAL DAG AUTOMORPHISMS
Graph metanetworks are permutation equivariant, and invariant architecture features restrict that behavior to neural DAG automorphisms. The theory section also frames expressive-power proofs through approximation of component functions.
B.6 EQUIVARIANCE OF GRAPH METANETWORKS TO NEURAL DAG AUTOMORPHISMS
- Graph metanetworks are equivariant to parameter permutations induced by neural DAG automorphisms.The result applies the permutation equivariance of graph neural networks to architecture graph symmetries.Source: §B.6 · pp. 20–21
- Additional node and edge features invariant to neural DAG automorphisms prevent the intended equivariance from extending indiscriminately to other graph permutations.The section notes that extra equivariances would reduce expressivity and states that the same argument applies to output edge features.Source: §B.6 · p. 21
C THEORY: EXPRESSIVE POWER OF GRAPH METANETWORKS
- If each component function can be universally approximated, their composition can also be universally approximated.This approximation principle supports the section’s strategy for proving expressive-power results for compositions of metanet operations.Source: C THEORY: EXPRESSIVE POWER OF GRAPH… · p. 21
C.1 PROOF: SIMULATING EXISTING METANETWORKS
The theory section shows that Graph Metanetworks can express StatNN and NP-NFN on MLP inputs, and explains how bias nodes support expressive updates. These results concern the stated MLP setting and the described bias-node constructions.
C.1.1 SIMULATING STATNN
- Graph Metanetworks can express StatNN on MLP input networks by aggregating continuous permutation-invariant weight and bias statistics before applying an MLP.The statistics described include the mean, variance, and percentiles at q ∈ {0, 25, 50, 75, 100}.Source: §C.1.1 · p. 22
C.1.2 SIMULATING NP-NFN
- Graph Metanetworks can express NP-NFN on MLP input networks by reproducing its updates through global, node, and edge features.The construction describes a graph with invariant node and edge features and shows the resulting edge representation matches the NP-NFN update.Source: §C.1.2 · pp. 23–24
C.1.3 EXPRESSIVE POWER WITH BIAS NODES
- Bias nodes let biases in the same layer communicate in one message-passing layer without a global feature, unlike encoding them as self-loops or node features.With a global feature, the described NP-NFN case allows communication between biases across layers.Source: §C.1.3 · p. 25
C.2 PROOF: SIMULATING FORWARD PASSES
The proposition states that computation-graph GMNs can express the forward pass of any input feedforward neural network. The proof uses induction on graph depth to compute each node’s activation.
C.2 PROOF: SIMULATING FORWARD PASSES
- Computation-graph GMNs can express the forward pass of any input feedforward neural network.Source: §C.2 · p. 25
- The proof inducts on graph depth, showing that a GMN of depth l computes activations for nodes whose maximum input-path distance is at most l.Source: §C.2 · p. 25
- At depth zero, input-node activations equal their inputs, while a node-update MLP can output 1 for bias nodes.Source: §C.2 · p. 25
- The induction handles output nodes similarly, replacing the hidden-node activation function with the identity map.Source: §C.2 · p. 26
D.1.1 CLASSIFIER TRAINING DETAILS
The accuracy experiments use two CIFAR-10 classifier datasets of about 30,000 networks each, with separate dataset sizes for diverse architectures and varying CNNs.
D.1.1 CLASSIFIER TRAINING DETAILS
- The accuracy experiments train two CIFAR-10 image-classification datasets of about 30,000 neural networks each.Source: §D.1.1 · p. 26
- For the Diverse Architectures dataset, about 6,000 networks are trained for each of five network types.Source: §D.1.1 · p. 26
- For the Varying CNNs dataset, about 30,000 2D-CNNs are trained.Source: §D.1.1 · p. 26
D.1.2 METANETWORK TRAINING DETAILS
The training details cover accuracy prediction, image editing with 2D INRs, and self-supervised learning with INRs. For small CNN accuracy prediction, GMNs outperform the other metanetworks in the reported comparison.
D.1.2 METANETWORK TRAINING DETAILS
- For OOD accuracy experiments, training and validation use the lowest hidden dimension for each architecture, and testing uses networks with higher hidden dimensions.The held-out dimensions include CNNs, DeepSets, ResNets, and ViTs.Source: §D.1.2 · p. 26
- GMNs outperform all other metanetworks on small-CNN accuracy prediction in the reported comparison.Source: Table 6 · p. 28, §D.1.2 · p. 28
- Figure 8 plots actual against predicted test accuracy for OOD networks, separating validation-split inputs in the top row from test-split inputs in the bottom row.Source: Figure 8 · p. 28
- Metanets are sized to about 750,000 trainable parameters, tuned by validation R2 across five learning rates, and reported over five seeds.Training uses Adam with binary cross-entropy and a sigmoid output to constrain predicted accuracy to [0, 1].Source: §D.1.2 · p. 28
D.2 EDITING 2D INRS DETAILS
- For 2D INR editing, updated parameters are θ plus a learned per-parameter scale times the metanet output, with scales initialized to 0.01.The metanets train for 50,000 iterations with batch size 32; the message-passing model has four layers.Source: §D.2 · pp. 28–29
D.3 SELF-SUPERVISED LEARNING WITH INRS DETAILS
- For self-supervised INR learning, the graph metanet uses three GNN layers and hidden dimension 76, with contrastive training on noisy and masked weights.Results are averaged over three random seeds.Source: §D.3 · p. 29