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

Soutrik Sarangi; Yonatan Sverdlov; Adir Dayan; Haggai Maron; Nadav Dym

cs.LGSep 27, 2026

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

Authors: Soutrik Sarangi, Yonatan Sverdlov, Adir Dayan, Haggai Maron, Nadav Dym

Organizations: Microsoft · Faculty of Mathematics, Technion – Israel Institute of Technology · Faculty of Electrical and Computer Engineering, Technion – Israel Institute of Technology · NVIDIA

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.

Figures & tables

Figure 1: Overview of HiddenProbe. A set of learnable probes $p_{1},\ldots,p_{N}$ is evaluated on a target neural network $f_{\theta}$ , producing intermediate activations $x_{i}(p_{j},\theta)$ and final outputs $x_{L}(p_{j},\theta)$ , which are processed by an invariant model and an MLP to predict a property of the target network.

Figure 2: For any finite set of probes, a nonzero ReLU ‘hat’ function can be constructed whose support avoids all probes.

Method	MNIST	FMNIST	CIFAR-10	CIFAR-10 Aug
ScaleGMN	$0.966\pm 0.002$	$0.808\pm 0.001$	$0.388\pm 0.001$	$0.570\pm 0.005$
Neural Graphs (128 probes)	$0.976\pm 0.001$	$0.745\pm 0.008$	–	–
NG-T	$0.924\pm 0.003$	$0.727\pm 0.006$	–	–
MAGEP-NFN	–	–	$0.3718\pm 0.003$	–
NFT	–	–	–	$0.634\pm 0.000$
ProbeGen	$0.984\pm 0.001$	$0.877\pm 0.003$	$0.573\pm 0.007$	$0.563\pm 0.003$

Table 1: Classification accuracy (

\uparrow

) on INR weight-space benchmarks. Uncertainties denote the standard error over five seeds. Baseline numbers are copied from the respective papers

Method	MNIST	FMNIST	SVHN	CIFAR-WP	CIFAR-GS
NFN	$0.942\pm 0.001$	$0.935\pm 0.000$	$0.931\pm 0.005$	–	$0.934\pm 0.001$
ScaleGMN	–	–	–	–	$0.941\pm 0.000$
NG-T	–	–	$0.872\pm 0.001$	$0.817\pm 0.007$	$0.935\pm 0.000$
DNG-Encoder	–	–	$0.867\pm 0.002$	$0.874\pm 0.002$	$0.936\pm 0.000$
Neural Graphs	–	–	–	$0.885\pm 0.005$	$0.938\pm 0.000$
ProbeGen	$0.953\pm 0.002$	$0.948\pm 0.005$	$0.876\pm 0.005$	$0.932\pm 0.006$	$0.957\pm 0.001$

Table 2: Test-accuracy prediction performance measured by Kendall’s

\tau

(

\uparrow

). Uncertainties denote standard error over five seeds. Baseline numbers are copid from the respective papers.

Method	MNIST-Transformers	AGNews-Transformers
MLP	$0.866\pm 0.002$	$0.879\pm 0.006$
STATNN ( Unterthiner et al., 2020 )	$0.881\pm 0.001$	$0.841\pm 0.002$
XGBoost ( Chen and Guestrin, 2016 )	$0.860\pm 0.002$	$0.859\pm 0.001$
LightGBM ( Ke et al., 2017 )	$0.858\pm 0.002$	$0.835\pm 0.001$
Random Forest ( Breiman, 2001 )	$0.772\pm 0.002$	$0.774\pm 0.003$
Transformer-NFN ( Tran et al., 2025 )	$0.905\pm 0.002$	$0.910\pm 0.001$

Table 3: Test accuracy prediction on the MNIST-Transformers and AGNews-Transformers benchmarks, evaluated on the full datasets. Performance is measured by Kendall’s

\tau

, reported as mean

\pm

standard error over five seeds. The best and second-best results are highlighted.

Figure 3: Performance of HiddenProbe and ProbeGen as the number of probes

N

varies. FMNIST INR classification is evaluated by test accuracy, while regression tasks are evaluated by test Kendall’s

\tau

. Higher is better in all panels.

Appendix figures & tables9 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 4 — Figure 4: Reconstruction of Analytic Net from $\Psi_{\text{final}},\Psi_{\text{all}}$ .

Panel of Figure 5 — Figure 5: Reconstruction of ReLU Net from $\Psi_{\text{final}},\Psi_{\text{all}}$ .

