Finite Probes Suffice: Identifiability and Universality for Weight-Space Learning

2026-09-27

AI Summary

This study addresses the lack of theoretical foundations for finite probe representations in neural network property learning, where reliance solely on final outputs yields insufficient information. To bridge this gap, we establish identifiability and universality theories for probe learning, deriving the first sufficiency bounds for finite probes and demonstrating that intermediate hidden-layer representations are superior to final outputs. Guided by these theoretical insights, we propose HIDDENPROBE, a minimalist yet highly efficient architecture. We apply this framework, supported by rigorous theoretical analysis, to both MLPs and Transformers. Extensive evaluations across multiple neural functionality benchmarks show that HIDDENPROBE consistently outperforms existing methods, achieving state-of-the-art performance. The source code has been made publicly available.

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.

Problem

• neural functionals
• weight-space learning
• probe-based representations
• identifiability
• universality

Innovation

• Probe-based Representations
• Identifiability
• Universality
• Hidden Probes
• Neural Functionals