Towards Scalable and Versatile Weight Space Learning

TL;DR

The SANE approach to weight-space learning overcomes previous limitations by learning task-agnostic representations of neural networks that are scalable to larger models of varying architectures and that show capabilities beyond a single task.

Abstract

Learning representations of well-trained neural network models holds the promise to provide an understanding of the inner workings of those models. However, previous work has either faced limitations when processing larger networks or was task-specific to either discriminative or generative tasks. This paper introduces the SANE approach to weight-space learning. SANE overcomes previous limitations by learning task-agnostic representations of neural networks that are scalable to larger models of varying architectures and that show capabilities beyond a single task. Our method extends the idea of hyper-representations towards sequential processing of subsets of neural network weights, thus allowing one to embed larger neural networks as a set of tokens into the learned representation space. SANE reveals global model information from layer-wise embeddings, and it can sequentially generate unseen neural network models, which was unattainable with previous hyper-representation learning methods. Extensive empirical evaluation demonstrates that SANE matches or exceeds state-of-the-art performance on several weight representation learning benchmarks, particularly in initialization for new tasks and larger ResNet architectures.

Figures (11)

• Figure 1: Aggregated results of 56 experiments showing (left:) four discriminative downstream tasks in $R^2$, and (right:) four generative downstream tasks in accuracy, each evaluated on (bottom:) CNNs model zoos trained on 4 datasets (NMIST, SVHN, CIFAR-10, STL) and (top:) ResNet18 model zoos trained on three datasets (CIFAR-10, CIFAR-100, Tiny-ImageNet). The colors indicate performance of Red: raw NN weights, Orange: weight statistics from unterthinerPredictingNeuralNetwork2020, Green: trained hyper-representations from schurholtSelfSupervisedRepresentationLearning2021schurholtHyperRepresentationsGenerativeModels2022, and Blue:SANE (ours). While some methods perform well on specific tasks, or are restricted by the size of the underlying models, SANE can deliver excellent performance on all tasks and model sizes.
• Figure 2: Given model zoos trained on different classification tasks, we extract and sequentialize the model weights. SANE trains hyper-representations on weights subsequences, i.e., individual layers. SANE can be used for multiple downstream tasks, either using the encoder for discriminative tasks such as the prediction of model accuracy, or the decoder for generative tasks such as sampling of new models.
• Figure 3: Comparison between WeightWatcher (WW) features (left) and SANE (right). Features over layer index for ResNets from pytorchcv of different sizes.
• Figure 4: Comparison between WeightWatcher features (left) and SANE (right). Accuracy over model features for ResNets and VGGs from pytorchcv of different sizes. Although SANE is pretrained in a self-supervised fashion, it preserves the linear relation of a globally-aggregated embedding to model accuracy.
• Figure 5: Comparison between sampled models and random initialization trained for 5 epochs Tiny-Imagenet. Different architectures are sampled from SANE pretrained on a ResNet-18 CIFAR-100 zoo. Although both models and tasks are changed, sampled models perform better.