Figure 6: Reconstruction of MHA Stack from $\Psi_{\text{final}},\Psi_{\text{all}}$ .

Method	MNIST	FMNIST	CIFAR-10	CIFAR-10 Aug
StatNN	$0.398\pm 0.001$	$0.418\pm 0.002$	–	–
DWS	$0.857\pm 0.006$	$0.671\pm 0.003$	$0.3445\pm 0.004$	$0.4127\pm 0.001$
NFN	$0.791\pm 0.008$	$0.689\pm 0.006$	$0.3341\pm 0.001$	$0.4660\pm 0.001$
ScaleGMN	$0.966\pm 0.002$	$0.808\pm 0.001$	$0.3882\pm 0.001$	$0.5695\pm 0.005$
NG-GNN	$0.914\pm 0.006$	$0.680\pm 0.002$	$0.3604\pm 0.004$	$0.4570\pm 0.002$
NG-T	$0.924\pm 0.003$	$0.727\pm 0.006$	–	–

Appendix

Table 4: Classification accuracy (

\uparrow

) on INR weight-space benchmarks.

Method	MNIST	FMNIST	SVHN	CIFAR-WP	CIFAR-GS
MLP	$0.878\pm 0.001$	$0.874\pm 0.001$	$0.809\pm 0.003$	–	$0.880\pm 0.000$
StatNN	$0.926\pm 0.000$	$0.915\pm 0.000$	$0.843\pm 0.000$	$0.719\pm 0.010$	$0.914\pm 0.000$
SNE	$0.941\pm 0.000$	$0.928\pm 0.001$	$0.858\pm 0.003$	–	$0.927\pm 0.000$
NFN NP	$0.937\pm 0.000$	$0.922\pm 0.001$	$0.856\pm 0.001$	–	$0.922\pm 0.001$
NFN HNP	$0.942\pm 0.001$	$0.935\pm 0.000$	$0.931\pm 0.005$	–	$0.934\pm 0.001$
NFT	–	–	$0.858\pm 0.000$	–	$0.926\pm 0.001$

Appendix

Table 5: Full comparison of test-accuracy prediction performance measured by Kendall’s

\tau

(

\uparrow

). Results that were not reported for a particular benchmark are denoted by “–”.

Accuracy threshold	No threshold	20%	40%	60%	80%
MLP	$0.866\pm 0.002$	$0.873\pm 0.001$	$0.874\pm 0.003$	$0.874\pm 0.006$	$0.873\pm 0.007$
STATNN ( Unterthiner et al., 2020 )	$0.881\pm 0.001$	$0.872\pm 0.001$	$0.868\pm 0.001$	$0.860\pm 0.001$	$0.856\pm 0.001$
XGBoost ( Chen and Guestrin, 2016 )	$0.860\pm 0.002$	$0.839\pm 0.004$	$0.869\pm 0.003$	$0.846\pm 0.001$	$0.884\pm 0.001$
LightGBM ( Ke et al., 2017 )	$0.858\pm 0.002$	$0.835\pm 0.001$	$0.847\pm 0.001$	$0.822\pm 0.001$	$0.830\pm 0.001$
Random Forest ( Breiman, 2001 )	$0.772\pm 0.002$	$0.758\pm 0.004$	$0.769\pm 0.001$	$0.752\pm 0.001$	$0.759\pm 0.001$
Transformer-NFN ( Tran et al., 2025 )	$0.905\pm 0.002$	$0.899\pm 0.001$	$0.895\pm 0.001$	$0.895\pm 0.002$	$0.888\pm 0.002$

Appendix

Table 6: Performance measured by Kendall’s

\tau

on the MNIST-Transformers dataset. Uncertainties indicate the standard error over 5 seeds.

Accuracy threshold	No threshold	20%	40%	60%	80%
MLP	$0.879\pm 0.006$	$0.875\pm 0.001$	$0.841\pm 0.012$	$0.842\pm 0.001$	$0.862\pm 0.006$
STATNN ( Unterthiner et al., 2020 )	$0.841\pm 0.002$	$0.839\pm 0.003$	$0.812\pm 0.003$	$0.813\pm 0.001$	$0.812\pm 0.001$
XGBoost ( Chen and Guestrin, 2016 )	$0.859\pm 0.001$	$0.852\pm 0.002$	$0.872\pm 0.002$	$0.874\pm 0.001$	$0.872\pm 0.001$
LightGBM ( Ke et al., 2017 )	$0.835\pm 0.001$	$0.845\pm 0.001$	$0.837\pm 0.001$	$0.835\pm 0.001$	$0.820\pm 0.001$
Random Forest ( Breiman, 2001 )	$0.774\pm 0.003$	$0.801\pm 0.001$	$0.797\pm 0.001$	$0.798\pm 0.002$	$0.773\pm 0.001$
Transformer-NFN ( Tran et al., 2025 )	$0.910\pm 0.001$	$0.908\pm 0.001$	$0.897\pm 0.001$	$0.896\pm 0.001$	$0.890\pm 0.001$

