Deep Linear Probe Generators for Weight Space Learning
This work addresses the problem of weight space learning for researchers and practitioners in machine learning, offering an incremental improvement over existing probing methods.
The paper tackles the challenge of extracting information from neural network weights by proposing Deep Linear Probe Generators (ProbeGen), a modification to probing approaches that adds a shared generator module to reduce overfitting. The result is a method that significantly outperforms state-of-the-art techniques while being 30 to 1000 times more efficient in terms of FLOPs.
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.