PaperlayerView exact revision on arXiv

Source-linked AI summary

Graph Neural Networks for Learning Equivariant Representations of Neural Networks

Miltiadis Kofinas, Boris Knyazev, Yan Zhang, Yunlu Chen, Gertjan J. Burghouts, Efstratios Gavves, Cees G. M. Snoek, David W. Zhang

arXiv:2403.12143v3cs.LGcs.AIstat.ML

TL;DR

Existing approaches to processing neural-network parameters can overlook neuron-permutation symmetry or rely on architecture-specific equivariant designs. This paper represents networks as neural graphs for graph neural networks and transformers, and reports strong results across diverse tasks, outperforming state-of-the-art approaches.

  • Problem

    Methods for processing neural-network parameters either overlook neuron-permutation symmetry or require architecture-specific weight sharing while assuming a fixed architecture.

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

  • Method

    The method represents neurons as graph nodes and connections as edges, with biases and weights as features, then applies graph neural networks and transformers to these neural graphs.

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

  • Results

    Across a wide range of tasks, the proposed method outperforms state-of-the-art approaches by a large margin.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Neural graphs provide a flexible framework for processing networks with varying architectures and for adding domain-specific inputs such as gradients.

    Source: §6 · p. 9

    Evaluation covered only MLPs and CNNs, transformers were demonstrated theoretically, and strong INR performance was confined to 2D images.

    Source: §6 · p. 9

Abstract

from arXiv · show

Neural networks that process the parameters of other neural networks find applications in domains as diverse as classifying implicit neural representations, generating neural network weights, and predicting generalization errors. However, existing approaches either overlook the inherent permutation symmetry in the neural network or rely on intricate weight-sharing patterns to achieve equivariance, while ignoring the impact of the network architecture itself. In this work, we propose to represent neural networks as computational graphs of parameters, which allows us to harness powerful graph neural networks and transformers that preserve permutation symmetry. Consequently, our approach enables a single model to encode neural computational graphs with diverse architectures. We showcase the effectiveness of our method on a wide range of tasks, including classification and editing of implicit neural representations, predicting generalization performance, and learning to optimize, while consistently outperforming state-of-the-art methods. The source code is open-sourced at https://github.com/mkofinas/neural-graphs.

1 INTRODUCTION

Processing neural-network parameters could support tasks such as predicting generalization error and classifying implicit representations, but neuron permutations make flattened parameter inputs symmetry-sensitive. The paper represents networks as neural graphs so graph neural networks and transformers can process varied architectures while preserving permutation symmetry.

1 INTRODUCTION

  • Neuron reordering can preserve a network’s function, but flattened parameter inputs overlook this symmetry and may produce different predictions for equivalent networks.The motivation includes predicting generalization error, generating weights, and classifying or generating implicit neural representations.Source: §1 · p. 1
  • Existing equivariant weight-sharing patterns require manual adaptation to each architecture and limit a model to one fixed architecture.The proposed neural-graph approach instead targets heterogeneous computational graphs, including differences in depth, width, nonlinearities, and residual connections.Source: §1 · p. 1
  • Neural graphs map biases to node features and weights to edge features connecting adjacent layers.The paper presents this representation as a way to integrate graph structure into models processing neural-network parameters.Source: §1 · p. 1, Figure 1 · p. 2
  • The contributions include adapting graph neural networks and transformers to neural graphs, introducing probe features for neuron activations, and evaluating the approach across tasks.The paper reports outperforming state-of-the-art approaches and frames neural graphs as a benchmark for graph neural networks.Source: §1 · p. 2

2 NEURAL NETWORKS AS NEURAL GRAPHS

Neural graphs encode both neural-network parameters and architecture, allowing graph models to respect neuron permutation symmetries while representing diverse network designs. The section extends this representation from MLPs to CNNs and heterogeneous architectures, with optional node and edge features.

2.1 MLPS AS GRAPHS

  • Neural graphs map neurons to nodes and connections to edges, with biases and weights as features; their permutation symmetries match neuron permutations, and the MLP mapping is one-to-one with its computation graph.For CNNs, channels are nodes and kernels provide edge features after spatial dimensions are zero-padded to a common size and flattened.Source: §2 · p. 2, §2.1 · p. 3, §2.2 · pp. 3–4