Appendix

Table 7: Performance measured by Kendall’s

\tau

on the AGNews-Transformers dataset. Error bars indicate the standard error over 5 seeds.

Method	INR classification			CNN regression	Transformer zoo
	MNIST	FMNIST	CIFAR-10	CIFAR-10-GS	MNIST-Trans.
MLP (Monomial-NFN setup)	$2.00$	$2.00$	$2.00$	–	–
MLP (Transformer setup)	–	–	–	–	$0.933$
StatNN	–	–	–	$1.06$	$0.203$
DWS	–	$0.553$	–	–	–
NG-GNN (no probes)	–	$0.331$	–	–	–

Appendix

Table 8: Number of trainable parameters (millions, M) for published baselines on three INR classification tasks, CIFAR-10-GS regression, and MNIST-Transformers. INR and CNN counts are taken from Tran et al. (2024) ; Tran et al. (2026) and the ScaleGMN supplement ( Kalogeropoulos et al., 2024 ) ; MNIST-Transformers counts are from Tran et al. (2025) ; Tran et al. (2026) . Published model configurations may vary across sources and need not exactly match those used in our accuracy comparisons. ProbeGen and HiddenProbe counts are provided by the authors. A dash indicates that a count has not been verified for the specified task.

Method	Time (s)	Relative to MLP
MLP	103.461	1.00 $\times$
NFN	90.755	0.88 $\times$
ProbeGen	109.122	1.05 $\times$
HiddenProbe	155.508	1.50 $\times$
NFT	164.774	1.59 $\times$
DWS	174.766	1.69 $\times$

Appendix

Table 9: Runtime on the full MNIST-INR training split (55,000 target INRs, batch size 32).

Explore similar work

Deep Minds and Shallow Probes
May 12, 2026Su Hyeong Lee, Risi KondorNeural Representations Hidden States
Department of Statistics, University of Chicago. · Department of Statistics and Department of Computer Science, University of Chicago.
What Linear Probes Miss: Multi-View Probing for Weight-Space Learning
May 22, 2026Eunwoo Heo, Kyeongkook Seo, Jaejun YooWeight-Space Multi-View
Graduate School of Artificial Intelligence, Ulsan National Institute of Science and Technology, Ulsan, Republic of Korea.
Loss-Invariant Projections as Passive Probes of Learned Representations
Oct 5, 2026Akshay Chandrasekhar, Pavlo Melnyk
Linköping University

May 12, 2026, cs.LGJ/K move · Enter open · S save

Deep Minds and Shallow Probes

Su Hyeong Lee, Risi Kondor

Department of Statistics, University of Chicago. · Department of Statistics and Department of Computer Science, University of Chicago.

Neural Representations Hidden States Symmetry

Open paper AskPDF arXiv

Neural representations are not unique objects. Even when two systems realize the same downstream computation, their hidden coordinates may differ by reparameterization. A probe family intended to reveal structure already present in a representation should therefore be stable under the relevant representation symmetries rather than be tied to a particular basis. We prove that, under standard nondegeneracy assumptions, behaviorally equivalent LLMs have hidden representations related by an invertible affine transformation. The resulting symmetry principle singles out a unique hierarchy of shallow coordinate-stable probes, with linear probes as its degree-1 member. We also show that a natural object for cross-model probe transfer is a shared probe-visible quotient--the representation modulo directions invisible to the probe family--rather than the full hidden state. Experiments on synthetic and real-world tasks support both predictions, showing where degree-2 probes help beyond linear ones and how quotient-based transfer enables coverage-aware monitor portability across model families. These results point toward a broader geometric representation theory of neural probing, with coverage-aware monitor transfer as a concrete operational consequence.

Abstract

Explore similar work

Deep Minds and Shallow Probes

What Linear Probes Miss: Multi-View Probing for Weight-Space Learning

Loss-Invariant Projections as Passive Probes of Learned Representations