Deep Linear Probe Generators for Weight Space Learning

TL;DR

This work tackles weight-space learning by improving probing methods to predict undocumented model attributes from a network's responses to fixed inputs. The authors identify suboptimal learned probes and introduce ProbeGen, a Deep Linear Probe Generator that shares a linear generator across probes to impose inductive bias and reduce overfitting. Empirical results across INR and CNN generalization tasks, including a large ResNet18 Model Zoo, show ProbeGen achieving state-of-the-art performance with 30–1000× fewer FLOPs than graph-based approaches. The approach demonstrates strong scalability, interpretable probe representations, and applicability to black-box models, with potential extensions to other modalities and adaptive probing. Overall, ProbeGen offers a simple yet powerful mechanism to extract meaningful, efficient insights from neural network weights.

Abstract

Weight space learning aims to extract information about a neural network, such as its training dataset or generalization error. Recent approaches learn directly from model weights, but this presents many challenges as weights are high-dimensional and include permutation symmetries between neurons. An alternative approach, Probing, represents a model by passing a set of learned inputs (probes) through the model, and training a predictor on top of the corresponding outputs. Although probing is typically not used as a stand alone approach, our preliminary experiment found that a vanilla probing baseline worked surprisingly well. However, we discover that current probe learning strategies are ineffective. We therefore propose Deep Linear Probe Generators (ProbeGen), a simple and effective modification to probing approaches. ProbeGen adds a shared generator module with a deep linear architecture, providing an inductive bias towards structured probes thus reducing overfitting. While simple, ProbeGen performs significantly better than the state-of-the-art and is very efficient, requiring between 30 to 1000 times fewer FLOPs than other top approaches.

Figures (11)

• Figure 1: Overview of Our Method. We optimize a deep linear probe generator to create suitable probes for the model. Meaning, our generator includes no activations between its linear layers, yet the addition of linear layers reinforces a desired structure for the probes. We then gather the models responses over all probes, and train a classifier to predict some attribute of interest about the model.
• Figure 2: A Few Examples from the Dead Leaves Dataset. We show these images are synthetic and highly dissimilar to real images.
• Figure 3: Learned Probes. Optimized probes by (a) Latent Optimization and (b) ProbeGen for the CIFAR10 Wild Park benchmark. We show the same probes for identical runs (including the seed), except the generator module.
• Figure 4: Overfitting for Different Generators. We compare ProbeGen to linear generators with fully connected layers, and non-linear convolutional generators. Overfit is measured by the generalization gap of each method. Results are averaged over $5$ seeds.
• Figure 5: Ablation Studies. We compare the performance of ProbeGen, linear generators with fully connected layers, and non-linear convolutional generators, for varying numbers of layers. With less than $2$, all versions are equivalent to ProbeGen. We average over $5$ seeds.