Draft:Learning in Weight Spaces

Learning in Weight Spaces

Learning in weight spaces is an area in machine learning, and specifically in Deep learning, that focuses on processing the weights or parameters of neural networks as data. This approach represents a paradigm shift from traditional machine learning methods by treating the weights of a trained neural network as a distinct data modality and as the input domain for another learning algorithm. By conceptualizing weights as a data modality in their own right, alongside more conventional modalities like images or text, this field opens up new possibilities for analyzing, manipulating, and understanding neural networks at a fundamental level. This perspective allows researchers to apply machine learning techniques directly to the weight space, enabling novel approaches to model analysis, transfer learning, and meta-learning.

Overview

In traditional machine learning, neural networks are trained to map inputs to desired outputs by adjusting their weights and biases. Learning in weight spaces, however, uses these trained weights and biases as inputs for secondary models. These secondary models can be designed to predict properties of the original network, modify its behavior, or extract meaningful representations from its parameters.

The field encompasses a wide range of techniques and applications, including:

• Hyperparameter optimization: Using weight space features to improve efficiency in selecting optimal hyperparameters for neural networks.
• Model Selection and Fine-tuning: Enhancing the process of selecting and adapting pretrained models for specific tasks.
• Processing Implicit Neural Representations (INRs): Classifying and editing INRs of images and 3D shapes.
• Predicting Neural Network Properties: Estimating test accuracy, generalization gap, or runtime without full evaluation of the network.
• Data augmentation: Developing novel augmentation techniques specific to weight spaces.
• Self-supervised learning: Applying contrastive learning techniques in weight spaces to learn meaningful representations.
• Transfer learning and Few-shot learning: Leveraging weight space information to improve adaptation to new tasks with limited data.

A key challenge in learning in weight spaces lies in designing models that can effectively capture and utilize the information encoded in neural network parameters. This requires careful consideration of the high-dimensional nature of weight spaces and the particular ways in which neural network weights encode information about their learned tasks. Recent advances in this field have led to the development of specialized architectures that incorporate the inherent symmetries of neural networks, such as permutation and scale invariance.

Formal Definition

Weight space for MLPs

For a multilayer perceptron (MLP) with $L$ layers, parameterized by weights $W=\{W^{(i)}\in \mathbb {R} ^{n_{i}\times n_{i-1}}\mid i\in [1..L]\}$ and biases $v=\{v^{(i)}\in \mathbb {R} ^{n_{i}}\mid i\in [1..L]\}$, where $n_{i}$ is the number of neurons in layer $i$, the function $f$ can be expressed as: $f(x)=x_{L},{\text{ where }}x_{m+1}=\sigma (W^{(m+1)}x_{m}+v^{(m+1)}),x_{0}=x$ Here, $\sigma$ is a pointwise non-linearity (e.g., ReLU or sigmoid function). The weight space $U$ of this MLP is defined as: $U:=\bigoplus _{m=1}^{L}(W_{m}\oplus B_{m})$ Where $W_{m}:=\mathbb {R} ^{n_{m}\times n_{m-1}}$ and $B_{m}:=\mathbb {R} ^{n_{m}}$.

Weight space for general models

For any parametric model $f_{\theta }:X\to Y$ with parameters $\theta \in \Theta$, the weight space is defined as $U:=\Theta$.

Learning Problem

Given a dataset $D={(U_{i},z_{i})}_{i=1}^{N}$, where $U_{i}\in U$ and $z_{i}\in Z$, the goal is to learn a function $g:U\to Z,$ that minimizes some loss $L$: $g^{*}={\underset {g}{\operatorname {argmin} }}\mathbb {E} _{(U,z)\sim D}[L(g(U),z)]$.

Approaches

Pioneering Methods

Initial approaches to learning in weight spaces focused on using standard neural network architectures and feature-based methods to process the weights of other networks. These methods included:

• Feature-based methods: Computing statistical features from weights and using them as inputs to machine learning models. Unterthiner et al. demonstrated that simple statistics of weights (mean, variance, and quintiles) could be used to predict neural network accuracy with surprising effectiveness.
• Convolutional Neural Networks (CNNs): Using 1D CNNs to process weight matrices and tensors. Eilertsen et al. used this approach to predict various properties of neural networks, including hyperparameters and performance metrics.