2.2 CNNS AS GRAPHS

  • CNN graphs treat channels as nodes and kernels as edge features, padding and flattening kernels to support different kernel sizes; adaptive pooling needs no graph-construction changes.Linear layers can use an MLP-style representation or be treated as 1 × 1 convolutions for a unified representation with convolutional layers.Source: §2.2 · pp. 3–4

2.3 MODELLING HETEROGENEOUS ARCHITECTURES

  • A single graph network can process varying layer counts, layer dimensions, and connectivities, with nonlinearities encoded as node features and residual connections as edges.Transformer feedforward components are treated as MLPs, and converted components can be composed into a transformer neural graph.Source: §2.3 · p. 4

2.4 NODE AND EDGE REPRESENTATION

  • The representation can add learned probe features by running shared learned inputs through each network and storing intermediate activations and outputs as node features.Probe features remain invariant to parameter augmentations that preserve the exact function; the section also proposes reversed or undirected edge features, layer-level weight and bias normalization, and positional embeddings that distinguish input and output nodes while preserving hidden-layer symmetry.Source: §2.4 · p. 5

3 LEARNING WITH NEURAL GRAPHS

The paper adapts GNNs and transformers to process neural graphs with edge features. Its designs update edges and enable multiplicative interactions between node and edge representations.

GNN

  • The GNN uses PNA with edge-feature updates at each layer based on incident node features, producing representations that depend on graph structure.PNA does not update edge features by default, so the method adds a lightweight neural network for that update.Source: §3 · p. 5
  • Message passing uses FiLM to add multiplicative interactions between node and edge features, with scaling factors computed from edge features.The method motivates this design by the shared edge-wise weight-times-activation operation in a neural network's forward pass.Source: §3 · pp. 5–6

Transformer

  • The transformer uses relational attention to incorporate and update edge features, then modulates its value-matrix update to enable node-edge multiplicative interactions.Source: §3 · p. 6

4 EXPERIMENTS

The experiments evaluate neural-graph GNN and transformer variants across tasks with global or per-parameter outputs. The graph-based methods outperform baselines across INR classification and editing, CNN generalization prediction, and learning to optimize, with the strongest variant depending on the task.

4.1 INR CLASSIFICATION AND STYLE EDITING

  • The INR classification datasets pair images with separately optimized MLP representations, while style editing predicts weight updates to enlarge represented MNIST digits.The experiments assess classification on MNIST and Fashion MNIST and dilation-based editing on MNIST INRs.Source: §4.1 · p. 6
  • Without positional embeddings, MNIST INR accuracy is 83.9±0.3% for NG-GNN and 77.9±0.7% for NG-T, decreases of 7.5 and 14.5 points, respectively.The ablation identifies positional embeddings as important for this classification task.Source: §4.1 · p. 6
  • NG-GNN and NG-T outperform equivariant baselines by up to +11.6% on MNIST and +7.8% on Fashion MNIST INR classification, and improve style editing.Both methods outperform DWSNet and NFN even with 0 probe features, and improve further as probe features increase.Source: §4.1 · p. 6, Figure 2 · p. 7

4.2 PREDICTING CNN GENERALIZATION FROM PARAMETERS

  • On CIFAR10-GS and CNN Wild Park, NG-T outperforms baselines; the gap over StatNN is narrow on CIFAR10-GS but significant on Wild Park.NG-N performs slightly worse than NG-T, and the authors report that architecture provides crucial information on Wild Park.Source: Table 1 · p. 7, §4.2 · p. 7
  • Removing non-linearity features lowers CNN Wild Park performance by 2.6 points for NG-GNN and 8.9 points for NG-T.The reported ablation indicates that these features are especially important for the Transformer.Source: §4.2 · p. 7

4.3 LEARNING TO OPTIMIZE

  • NG-T performs best on FashionMNIST validation, while NG-GNN leads on CIFAR-10 testing, where both graph-based optimizers outperform baselines.CIFAR-10 was not used for model or hyperparameter selection; the curves show NG-GNN performing best on that task.Source: §4.3 · p. 8, Figure 3 · p. 9

5 RELATED WORK

Prior approaches to representing neural networks face trade-offs between retaining parameter detail, respecting neuron-permutation symmetry, and supporting different architectures.

