Graph Neural Networks for Learning Equivariant Representations of Neural Networks

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

What

This paper introduces a novel approach to representing neural networks as computational graphs of parameters called ‘neural graphs’, which allows for leveraging powerful graph neural networks and transformers while preserving permutation symmetry.

Why

This research is important because it addresses limitations of existing methods that process neural network parameters, such as overlooking inherent permutation symmetry or relying on complex weight-sharing patterns. By representing neural networks as graphs, this approach allows a single model to learn from diverse architectures and opens up new possibilities for applications like neural network analysis, generation, and optimization.

How

The authors represent neural networks as graphs by mapping neurons to nodes and connections to edges, with weights and biases as edge and node features respectively. This representation is then used as input to graph neural networks (GNNs) or transformers, adapted to incorporate inductive biases from the neural graph structure. They validate their approach on various tasks including implicit neural representation classification and editing, predicting generalization performance of CNNs, and learning to optimize.

Result

The proposed neural graph approach consistently outperforms state-of-the-art methods on tasks like INR classification and style editing, showing significant improvement over previous methods like DWSNet and NFN. It also demonstrates superior performance in predicting CNN generalization, especially when dealing with diverse architectures where accounting for both parameters and architecture is crucial. Furthermore, the method shows promise in the field of learning to optimize, achieving strong performance on both validation and test tasks.

LF

The authors acknowledge limitations in terms of architectural diversity explored, focusing mainly on MLPs and CNNs. Future work could investigate the representation of other architectures like transformers. Additionally, the strong performance on INRs is currently limited to 2D images, and extending it to handle 3D representations like neural radiance fields is an area for further exploration.

Abstract

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.