Improved Generalization of Weight Space Networks via Augmentations

TL;DR

These strategies for data augmentation in weight spaces are explored and a MixUp method adapted for weight spaces is proposed that improves performance similarly to having up to 10 times more data and yields substantial 5-10% gains in downstream classification.

Abstract

Learning in deep weight spaces (DWS), where neural networks process the weights of other neural networks, is an emerging research direction, with applications to 2D and 3D neural fields (INRs, NeRFs), as well as making inferences about other types of neural networks. Unfortunately, weight space models tend to suffer from substantial overfitting. We empirically analyze the reasons for this overfitting and find that a key reason is the lack of diversity in DWS datasets. While a given object can be represented by many different weight configurations, typical INR training sets fail to capture variability across INRs that represent the same object. To address this, we explore strategies for data augmentation in weight spaces and propose a MixUp method adapted for weight spaces. We demonstrate the effectiveness of these methods in two setups. In classification, they improve performance similarly to having up to 10 times more data. In self-supervised contrastive learning, they yield substantial 5-10% gains in downstream classification.

Figures (9)

• Figure 1: Effect of our Alignment + MixUp augmentation on ModelNet40: MixUp without augmentation (blue) shows major signs of overfitting from the increasing test loss, while MixUp + alignment (maroon) mitigates this overfitting and improves the DWS's accuracy.
• Figure 2: Illustration of multiview INR generation and INR classification task. The INR is trained to receive $x,y$ coordinates and map them to the corresponding grayscale value in the original raw data (top panel). The trained INR is fed into the weight space architecture to perform classification (bottom panel).
• Figure 3: Overfitting of weight space architectures on the FMNIST dataset: We visualize train and test losses for DWS navon23dws on the ModelNet40 dataset with 1 or 10 trained input networks per point cloud (views). Notably, DWS tends to overfit early during training, even when using more data.
• Figure 4: Internal vs. external generalization: We visualize DWS performance on internal and external FMNIST splits with varying a number of neural views. There is a relatively small difference between the two types of generalization.
• Figure 5: Representations in weight space: Visualization of the 2D feature space attained through SimCLR contrastive learning.