5 RELATED WORK

  • Flattened parameters can yield different predictions for functionally identical networks, while parameter statistics respect symmetry but discard fine-grained detail.Permutation augmentation has also been proposed, but can require substantially more training compute because of the size of the permutation group.Source: §5 · p. 9
  • Equivariant linear layers address the symmetry problem through intricate weight-sharing patterns, but these approaches have reported an INR-classification gap relative to CNNs on images.The cited approaches compose such layers into invariant and equivariant networks for networks.Source: §5 · p. 9
  • Unlike prior methods limited to homogeneous architectures, the paper's graph-based approach can process heterogeneous architectures while maintaining neural-network symmetry.The related-work discussion contrasts this with methods designed around fixed architectural settings.Source: §5 · p. 9

6 CONCLUSION AND FUTURE WORK

The conclusion presents neural graphs as a flexible way to process neural networks, while noting that the investigation covers only a limited range of architectures and INR settings.

6 CONCLUSION AND FUTURE WORK

  • The method represents input neural networks as neural graphs, with a framework flexible enough to add domain-specific inputs such as gradients.Direct graph use also opens applications requiring varying architectures and new graph-neural-network benchmarks.Source: §6 · p. 9
  • The investigation covers MLPs and CNNs, demonstrates transformer representations only theoretically, and confines strong INR performance to 2D images.The authors identify neural radiance fields as a possible extension that would broaden utility.Source: §6 · p. 9

A.1 NEURAL GRAPH SYMMETRIES

  • A standard message-passing graph neural network is equivariant because its individual operations are designed to preserve equivariance.The graph symmetries are described using the symmetric group over the total number of graph nodes.Source: §A.1 · p. 13

A.2 NEURON PERMUTATION GROUP

The discussion compares neuron-permutation symmetries with permutations of all neural-graph nodes, showing why the broader graph symmetry supports multiple architectures.

A.2 NEURON PERMUTATION GROUP

  • The neuron-permutation group is a product of symmetric groups for the network layers, with each group element specifying a permutation per layer.The paper focuses on this group to simplify exposition and notes that its version also permits input and output neuron permutations, disambiguated with positional embeddings.Source: §A.2 · p. 13
  • The neuron-permutation group is a subgroup of the symmetric group over all graph nodes, so equivariance to the latter implies equivariance to the former.The paper expresses the layerwise permutations as a block permutation matrix to connect the two group actions.Source: §A.2 · p. 13

A.3 DISCUSSION

The discussion reports better generalization for NG-T on MNIST INR classification and describes neural graphs as a route to processing networks with varied architectures.

A.3 DISCUSSION

  • Using the symmetric group over all graph nodes makes one model equivariant to different layer-size choices, unlike prior models restricted to a fixed architecture.Layer embeddings let graph networks distinguish nodes by layer, addressing the possibility that shared update parameters are too restrictive for a fixed architecture.Source: §A.3 · p. 14
  • At matched training losses, NG-T has lower test loss than DWSNet on MNIST INR classification, and the figure reports better generalization for NG-T.The authors say parameter sharing from symmetric-group equivariance does not harm performance and might benefit generalization.Source: Figure 4 · p. 14, §A.3 · p. 14

B EXPRESSIVITY OF GNNS FOR NEURAL GRAPHS

  • An MPNN can approximate an input MLP’s forward pass layer by layer, using edge weights and node inputs and biases to reproduce multiplication, addition, and activation.Layer embeddings indicate which MLP layer to simulate, and the final MPNN node features match the simulated output.Source: B EXPRESSIVITY OF GNNS FOR NEURAL GRAPHS · pp. 14–15

C.1 MLPS AS GRAPHS

  • In the MLP graph, edges encode weighted input contributions and each node sums its incoming contributions with its bias.The described edge operation multiplies the tail-node input by the edge weight before the destination node sums incoming terms.Source: §C.1 · p. 15, §C.2 · p. 15

C.2 CNNS AS GRAPHS

  • CNN graphs preserve filter-channel permutation symmetry, while convolutional kernels are zero-padded and flattened and flattening layers require resolution-specific graph structure.Two flattening options are repeating final-layer nodes for spatial locations or adding a virtual layer of spatial nodes.Source: Figure 5 · p. 15, §C.2 · pp. 15–16, Figure 6 · p. 16