• Model Zoos: Collections of trained neural network models that serve as datasets for weight space learning research. Schürholt et al. introduced a large-scale, standardized dataset of diverse neural network populations. This resource facilitates research on model analysis, learning dynamics, representation learning, and model generation in weight spaces. In their work, the authors explored various architectures for processing weights and predicting model properties, including linear models, gradient boosting machines (particularly LightGBM), deep neural networks with fully-connected layers and ReLU activations, and 1D convolutional neural networks adapted to process the sequence of flattened weights.
• Hyper-Representations: Schürholt et al. introduced the concept of hyper-representations - learned representations of neural network weights. They demonstrated that these representations can be used to predict various model characteristics and even generate new model weights with desired properties. This work showed the potential of self-supervised learning in weight spaces. The authors introduced several architectures for learning hyper-representations, including attention-based encoders using multi-head self-attention modules, different weight encoding methods (individual weight encoding and neuron-level encoding), and various compression strategies (sequence aggregation and compression token). These architectures were combined with different self-supervised learning objectives, including reconstruction and contrastive learning, to create effective representations of neural network weights.

Incorporating Weight Space Symmetry

A key aspect of learning in weight spaces is understanding and leveraging the inherent symmetries present in neural network architectures. For Multilayer Perceptrons (MLPs), one crucial symmetry is permutation symmetry.

Definition of Permutation Symmetry in MLP Weight Spaces: Let $W=\{W^{(i)}\in \mathbb {R} ^{n_{i}\times n_{i-1}}\mid i\in [1..L]\}$ be the weights of an $L$-layer MLP. The permutation symmetry of the weight space can be defined as follows:

For any layer $i$ and any permutation $\pi$ of $n_{i}$ elements, the transformation:

${\begin{aligned}W^{(i)}&\rightarrow P_{\pi }W^{(i)}\\W^{(i+1)}&\rightarrow W^{(i+1)}P_{\pi }^{-1}\end{aligned}}$

where $P_{\pi }$ is the permutation matrix corresponding to $\pi$, does not change the function computed by the network. This symmetry arises because neurons within a layer can be arbitrarily reordered without affecting the network's output, as long as the connections to the subsequent layer are adjusted accordingly. Symmetries like these provide important information about the data type we are processing in weight spaces. Leveraging these symmetries can lead to the design of much more effective and efficient architectures for learning in weight spaces. By incorporating symmetry information directly into the model architecture, we can:

• Reduce the hypothesis space, potentially leading to faster learning and better generalization.
• Ensure that the model's predictions are invariant to irrelevant transformations of the input weights.
• Improve sample efficiency by enforcing correct inductive biases.

More advanced approaches incorporate weight space symmetries directly into the network architecture. Notable examples include:

• Deep Weight Space Networks (DWSNets) and Permutation Equivariant Neural Functionals

These networks respect neuron permutation symmetries through parameter-sharing constraints.

• Neural Functional Transformers (NFTs): NFTs incorporate attention mechanisms while respecting neuron permutation symmetries.

Graph-based Approaches

Recent work has explored representing neural networks as graphs to enable processing by graph neural networks (GNNs). Key developments include:

• Graph Metanetworks (GMNs): GMNs represent input neural networks as parameter graphs.
• Neural Graphs: This approach encodes weights as edge features and biases as node features.
• Scale Equivariant Graph MetaNetworks (ScaleGMNs): These networks incorporate scaling symmetries in addition to permutation symmetries. Scale symmetries in this context refer to the invariance of neural network function under certain transformations of the weights. For example, in networks with ReLU activations, multiplying all weights and biases in a layer by a positive constant and dividing those in the subsequent layer by the same constant does not change the network's output. ScaleGMNs are designed to respect these symmetries, potentially leading to more efficient and effective processing of neural network weights.

Comparison with Hypernetworks

The concept of hypernetworks was introduced by Ha et al. in their seminal paper "HyperNetworks" (2016). This work laid the foundation for a new approach to neural network architecture, which has since been explored and expanded upon in various directions, including some that intersect with the goals of learning in weight spaces.

Learning in weight spaces shares some conceptual similarities with hypernetworks, as both involve meta-learning approaches that operate on neural network parameters. However, they differ significantly in their goals, methodologies, and applications.

The primary distinction lies in their fundamental purpose and directionality. Learning in weight spaces focuses on analyzing and processing the weights of already trained networks to extract information or predict properties. It typically operates on the high-dimensional weight space of trained networks, with the flow of information going from weights to some output, such as network performance predictions, property estimations or in some cases, new weights. In contrast, hypernetworks are designed to generate weights for another network (the primary network) based on some input. They work in the opposite direction, mapping from a lower-dimensional input space to the high-dimensional weight space of the primary network.

Challenges and Future Directions

While learning in weight spaces has shown promise, several challenges remain:

• Scaling to very large models or diverse architectures
• Developing more efficient and expressive representations of weight spaces

Future research may focus on addressing these challenges and exploring new applications in areas such as model compression, neural architecture search, and continual learning.