Finite Probes Suffice: Identifiability and Universality for Weight-Space Learning
TL;DR
This work studies when finite probe-based representations are sufficient for learning neural functionals, establishes general identification and universality results for probing, and shows that using intermediate hidden representations can provide significantly more informative representations than relying only on final outputs.
Abstract
Learning properties of neural networks has recently attracted growing interest, with existing approaches operating either directly on network parameters or through probe-based representations of network behavior. While probing methods have shown strong empirical performance, their theoretical foundations remain limited. In this work, we study when finite probe-based representations are sufficient for learning neural functionals. We establish general identification and universality results for probing, and show that using intermediate hidden representations can provide significantly more informative representations than relying only on final outputs. Motivated by these results, we introduce HIDDENPROBE, a simple architecture for learning from hidden probe responses. Across a range of neural functional benchmarks, including both MLPs and Transformers, HIDDENPROBE consistently improves over existing probing methods and achieves state-of-the-art performance. Our code is publicly available on GitHub.