C.3 NORMALIZATION LAYERS AS GRAPHS

  • Normalization is represented as a linear graph layer, with beta terms on node features and gamma terms on edges connecting corresponding input and output dimensions.The representation uses one output node per dimension and only connects each input dimension to its corresponding output.Source: §C.3 · p. 16

C.4 TRANSFORMERS AS GRAPHS

  • Transformer graphs encode the three linear projections with multidimensional edge features, while the graph network approximates parameter-free dot-product attention.The graph construction also connects head nodes through output weights for concatenation and final projection.Source: §C.4 · p. 16

D EXPERIMENTAL SETUP DETAILS

The appendix lists model configurations and training choices for INR classification and style editing, CNN generalization prediction, and probe-feature use.

D.1 INR CLASSIFICATION AND STYLE EDITING

  • The INR classification NG-T and NG-GNN models each use four layers and 64 hidden dimensions, with dropout probability 0.2 and classification heads on final-layer node representations.Their reported parameter counts with the classification head are 378,090 for NG-T and 348,746 for NG-GNN.Source: §D.1 · p. 17
  • For style editing, the models omit final aggregation and the classification head, while NG-GNN additionally uses reversed edge features.Training runs for 100 epochs on MNIST and 150 on FashionMNIST, using the same early-stopping protocol as Navon et al. (2023).Source: §D.1 · p. 17

D.2 PREDICTING CNN GENERALIZATION FROM PARAMETERS

  • CNN generalization experiments use separate NG-T and NG-GNN configurations for Small CNN Zoo and CNN Wild Park, with task-specific training losses and epoch counts.The Small CNN Zoo models train for 300 epochs; the CNN Wild Park models train for 10 epochs.Source: §D.2 · p. 17

D.3 PROBE FEATURES

  • Probe features are used for INR classification and style editing, but not for CNN generalization prediction or learning to optimize.Source: §D.3 · p. 17

E.1 SMALL CNN WILD PARK

CNN Wild Park contains checkpoints and test accuracies from 2,800 CNNs trained on CIFAR10, with architectures varying across depth, width, kernels, activations, and skip connections.

E.1 SMALL CNN WILD PARK

  • 2,800 CNNs were trained for 200–1,000 epochs on CIFAR10, with parameter checkpoints retained every 10 steps and test accuracy recorded.Train, validation, and test splits keep checkpoints from the same run out of both train and test splits.Source: §E.1 · pp. 17–18
  • Architectures varied in 2–5 layers, 4–32 channels per layer, and convolution kernels of size 3, 5, or 7.Source: §E.1 · p. 18
  • Each layer used one of six activation options: ReLU, GeLU, tanh, sigmoid, leaky ReLU, or identity.Source: §E.1 · p. 18
  • Architectures allowed at most one incoming skip connection per layer, between layers with at least one intervening layer.When skip-connected layers had different channel counts, the connection was added only to the min(c_n, c_m) nodes.Source: §E.1 · p. 18

F EXPERIMENT RESULTS

Graph-based models outperform baselines on MNIST and Fashion MNIST INR classification; the section also reports INR dilation and feature ablations, with exact results provided for reference.

F EXPERIMENT RESULTS

  • All graph-based models outperform baselines on MNIST and Fashion MNIST INR classification.The exact experimental results are provided for reference and easier comparison.Source: Table 3 · p. 18, Table 4 · p. 18, F EXPERIMENT RESULTS · p. 18
  • The INR dilation results use mean-squared error between reconstructed and dilated ground-truth images, where lower is better.Source: Table 5 · p. 19
  • Removing non-linearity embeddings decreases CNN Wild Park Kendall’s τ by 2.6 for NG-GNN and 8.9 for NG-T.The ablation identifies non-linearity features as crucial, especially for the Transformer variant.Source: Table 6 · p. 19
  • Removing positional embeddings decreases MNIST INR classification performance by 7.5 points for NG-GNN and 14.5 for NG-T.The ablation highlights the importance of positional embeddings.Source: Table 7 · p. 19

G VERSION HISTORY

  • In version 2, rerunning the Small CNN Zoo experiments changed NG-GNN performance from 0.930±0.001 to 0.931±0.002.The paper also added Small CNN Zoo hyperparameters and made minor style edits.Source: G VERSION HISTORY · p. 19

Related papers

Loading 2403.12143v3…