Deep Linear Probe Generators for Weight Space Learning
Authors: Jonathan Kahana
Eliahu Horwitz
Imri Shuval
Yedid Hoshen
Venue: ICLR 2025 Poster (2025)
Keywords: Deep Weight Space
Weight Classification
Primary area: transfer learning, meta learning, and lifelong learning

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.

Published as a conference paper at ICLR 2025
DEEP LINEAR PROBE GENERATORS FOR WEIGHT SPACE LEARNING
Jonathan Kahana, Eliahu Horwitz, Imri Shuval, Yedid Hoshen School of Computer Science and Engineering The Hebrew University of Jerusalem, Israel jonathan. kahana@mal i hhuji.ac.in
Project page: /ttps://vision. wwwi. ac. i1 il/probegen/
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 top of the corresponding outputs. Although
the model, and training predictor probing is typically not used as a stand alone approach, our preliminary experi- ment found that vanilla probing baseline worked surprisingly well. However, we discover that current probe learning strategies are ineffective. We therefore propose Deep Linear Probe Generators (ProbeGen), simple and effective mod- ification to probing approaches. ProbeGen adds a shared generator module with 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 FLOOSs than other top approaches.
INTRODUCTION
The growing importance and popularity of neural networks has led to the development of several model hubs (e.g. HuggingFace, CivitAl), where more than million models are now publicly avail- able. Treating them as new data modality presents new opportunities for machine learning. Specif- ically, as not all neural networks include information about their training, developing methods to au- tomatically learn from weights becoming important. For instance, given an undocumented model, it is interesting to know its generalization eтror (Unterthiner al., 2020) its training dataset (Dupont et al. 2022). While some of these questions could be answered by evaluating the model on many labelled samples, this is often impractical the data may be unavailable or unknown. Here, we want to answer these questions without access to the models true data distribution and under minimal computational effort.
Learning from weights is essentially similar to the well studied problem of binary code analysis Shoshitaishvili al. (2016), where the goal is predict the function of unknown software using its binary code. In both cases, the task is to understand an unknown complex function specified by many parameters. Binary code analysis approaches generally fall into two categories: static and dynamic. Static methods Shoshitaishvill et al. (2116) aim to understand function without running the binary code. Dynamic code analysis Egele et (2008): Bayer al.) (2066) runs the code on inputs provided by the user and analyzes its outputs to understand what the code does. Similarly, in weight space learning, there are two main types of methods: mechanistic approaches CaaroacchsUUntertinnee (Unterthiner iet al. (2020): : Navon al. (2023a): Zhou et al. (2024a); Lim et al. (2023); Kalogeropoulos et al. (2024) aim to understand model weights without running the model, while probing approaches Hermann al. (2024): Kofinas al. (2024) represent models by their responses to set of well selected inputs. Despite dynamic methods often performing better for binary code analysis, in the context of weight space learning, probing is still under utilized.
Motivated by the success of dynamic code analysis, in this paper we focus on advancing probing methods for learning from weights. First, we support this intuition by showing that simple prob-

---

Published as a conference paper at ICLR 2025
4
(Pp)
<2
Deep Linear Probe (G)
Evaluated
(P,)
Ps
Neural
((rs)
Classifier
Network
If)
(Pg)
Px
Figure 1: Overview of Our Method. We optimize deep linear probe generator 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 interest about the model.
ing baseline, with no bells-and-wiistle, achieves comparable or better results than state-of-the-ar mechanistic approaches. However, we discover that despite the good performance, current probing methods learn probes that perform comparably to random probes sampled from simple, unlearned statistical distributions. This suggests that current learned probes are suboptimal.
To this end, we propose Deep Linear Probe Generators (ProbeGen) as simple and effective so- lution. ProbeGen factorizes its probes into two parts, a per-probe latent code and a global probe generator. The generator offers two key benefits: (i) It helps sharing information across multiple probes, and (ii) can implicitly introduce an inductive bias into the probes. For example, im- ages, hierarchical and convolutional layers create a local structure. Finally, by observing the learned probes we hypothesize they are not necessarily semantic, and owe some of their expression to low- level structures. We then find that the non-linear activation functions, which increase expressivity, actually degrade the learned probes. Our final approach therefore consists of a deep linear network (Arora et al.) 2019), with data-dependent biases. Our linear generators produce probes that achieve state-of-the-ae art performance on common weight space learning tasks.
RELATED WORK
Weight Space Learning. A recent line of works have focused training models process the weights of diverse model populations (known as Model Zoos, c.g., (Schürholt et al. 2022)) to predict undocumented properties of a model. These properties include classifying the training dataset (e.g.,
the class of an image used to train an Implicit Neural Representation) or predicting the generalization error. Mechanistic methods (Lim al. 2023 Navon et al., 2023a] Schürholt et al., 2024 Eilertsen
et al. 2020 Unterthiner et al., 2020) represent models using the mechanics of their inner workings. One approach (De Luigi et al. 2023; Schürholt et al. 2024 202T) to learn standard architectures over raw weights. However, weights exhibit permutation symmetries between neurons (Navon et al.) 2023a) which these architectures do not explicitly account for, although some of these methods use augmentations (Schürholt et al. 2024 2021) to encourage permutation invariance. Other methods proposed specialized augmentations (Shamsian et al. 2024) which can be applied to aligned net-
works (Shamsian et al., 2024 Navon al. 2023b] Peña et al. 2023 Ainsworth et al., 2022 to
already aligned weight spaces (Lim et al. 2024b ap. Another approach (Unterthiner et al. 2020
Dupont et al., 2022; Salama et 2024) maps the weights a low dimensional embedding by of weight statistics, which are completely invariant to permutations, but have limited expressivity
as they ignore the inner relations between neurons. Recently. line of works (Navon al.) 2023a
Zhou et al.] 2024a]] b, Kofinas et al,, [2024] Lim et al., (2023) Tran et al., 2024 Kalogeropoulos et al. 2024) focus on specially designing permutation equivariant architectures for processing neural net- works. A dominant approach uses graph based architectures (Kořnnas et al., 2024 Lim et al.) 2023
Kalogeropoulos et al., (2024) modeling a neural network as a computational graph, where every neu- ron node. They then train Graph Neural Network (GNN) (Gilmer et al. 2017: Kipf & Welling 2016) or Transformer (Vaswani 2017 Diao & Loynd 2022) modules, which are equivariant by design, to analyze the computational graph. Most recently, some equivariant approaches (Kofinas) ctt al. 2024; Herrrmann et al. 2024) also included learned probes. In this work we take a deep look into probing methods, and their failure points.
2

---

Published as a conference paper at ICLR 2025
Another research thrust developed new applications for weight space learning. Some works (Ha ct al.
2016 Ashkenazi et al.
2022 Peebles et al., (022) encode the parameters of neural networks,
manly for generating, moditying
compressing weights matrices. Others (Erkoç al. 2023 2023) use weights for advanced image generation capabilities. More
Dravid et al. 2024 Shah et al.
recently, Carlini al. (2024) proposed recovering the exact black-boxed weights of neural net-
(2024a) demonstrated recovering entire pre-trained models when the weights of multiple fine-tuned models of the same Model Tree Horwitz et al. (2024b) are available.
work layer, and Horwitz et al.
Implicit Neural Representations (INRs). In recent years, INRs (Sitzmann et al. 2020 Tan- cik et al. 2020) have emerged as new paradigm for representing data points with neural net- works. These networks are used in various data modalities such as images (Ha 2016). 3D shapes (Mescheder al. 2019 Chen & Zhang 2019). 3D scenes (Mildenhall et al. 2021), video (Li ct al., 2021), audio (Sitzmann et al. 2020), etc. Specifically, in the context of images, INRs are neural networks trained on single image, which given a pixel location (z. y) return the value of that pixel in the training image. In this work, we classify the class of a training image an INR network was trained on, using only the INR's weights.
BACKGROUND
Definition: Weight Classification. The learner receives as input a set of n models f2. fn
where each has a corresponding label y1 ./....... Yn- Each model takes as input some tensor æ and outputs (1x), where the output can be vector of logits, probabilities or other variables. Each model fully specified by its weights and architecture. The task of the learner to train a classifier C, that takes as input the model f and predicts its label y.
The Challenge. While a naive solution would be to apply a standard neural architecture directly to the weights, this idea encounters serious setbacks. The dimension of the weight vector is very high, but more fundamentally, the ordering of the neurons in each layer is arbitrary. Hence, different permutations of the neurons result in functionally identical models. Standard architectures such as MLPs and CNNs do not respect such symmetries. One popular approach by Unterthiner et al.) 2020) computes permutation-invariint statistics for the flattened weights and biases of each layer, then trains standard classifier on them. These statistics lose much of the information contained in the weights, limiting the potential of this approach. Another direction is learning with equivariant architectures (Navon et al. al. [2023a] Zhou et al. 2024a) which respect the permutation invariance, e.g.. graph neural networks Kofinas et al. (2024); Lim et al. (2023). However, these methods (Kofinas) et al. 2024) treat each neuron of the model as token, and scaling them to large architectures is challenging (see Sec. 5.I.
Probing. Probing represents a model by running it several fixed inputs and noting the responses received on them. The learner can then train a classifier to map the model responses to the label. This approach avoids the issue of weight permutation invariance both the orders of input dimensions (c.g., image pixels) and output dimensions (c.g., class logits) are consistent across models.
Assuming that we want to predict attribute y of a network f. probing methods (Herrmann al.) 2024 Kofinas et al. 2024) optimize a set probes (PI...,,kk) and feed them into the network. They then train classifier C on the concatenation of the outputs. The prediction ÿ is:
i = C(F(p),, f(kk)
Probing methods learn the parameters of each probe p directly by latent optimization (Bojanowski et al. 2017). Each probe provides some information about the model attributes, and learning diverse and discriminative probes is key for obtaining useful representation. The classifier C' leverages information from all probes, and is typically trained by cross-entropy for classification and mean squared error for regression.
4 WEIGHT SPACE LEARNING WITH DEEP LINEAR PROBE GENERATORS
Our initial hypothesis is that probing methods, when done right, hold significant potential. Much like binary code files, neural networks are unknown and highly complex functions. Drawing inspiration
3

---

Published as a conference paper at ICLR 2025
Table Simple Probing vs. Other Approackes. We compare simple probing approach previous graph based and mechanistic approaches. We average the results over 5 different seeds. For probing, we experiment with different numbers of probes (in brackets).
seeds.
Accuracy
Kendall's s (1)
Method
MNIST
FMNIST
CIFARIO-GS
CIFARIO 10 Wild Park
StatNN (0)
0.398 =0.001
0.418
+00002
0,914 =0.000
0.719 #0.010
Neural Graphs (0)
0.923 10.003
0.727
1.0006
0.935 10.000
0.817 0.0007
Neural Graphs (64)
0.967 000002
0.736
:00012
0,938
10.00
0,888 000009
Neural Graphs (128)
0.976 00.001
0.745
100008
0.938
#0.000
0.885 0.0005
Vanilla Probing (64)
0.873
+2.026
0.784
+0.017
0.933 #0.001
0.885 =
Vanilla Probing (128)
0.955
#0.000
0.808
=:0.000
0.936
00.001
0.889
000008
Figure 2: A Few Examples from the Dead Leaves Dataset. We show these images are synthetic and highly dissimilar to real images.
from binary code analysis, where dynamic approaches Egele et al. (2008): Bayer al. (2006) are more common than static ones Shoshiittta et al. (2016), we believe that running neural networks, i.e. probing, is a promising approach for weight space learning. We begin with preliminary experiments to test the quality and potential of probing approaches.
4.1 A SIMPLE PROBING BASELINE
As we believe probing should be an effective way of analyzing neural networks, we begin by testing its raw capabilities. We take a vanilla probing approach, without any enhancements or modifica- tions, that optimizes all probes Pl, Pt..,,..,, with latent optimization (i.e., optimizing their values directly) and uses a simple MLP classifier for C. We compare this vanilla probing to its top competi- tors. First, graph based approach (Neural Graphs) (Kofinas et al. 2024) which treats each neuron of the network as node and operates a transformer on the resulting computational graph. Second, weight statistics (StatNN) (Unterthiner et al. 2020) which extracts simple statistics from the flat- tened weights and biases of a network and trains a simple predictor over them. We test all approaches on popular benchmarks. For dataset classification, we measure the accuracy for MNIST LeCun iet al. (1998) INRs digit prediction and Fashion-MNIST Xiao et al. (0117) INRs class prediction, both provided by Navon al. (2023a). For generalization error prediction, (Unterthiner et al. (2200) and CIFARIO SEAARRO--ild-PPP (0--Wild-Park (Kofinas et al. 2024) benchmarks, measuring the Kendall's metric as common in weight-space learning evaluation. The Kendall's 7 is a statistic used to measure the correlation agreement between two rankings, where 1 indicates perfect correla- tion, and indicates perfect negative correlation. The results are presented in Tab. T It is clear that (i) with enough probes, vanilla probing is able to perform better than graph based approach that does not use probing. (ii) Graph based approaches become comparable to vanilla probing only when incorporating probing features. This shows the promise of probing methods. Additionally, in Sec. [.I we demonstrate that probing also requires much less computational resources than graph based methods.
use the CIFAR
4.2 LEARNED VS. UNLEARNED PROBES
Having shown the merit in probing approaches, it is interesting to understand the quality of the probes themselves. To do so, we replace the learned probes by a set of unlearedd inputs, training only the classifier C. We test both: (i) probes with knowledge about the training data, selected from random synthetic data, and (ii) probes from in-distribution data, selected from the training set of the networks. We fix the number of probes k in all cases. In the MNIST and FMNIST INRs
4

---

Published as a conference paper at ICLR 2025
Table 2: Learned vs. Out of Distribution Probes. Comparison of learned probes, in-distribution
data probes and probes from randomly selected data. We average the results over 5 different seeds.
Accuracy
Kendall's T (t)
# Probes
Method
MNIST
FMNIST
CIFARIO-GS CIFARIO Wild Park
Learned Probes
0.873 #0.026
0.784 #0.017
0.933
=0.001
0,885 +0.008
64
Synthetic Probes
0.899
#0.022
0.832
#0.010
0.918
=0.001
0,882 #0.007
In-Dist. Probes
0.899
1:0.022
0.832 1:00100
0.939 20.000
0.922 1.0008
Learned Probes
0.955 0.0005
0.808
1:0.006
0.936
10.001
0.889 1.0.008
128
Synthetic Probes
0.959
000006
0.856
:00007
0.917
=0.001
0.893 20.014
In-Dist. Probes
0.959 000006
0.856
00.007
0.941
10.001
0.930
10.011
Qpay e5
Qvery 1M
OueyO
Deeyy 1ub
Qony
Qusyyt
Dany
Cosy
Cuayy
(a) Latent Optimization
(b) ProbeGen
Figure 3: Learned Probes. Optimized probes by (a) Latent Optimization and (b) ProbeGen for the CIFAR1O Wild Park benchmark. We show the same probes for identical runs (including the seed), except the generator module.
tasks, synthetic and in-distribution probes are the same, where we choose points from a uniform distribution between [[-1, 11,1,11111 1]]. In the CIFAR1O cases, select i random synthetic images from the Dead Leaves [Baradad Jurjo et al. 2021 Lee et al., 2001) dataset as synthetic probes, and k random images from CIFARIO in in-distribution probes. As seen in Fig. [ Dead Leaves images are synthetic, unlearned and highly dissimilar to real images. We therefore choose these images as our synthetic probes they include some structure, yet still far from being realistic. The results are presented in Tab. 2 It is clear that random probes are comparable to learned ones. We conclude that current probing techniques find suboptimal probes. To understand why these learned probes perform worse, we first observe them in Fig. 3a The probes show low-level, almost adversarial patterns, which can be highly expressive (Goodfellow et al.] 2014). We therefore hypothesize vanilla probing tends to overfit. In Sec. [.2 we demonstrate that our final method indeed overfits less than vanilla probing, and that more expressive methods hurt performance in these tasks.
4.3 DEEP LINEAR PROBE GENERATORS
We propose Deep Linear Probe Generators (ProbeGen) for learning better probes. ProbeGen op- timizes a deep generator module limited to linear expressivity, that shares information between the different probes. It then observes the responses from all probes, and trains an MLP classifier on them. While simple, in Sec. other approaches by a large margin.
Sec.5.1l [.lle show it greatly enhances probing methods, and also outperforms
Shared Deep Generator. Learning the probes through latent optimization prevents them from sharing useful patterns, as they do not have any shared parameters. A straightforward way for overcoming this is by factorizing each probe p; into two parts: (i) a latent code 2, learned using latent optimization and (ii) deep generator network G that all probes share. Formally,
Pt G(zs)
(2)
5

---

Published as a conference paper at ICLR 2025
where the latent code can have either higher lower dimension than p.- This factorization introduces a dependence between the probes as they all share G. It also reduces their expressivity, as G' may not be able to express all possible outputs.
Deep Linear Networks. Non-linear activations such as ReLU are the engine that makes deep networks very expressive, However, in this case, we would like to add regularization (see Sec. rather than increasing expressivity. We therefore remove the activations from our generator, keeping only stacked linear layers. Work by Arora et al. (2019) showed that when using SGD, deep linear networks have an implicit regularization effect. We therefore use deep linear networks in our approach, i.e., stacked linear layers but without the non-linear activations between them. In Sec. 5.2 we show that removing the activations reduces overfitting and therefore also enhances performance.
4.2
Data-Specfic Inductive Bias. In the case that our target model f takes structured inputs such as images, we hypothesize that presenting probes from a similar distribution will achieve higher accu- racy and reduce overfitting. As we do not know the data distribution, we cannot learn an accurate generative model for it. Still, example, we present generator for images, where we simply stack 2D convolutional layers, simi- larly DCGAN (Radford al., 2015). While this does not guarantee natural image statistics, it at least encourages some local structure and multi-scale hierarchy, both of which are some of the most important image characteristics.
introduce data-specific inductive biases into the generator. For
5 EXPERIMENTS
Our evaluation follows the standard protocol for weight space learning. We evaluate on two tasks: (i) CNNs generalization error prediction and (ii) detecting the training classes of images based on INR networks trained on them. We include experiments on small-scale established benchmarks as well as new larger-scale Model Zoo which we present, using ResNet 8HHe et al.) 2016) models.
Baselines. We compare our method, ProbeGen, with StatNN (Unterthiner et al. 2020 DWvoN et al. 2023a), NFN (Zhou et al. (024a).. ScaleGMN (Kalogeropoulos et al. 2024) and Neural Graphs Kofinas al. 2024). StatNN computes statisties for the weights and biases of each layer, concatenates them and trains a gradient boosted tree method on this representation. DWS (Navon et al. 2023a) and NFN (Zhou et al.) (0244) train linear, permutation equivariant layers with non-linearitie between them, but DWS is not applicable for CNNs. ScaleGMN (Kalogeropoulos al. 2024) treats neural networks as computational graphs, accounting for permutation symme- tries and scaling symmetries. Neural Graphs (Kofinas et al. 2024) is another graph-based approach, where each bias a node and the weights are the matching edges between these nodes. This method then trains transformer on the created graph, so the attention score between a pair of neurons de- pends on the weight that connects them. While Neural-Graphs do not account for scaling symmetries like the ScaleGMN baseline, it enriches each neuron's representation by using probing features.
Datasets. We evaluate 4 established datasets. For training data prediction we choose the MNIST and FMNIST implicit neural representation (INR) benchmarks (Navon et al. 2023a). Both datasets were formed by training an INR Sitzzaan al. (2020) model for each image of the origi- nal dataset. The goal is predicting the class of an image given its INR network. For generalization error prediction, we used the CIFAR 10-G0 (Unterthiner al., 2020 and CIFAR10 Wild Park tasks Koñnas et al. (2024). These datasets consists of thousands of small CNN models (3 5 layers), each trained separately on CIFARIO (Krizhevsky et al. 2009). We use accuracy for weight classification and Kendall's for regression.
Metrics. For INR class prediction, use simple accuracy to measure performance. For predict- ing the generalization error of CNNs, we follow standard evaluation protocol and use the Kendall's metric, which measures the agreement between two rankings. Similar to Pearson's 's correlation, with Kendall's t, values near 1 indicate strong positive correlation, values near -1 indicate strong negative correlation, and values near 0 indicate correlation.

---

Published as a conference paper at ICLR 2025
Table 3: Results for Small Scale Benchmarks. Comparison of ProbeGen, graph based, mecha- nistic approaches and latent optimized probes. We average the results over 5 different seeds.
Accuracy
Kendall's T (f)
# Probes
Method
MNIST
FMNIST
CIFARIO-GS
CIFARIO Wild Park
StatNN
0.398 +0.001
0,418
#0.002
0.914 +0.000
0.719 =0.010
DWS
0.857 #0.006
0.671 #0.003
0
NENHNP HNP
0.791 1:0008
0.689 100006
0,934
10.001
ScaleGMN
0,966 000002
0.808
0.0001
0,941 10.000
Neural Graphs
0.923 +0.003
0.727 20,006
0.935 10.000
0.817 00.007
Neural Graphs
0.967 #0.002
0.736 20.012
0,938 #0.001
0.888 00.009
64
Vanilla Probing
0.873 #0.026
0.784
#0.017
0.933
#0.001
0.885 20.008
ProbeGen
0,980
+0.001
0.861
#0.004
0,956 #0.000
0,933 =0.005
Neural Graphs
0.976 :00001
0.745
20.008
0.938
10.000
0.885 100005
128
Vanilla Probing
0.955 1:0.055
0.808
10.006
0.936
100001
0.889 1.0008
ProbeGen
0.984 #0.001
0.877
00.003
0.957 #00001
0.932 00.000
Table 5: Results for ResNet Scale. ProbeGen can successfully scale larger sized architec- tures, and outperforms both baselines.
Table 4: FLOPs Comparison. Using 128 probes and batch size of 64. ProbeGen much more efficient than graph approaches.
Kendall's (t)
Billion FLOPS (L)
# Probes
Method
ResNet18 Zoo
Method
MNIST CIFARIO-GS
StatNN
0.856
Neural Graphs
63.40
94.56
128
Vanilla Probing
0,843
ProbeGen
0.02
3.41
128
ProbeGen
0.910
5.1
MAIN RESULTS
Small-Scale Performance. Tab. 3 summarizes the results on the standard benchmarks. On INR dataset prediction tasks, ProbeGen outperforms all benchmarks significantly, even when the base- lines use the same number of probes. E.g.. in the FMNIST class prediction task. ProbeGen out- performs all other approaches by more than 6%, reaching an accuracy of 87.7%. ProbeGen also outperforms the baselines on predicting the generalization error of CNNs. While vanilla probing performs similarly to graph approaches quuality the probes, and achieves the highest result. On CIFARIO Wild Park ProbeGen also outperforms all previous approaches significantly, improving Kendall's from 0.889 to 0.933.
CIFARIO-GS, ProbeGen is clearly able to improve the
We compare the computational cost of ProbeGen and Neural- GraaalGraphs in terms of floating point operations (FLOPs). We test the MNIST INRs and CIFAR 10- D-GS datasets, with 128 probes and a batch size 64 for both approaches in Tab. We see that probing is a much more efficient approach, requiring between 1.5 orders of magnitude fewer FLOPs.
Computational Cost.
Large models. Knowing that probing methods are more efficient, we now wish to test probing methods can really scale to larger size models. We create new dataset of over 6, 000 ResNet18 He et al. (2016) models. Each ResNet model was trained on a randomly selected subset of Tiny- Imagenet Le & Yang (2015); Deng et al. (2009). We sampled the subset out of closed list of 10 subsets, that we created in advance. For each model, we record its generalization (test) error, and the objective is to predict this emror given the model weights. Graph based methods were too computationally expensive in this scale. Tab. (hshows the results. ProbGen achieves the best results reaching 0.91 kendall's 7, significantly outperforming both vanilla probing and statNN.

---

Published as a conference paper at ICLR 2025
28
Vartla Proting
0.58
Generaltazion Cap (x:10-)
Coreaceee Generator
Probe
22
Kendall's
2
18
0.9
Varila Proting
1.6
Cormictd
Generator
14
0.89
ProceGen
Layers
Layers
Figure 4: Overfitting for Different Generators. We compare ProbeGen to linear generators with fully connected layers, and non-linear convolu- tional generators. Overfit is measured by the generalization gap of each method. Results are averaged over 5 seeds.
Figure Ablation Studies. We compare the performance of ProbeGen, linear generators with fully connected layers, and non-linear con- volutional generators, for varying numbers of layers. With less than 2, all versions are equiva- lent ProbeGen. We average over seeds.
Table 6: Number of Probes. We compare ProbeGen to vanilla probing with different numbers of learned probes. We average over 5 different seeds.
FMNIST (1)
CIFARIO Wild Park (1)
# Probes
Vanilla
ProbeGen
Vanilla
ProbeGen
16
0.686 +0.043
0,805 #0.006
0,882 =0.010
0,926 #0.007
32
0.764 +0.013
0.844 +0.000
0.884 =0.008
0.930 #0.005
64
0.784 10.017
0.861
1:0.004
0.885 00000
0.933 1:0.00
128
0.808 1:0.006
0.877
1:0003
0.889
20.008
0.932 10.006
5.2 ABLATION STUDIES
Linear Generators. By observing Fig. 3 we hypothesized that removing the activations between the linear layers of the generator will reduce its overfitting. We therefore compare the results with and without non-linear activations. Indeed, seen in Fig. using non-linear activations results in a higher generalization gap, i.e., more overfitting. The amount overfitting even worse when not using a generator at all. In Fig. 5 we see that these activations also harm the model's performance, suggesting they lead to probes with reduced generalization abilities. In App. A we provide ablations demonstrating that, for these tasks, deep linear generator outperforms other regularizations.
Structure of Probes. First, we show several probes optimized via latent optimization and the same ones when optimized by ProbeGen (see Fig- 3). Although both not interpetable by humans, it is clear that ProbeGen probes have much more structure than latent-optimized ones. In Fig. 4 we demonstrate this has a regularizing effect, as ProbeGen significantly reduces the generalization gap compared to vanilla probing.
Additionally, we also test the contribution of the local inductive bias. By replacing the convolutional layers in generator to fully-connected ones, we completely remove this bias from the generator. Comparing the fully connected generator to our original ProbeGen helps isolate the effect of the local bias, as: (i) Both versions linear, meaning the fully-connected model is less restricted, and (ii) the size of the feature maps is kept similar throughout the generation process in both versions. Therefore, the primary difference is the inductive bias introduced by the convolutional layers. We compare the results in Fig. 5 for different numbers of layers. Indeed, we see the generators inductive bias important, as ProbeGen consistently outperforms the fully connected version.
Number of Probes. We compare ProbeGen to vanilla probing for differing numbers of probes. The results in Tab. 6 show that ProbeGen significantly improves over vanilla probing, even when using only fraction of the number of probes. In the FMNIST case, ProbeGen with 32 probes
8

---

Published as a conference paper at ICLR 2025
Digit
Digtt
Digit
True Digit
Digit
Dight
Thue Dugtt
Traa Dight
Trun
Thun Digt
Thue Digi
Thue Out
Thue Dight
Thue Digt4
Thue Digh
Trus Digit
True Dot3
True Dit
True Digit
True Digt
Digt?
True Digit
True Digh S
Trus Digit
Thue Digit
True Digit
True Desi
Trun Digt
True Dig
Thus Digt
Thus Dit
ProbeGen
(b) Vanilla Probing
Figure 6: MNIST Representations Extracted by ProbeGen vs. Vanilla Probing, We place the INR prediction in the probe locations. Other pixels are in gray.
already surpasses vanilla probing with 128 probes by almost 4% accuracy. In the CIFAR 10 case 16 ProheGen probes already significantly better than 128 latent optimized ones.
Representation Interpretability. ProbeGen represents each model as an ordered list of output values based on carefully chosen probes. These representations often have semantic meanings as the output space of the model (here, image pixels or logits) are semantic by design. For the MNIST INRs dataset, we visualize the inputs and outputs together, showing the partial image created by the probes. Fig- [ displays several representations for a few randomly selected images, comparing ProbeGen with vanilla probing. Vanilla probing chooses locations scattered around the image, in- cluding pixels far out of the image, where the behaviour of the INR is unexpected. ProbeGen on the other hand, chooses object centric locations, as suitable for this task. Indeed, one can easily identify the digits in the images despite only observing less than 20% of their pixels, hinting that this probe selection simplifies the task for the classifier module.
In Fig. 7 we visualize the representation (i.e., logits) of ProbeGen when trained the CIFAR10 Wild Park dataset. We can see that the values become more uniform as the accuracy of the models decreases, and sharper as it increases. This suggests that ProbeGen
some form of prediction entropy its classifier. We further test this hypothesis by training a classifier that only takes the entropy of each probe its features. We find that while not effective as ProbeGen, this classifier is still able to achieve a Kendall's r of 0.877. Hence, even when taken alone, prediction entropy is in fact highly discriminative for this task.
DISCUSSION
Black-box settings. We showed the potential of probing for learning from models. A side benefit of probing, compared to other approaches, is that it is suitable for inference on black-box models without any further adjustments. As probing only observes models responses, inference time it could simply probe a black-box model by its API, then make prediction based on the outputs.
Other modalities. ProbeGen incorporates inductive biases into the probe generator module. We tested it on images, showing its potential when having the right inductive bias, however, other modal- ities could require different structural biases. E.g.. in audio, consistency term would probably be helpful to simulate realistic recordings, and textual data may even require some pre-training for lin- guistic priors of the probes. App. we present a preliminary experiment using INRs trained on point clouds, demonstrating that ProbeGen's success generalizes other modalities well.
Adaptive probing. One interesting future direction for improving probing is being able to adap- tively choose the probes used to test each model. Probing models than lear policy for this adaptive selection, which would potentially improve accuracy for given certain probe budget. An early work by Herrmann et al. (2024) showed this could be effective for sequential models.

---

Published as a conference paper at ICLR 2025
True 0,200 Pred. 0.212
True 0.336 Prod. 0.335
True 0.464 Pred. 0.404
True 0.483 Pred. 0.472
Probability
True 0,501 Pred. 0.507
True 0.512 Pred. 0.519
True 0.512 Pred. 0,519
True 0.514 |I Pred. D. 506
0.8
True 0.532 Pred. 0.526
True 0.533 Pred. 0.533
True 0.534 Pred. 0.539
True 0.534 Pred. D.
0.6
True 0.535 || Prod.
True 0.537 Prod.
True Pred, 0,558
0.554
0.4
True 0,561 || Pred, 0.583
True D. 5.564 Pred, 0.580
True 0.592 Pred. 0,594
True 0.607 || Prod. D. 6.60
True 6.690 Pred. 0,682
True D. 721 Pred. 0.725
True 0 722 Pred. 0.751
True 0. 0.727 Pred. 0 7.726
Figure 7: CIFARIO Wild Park Representations Extracted by ProbeGen's Learned Probes. Each representation includes the 10 predicted probabilities (rows) given ProbeGen's 's 128 probes (columns), and above it are its true and predicted accuracies. Samples are sorted by true accuracy.
7 LIMITATIONS
Output space structure. While we demonstrated probing is a powerful tool for learning from neural networks, it requires the input and output dimensions to retain the same meaning across mod- els. There are important cases that do not satisfy this requirement, e.g., a repository of classification models with different classes in each model. Here, the output space for each varies and even if models share the same classes their order may still differ. Extending probing method to deal with the above cases is important direction for future work.
Weight generative tasks. Probing only looks at the input and output layers of each mode. There- fore, it cannot be used to give layer weight level predictions. That means that it is not suitable for weight generation tasks, such as editing or creating new neural networks.
Scalability. In Sec. 5. we we showed probing much more computationally efficient than previous graph based approaches. Still, it requires forwarding the entire model few times and computing the gradients through the model, with multiple models in each batch. This means that in order to infer about model using probing, one would need computational resources equivalent to training such a model. This would require a non-trivial solution for learning from larger models, e.g. CLIP (Radford et al. 2021) Stable Diffusion (Rombach et al. 2022). under limited compute budget.
CONCLUSIONS
This paper championed probing methods for weight space learning and improved them to achieve better than state-of-the-art performance. We first showed that vanilla probing approach. based on latent optimization, outperforms previous methods. However, we found that the learned probes are no better than randomly sampled synthetic data. To learn better probes, we proposed deep linear generator networks that significantly reduce overfitting through combination of implicit regulariza- tion and data-specific inductive bias. Beside consistently achieving the highest performance, often by large margin, our method requires 30 to 1000 fewer FLOPs than other top methods.
SOCIAL IMPACT
The goal of this paper is to further our understanding of neural networks. We expect this increased understanding to be helpful for reducing the social risks of Al such as bias and model safety.
10

---

Published as a conference paper at ICLR 2025
10 ACKNOWLEDGMENTS
This work was partially supported by the Israel Science Foundation (ISF), the Council for Higher Education (Vatat), the Center for Interdisciplinary Data Science Research (CIDR), the Israeli Cyber Authority, and KLA.
11 REPRODUCABILLITY
In this work, we presented new and light- framework for weight space learning. Our method simple to implement, and can be easily reproduced. To encourage future work in this direction, we provide short implementation our method the supplementary materials.
REFERENCES
Samuel K Ainsworth, Jonathan Hayase, and Siddhartha Srinivasa. Git re-basin: Merging models modulo permutation symmetries. arXiv preprint vrXiv:2209.04833 2022.
Sanjeev Arora, Nadav Cohen, Wei Hu, and Yuping Luo. Implicit regularization in deep matrix factorization. Advances in Neural Information Processing Systems, 32, 2019.
Maor Ashkenazi, Zohar Rimon, Ron Vainshtein, Shir Levi, Elad Richardson, Pinchas Mintz, and Eran Treister. Nern-learning neural representations for neural networks. arXiv preprint 12Xiv:2212.1355: 2022.
Manel Baradad Jurjo, Jonas Wulff, Tongzhou Wang, Phillip Isola, and Antonio Torralba. Learning to see by looking at noise. Advances in Neural Information Processing Systems, 34:2556-2569, 2021.
Uirich Bayer, Andreas Moser, Christopher Kruegel, and Engin Kirda. Dynamic analysis of mali- cious code. Journal in Computer Virology, 2:67-77, 2006.
Piotr Bojanowski, Armand Joulin, David Lopez-Paz, and Arthur Szlam. Optimizing the latent space 2017.
of generative networks. arXiv preprint
Nicholas Carlini, Daniel Paleka, Krishnamurthy Dj Dvijotham, Thomas Steinke, Jonathan Hayase, A Feder Cooper, Katherine Lee, Matthew Jagielski, Milad Nasr, Arthur Conmy, et al. Stealing part of a production language model. arXiv preprint arXiv:2403.06634 2024.
Angel X Chang, Thomas Funkhouser, Leonidas Guibas, Pat Hanrahan, Qixing Huang, Zimo Li, Silvio Savarese, Manolis Savva, Shuran Song, Hao Su, et al. Shapenet: An information-rice 3d model repository. arXiv preprint arXiv: 151.00302, 2015.
Zhiqin Chen and Hao Zhang. Learning implicit fields for generative shape modeling. In Proceedings of the JEEE/VV conference on computer vision and pattern recognition, pp. 5939-5948, 2019.
Luca De Luigi, Adriano Cardace, Riccardo Spezialetti, Pierluigi Zama Ramirez, Samuele Salti, and Luigi Di Stefano. Deep learning on implicit neural representations of shapes. arXiv preprint arXiv:2302.054338 2023.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large- scale hi- erarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, pp. 248-255. Ieee, 2009.
Cameron Diao and Ricky Loynd. Relational attention: Generalizing transformers for graph- structured tasks. arXiv preprint arXiv:2110005022 2022.
Amil Dravid, Yossi Gandelsman, Kuan-Chiel Wang. Rameen Abdal, Gordon Wetzstein, Alexei A Efros, and Kfir Aberman. Interpreting the weight space of customized diffusion models. arXiv preprint arXiv:2406. 4xXx:22406009113. 09413, 2024.
11

---

Published as a conference paper at ICLR 2025
Emilien Dupont, Hyunjik Kim, SM Eslami, Danilo Rezende, and Dan Rosenbaum. From data to functa: Your data point is a function and you can treat like one. arXiv preprint arXiv:2201.12204. 2022.
Manuel Egele, Theodoor Scholte, Engin Kirda, and Christopher Kruegel. A survey on automated dynamic malware-analysin techniques and tools. ACM computing surveys (CSUR), 44(2):1-42, 2008.
Gabriel Eilertsen, Daniel Jönsson, Timo Ropinski, Jonas Unger, and Anders Ynnerman. Classifying the classifier: dissecting the weight space of neural networks. In ECAI 2020, pp. 1119-1126. IOS Press, 2020.
Ziya Erkoç, Fangchang Ma, Qi Shan, Matthias Nießner, and Angela Dai. Hyperdiffusion: Generat- ing implicit neural fields with weight-space diffusion. In Proceedings of the IEEE/CVF interna- tional conference computer vision, pp. 14300-14310, 2023.
Justin Gilmer, Samuel Schoenholz, Patrick Riley, Oriol Vinyals, and George Dahl. Neural message passing for quantum chemistry. In International conference on machine learning, pp. 1263-1272. PMLR, 2017.
lan J Goodfellow, Jonathon Shlens, and Christian Szegedy. Explaining and harnessing adversarial examples. arXiv preprint arXiv:1412.6572, 2014.
David Ha. Generating large images from latent vectors. 2016.
David Ha, Andrew Dai, and Quoe V Le. Hypernetworks. arXiv preprint arXiv: 1rXii:1609009100 2016.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016.
Vincent Herrmann, Francesco Faccio, and Jürgen Schmidhuber. Learning useful representations of recurrent neural network weight matrices. arXiv preprint arXiv:2403. 2024.
Eliahu Horwitz, Jonathan Kahana, and Yedid Hoshen. Recovering the pre-fine-tuning weights of generative models. In ICML, 2024a. URL https:://ppeneeveeww.. wwtp:/reeeofrrrrmmeee
761Uxj0THB
Eliahu Horwitz, Asaf Shul, and Yedid Hoshen, On the origin of llamas: Model tree heritage recovery. arXiv preprint arXiv:2405. 2024b.
Ioannis Kalogeropoulos, Giorgos Bouritsas, and Yannis Panagakis. Scale equivariant graph metanet- works. arXiv preprint arXiv:2406. 106885 2024.
Thomas N Kipf and Max Welling, Semi-supervised classification with graph convolutional net- works. arXiv preprint arXiv: 1609000909 2016.
Miltiadis Kofinas, Boris Knyazev, Yan Zhang, Yunlu Chen, Gertjan Burghouts, Efstratios Gavves, Cees GM Snoek, and David W Zhang. Graph neural networks for learning equivariant represen- tations of neural networks. arXiv preprint arXiv:2403.1214: 2024.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
Ya and Xuan Yang. Tiny imagenet visual recognition challenge. CS 23/N, 7,7):3, 2015.
Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEEE, 86(11):2278-2 -2224,, 1998.
Ann B Lee, David Mumford, and Jinggang Huang. Occlusion models for natural images: A statis- tical study of a scale-invariaas dead leaves model. International Journal of Computer Vision, 41: 35-59, 2001.
12

---

Published as a conference paper at ICLR 2025
Zhengqi Li, Simon Niklaus, Noah Snavely, and Oliver Wang. Neural scene flow fields for space- time view synthesis of dynamic scenes. In Proceedings of the IEEE/CVE Conference on Computer Vision and Pattern Recognition, pp. 6498-6508, 2021.
Derek Lim, Haggai Maron, Marc T Law, Jonathan Lorraine, and James Lucas. Graph metanetworks for processing diverse neural architectures. arXiv preprint arXii:2332.04555 2023.
Derek Lim, Yoav Gelberg, Stefanie Jegelka, Haggai Maron, al. Learning on loras: GI- equivariant processing of low-rank weight spaces for large finetuned models. arXiv preprint 2oXh:22410.04207. 2024a.
Derek Lim, Theo Moe Putterman, Robin Walters, Haggai Maron, and Stefanie Jegelka. The empir- ical impact of neural parameter symmetries, or lack thereof. arXiv preprint arXiv:2405.202:: 2024b.
Lars Mescheder, Michael Oechsle, Michael Niemeyer, Sebastian Nowozin, and Andreas Geiger. Oc- cupancy networks: Learning 3d reconstruction in function space. In Proceedings of IÉÉE/CVF conference on computer vision and pattern recognition, pp. 4460-4470, 2019.
Ben Mildenhall, Pratul P Srinivasan, Matthew Tancik, Jonathan T Barron, Ravi Ramamoorthi, and Ren Ng. Nerf: Representing scenes as neural radiance fields for view synthesis. Communications of the ACM, 65(1):99-100) 106, 2021.
Aviv Navon, Aviv Shamsian, Idan Achituve, Ethan Fetaya, Gal Chechik, and Haggai Maron. Equiv- ariant architectures for learning in deep weight spaces. In International Conference on Machine Learning, pp. 25790-25816. PMLR, 2023a.
Aviv Navon, Aviv Shamsian, Ethan Fetaya, Gal Chechik, Nadav Dym, and Haggai Maron. Equiv- ariant deep weight space alignment. arXiv preprint arXiv:2310. /3397, 2023b.
Samuele Papa, Riccardo Valperga, David Knigge, Miltiadis Kofinas, Phillip Lippe. Jan-Jakob Sonke, and Efstratios Gavves. How to train neural field representations: A comprehensive study and benchmark. In Proceedings of the IEEE/CVE Conference on Computer Vision and Pattern Recog- nition, pp. 22616-22625, 2024.
William Peebles, Ilija Radosavovic, Tim Brooks, Alexei A Efros, and Jitendra Malik. Learning to learn with generative models of neural network checkpoints. arXiv preprint arXiv:2209. 12892, 2022.
Fidel A Guerrero Peña, Heitor Rapela Medeiros, Thomas Dubail, Masih Aminheidokhti, Eric Granger, and Marco Pedersoli. Re-basin via implicit sinkhorn differentiation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 20237-20246, 2023.
Alec Radford, Luke Metz, and Soumith Chintala. Unsupervised representation learning with deep convolutional generative adversarial networks. arxiv 2015. arXiv preprint arXiv:1511.0644 2015.
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In International conference on machine learning, pp- 8748 8763. PMLR, 2021.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEF/CVF confer- ence on computer vision and pattern recognition, pp. 10684-10695, 2022.
Mohammad Salama, Jonathan Kahana, Eliahu Horwitz, and Yedid Hoshen. Dataset size recovery from lora weights. arXiv preprint arXiv:2406. 19395, 2024.
Konstantin Schürholt, Dimche Kostadinov, and Damian Borth. Self-supervised representation leam- ing on neural network weights for model characteristic prediction. Advances in Neural Informa- fion Processing Systems, 34: 14:11641--2 (648--16493, 2021.
13

---

Published as a conference paper at ICLR 2025
Konstantin Schürholt, Diyar Taskiran, Boris Knyazev, Xavier Giró-i Nieto, and Damian Borth. Model zoos: A dataset of diverse populations of neural network models. Advances in Neural Information Processing Systems, 011588134-81444 2022.
Konstantin Schürholt, Michael W Mahoney, and Damian Borth. Towards scalable and versatile weight space learning. arXiv preprint arXiv:2406.099 099977 2024.
Viraj Shah, Nataniel Ruiz, Forrester Cole, Erika Lu, Svetlana Lazebnik, Yuanzhen Li, and Varun Jampani. Ziplora: Any subject in any style by effectively merging loras. arXiv preprint arXiv:2311.13600 2023.
Aviv Shamsian, Aviv Navon, David W Zhang, Yan Zhang, Ethan Fetaya, Gal Chechik, and Haggai Maron. Improved generalization weight space networks via augmentations. arXiv preprint arXiv:2402.04081, 2024.
Yan Shoshitaishvili, Ruoyu Wang. Christopher Salls. Nick Stephens, Mario Polino, Andrew Dutcher, John Grosen, Siji Feng. Christophe Hauser, Christopher Kruegel, et al. Sok:(state of) the art of war: Offensive techniques in binary analysis. In 2016 IEEE symposium on security and privacy (SP), pp. 138-157. IEEE, 2016.
Vincent Sitzmann, Julien Martel, Alexander Bergman, David Lindell, and Gordon Wetzstein. Im- plicit neural representations with periodic activation functions. Advances in neural information processing systems, 33:7462-7473, 2020.
Matthew Tancik, Pratul Srinivasan, Ben Mildenhall, Sara Fridovich-Keil, Nithin Raghavan, Utkarsh Singhal, Ravi Ramamoorthi, Jonathan Barron, and Ren Ng. Fourier features let networks learn high frequency functions in low dimensional domains. Advances in neural information processing systems, 33:7537-7547, 2020.
Viet-Hoang Tran, Thieu N Vo, An Nguyen The, Tho Tran Huu, Minh-Khoi Nguyen-Nhat, Thanh Tran, Duy-Tung Pham, and Tan Minh Nguyen. Equivariant neural functional networks for trans- formers. arXiv preprint arXiv:2410.04209, 2024.
Thomas Unterthiner, Daniel Keysers, Sylvain Gelly, Olivier Bousquet, and Ilya Tolstikhin. Predict- ing neural network accuracy from weights. arXiv preprint arXiv:2002. 1448, 2020.
A Vaswani. Attention is all you need. Advances in Neural Information Processing Systems, 2017.
Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-mnist: a novel image dataset for benchmark- 2017.
ing machine learning algorithms. arXiv preprint arXiv:
Allan Zhou, Kaien Yang. Kaylee Burns, Adriano Cardace, Yiding Jiang, Samuel Sokota, Zico Kolter, and Chelsea Finn. Permutation equivariant neural functionals. Advances in neural infor- mation processing systems, 36, 2024a.
Allan Zhou, Kaien Yang, Yiding Jiang, Kaylee Burns, Winnie Xu, Samuel Sokota, J Zico Kolter, and Chelsea Finn. Neural functional transformers. Advances in neural information processing systems, 36, 2024b.
14

---

Published as a conference paper at ICLR 2025
0.920
Accuracy
0.910
0.900
Vanilla Probing ProbeGen
0.890
0.880
500
1000
1500
2000
# of Probes
Figure 8: Ablating the Number of Probes used for Point Cloud INRs Classification. We show that ProbeGen exhibits much better scaling law when classifying point clouds, compared to Vanilla Probing. Notably, ProbeGen with 32 probes already outperforms Vanilla Probing with 2048 probes.
A OTHER REGULARIZATIONS
compare deep linear generators to a range of & weight decay regularization strengths. We present the results in Tab. 7 Our results show that while adding ful, is not as good as a deep linear generator, and requires aggressive hyper-parameter parameter tuning the regularization strength, including non-standard values (a weight decay of 10-7, 10-8), where deep linear generator does not. Moreover, these regularizations are not consistent between experiments, and there no single value that fits all experiments. We conclude that deep linear generator, as presented in ProbeGen, provides a stable and subtle regularization that fits weight-space analysis.
regularization might be help-
Table 7: Performance Comparison of Different Regularizations vs. ProbeGen.
Accuracy
Kendall's (f)
Method
MNIST
FMNIST
CIFARIO-GS
CIFAR10 10 Wild Park
Non-Linear Gen. + I.2 Reg. (10-2)
0.942
0.531
0.758
0.442
Non-Linear Gen. L. Reg. (10-3)
0.979
0.858
0.824
0.547
Non-Linear Gen. L2 Reg. (10-4)
0.981
0.867
0.861
0.737
Non-Linear Gen. + I.2 Reg. (10-5)
0.982
0.869
0.911
0.871
Non-Linear Gen. L.2 Reg. (10-6)
0.982
0.870
0.941
0.905
Non-Linear Gen. + L2 Reg. (10-7)
0.982
0.869
0.947
0.922
Non-Linear Gen. + L. Reg. (10-8)
0.981
0.870
0.948
0.928
ProbeGinn
0.984
0.877
0.957
0.932
B ANOTHER DATA MODALITY: POINT CLOUDS
We extend our evaluation to include a more complex data modality: point clouds. Specifically, we use the Neural-Field-Aree Arena (Papa et al. 2024). to evaluate ProbeGen's ability in classifying INRs trained on point clouds from the ShapeNet (Chang et al. 2015)) dataset. The goal is to detect the class of the point cloud each INR was trained on, directly from the INR's weights. As shown in Table 8 ProbeGen consistently outperforms all prior approaches, including Vanilla Probing.
Figure further examines the efficiency of probing methods with respect to the number probes that under highly restrictive budget of 32 probes, ProbeGen is able to surpass all baselines. Notably, this 32 probes version ProbeGen outperforms Vanilla Probing with 2048 probes, showing ProbeCien's 's better scaling laws.
used. We
15

---

Published as a conference paper at ICLR 2025
Table 8: Point Cloud INRs Classification. We show that ProbeGen can successfully learn from models trained on more complex data modelities, such as point clouds. Number of probes is in parentheses.
Accuracy
Method
Point Clouds INRs
DWS
0.911
Neural Graphs
0.903
Vanilla Probing (128)
0.913
ProbeGen
0.930
C IMPLEMENTATION DETAILS
Generator Architecture. For image generation, we use a transposed convolution based generator. With each layer, the feature maps spatial sizes doubles in each axis (and 4 times overall), while the number of neurons (channels) decrease by half. We choose a generator width multiplier of 16, i.c., our generators last layer has 16 input channels, and the channels multiply by 2 with each previous layer. Our image generators has 6 transposed convolutional operators for experiments on the CIFARIO-GS and ResNet18 cases, and 5 for experiments on the CIFARIO Wild Park dataset.
For our INR coordinates generators we use Fully-con Fully-connected layers, with a hidden size of 32. These generators use fully- conlecco layers for both FMNIST INRs and MNIST INRs.
Hyper-parameters. We use learning rate of 3- -4 and a batch size of 32 in all our experiments. Our MLP classifier C, uses layers with a hidden size of 256. The latent vectors each probe are of size 32. We trained all probing algorithms for 30 epochs, all experiments on the CIFAR 10-GS dataset for 150 epochs, and all experiments on our ResNet18 Model Zoo dataset for 100 epochs. Additionally, we trained all baselines using their default hyper-papn parameters.
the INR and CIFARIO Wild Park experiments
(a) Vanilla Probing (Latent Optimization)
Single Layer Generator
Demy
(c) 3 Layers Generator
(d) ProbeGen (5 Layers)
Figure 9: Learned Queries by Deep Linear Generators. Queries learned by deep linear generators and vanilla probing, for the CIFAR10 Wild Park dataset.
16

---

Published as a conference paper at ICLR 2025
Fully-Connected Image Generators. We find that optimizing fully connected generators with the exact same dimensions as ProbeGen leads to sever overfitting. Therefore, instead of having C' x H W hidden dimensions (where H and W are the spatial sizes created by the convolutions) in each layer, we choose x H x W which empirically worked much better.
D QUERIES OF DIFFERENT PROBING APPROACHES
We qualitatively compare the learned queries by different probing algorithms. We visualize the same random subset of the queries from each algorithm, learned for the CIFARIO Wild Park dataset.
In Fig. we provide queries learned by our ProbeGen using different numbers of layers. We see the queries gradually develop structure the number of layers increases. This goes in line with our hypothesis from Sec. 4.3
Next, provide visualization of the queries of ProbeGen with activations between its linear layers. These queries are provided in Fig. [00 We see more repetitive structure in these queries than in the standard ProbeGen, indicating the information may reside to more local patterns, which are not necessarily object centric as CIFARIO tends to be.
Finally, we observe the queries from fully-connecter generator. Presented in Fig. these queries show very little structure even when the generator have 5 layers. The structure is of local patterns, as there are no convolutions to present hierarchical order. This shows the convolutional operators indeed enforce the desired structure on its queries.
Query
Ouery 311
Guery
Query
Ouery 200
0L.07
Cecny 107
Detrr 10
Qurry de
Overs 106
Quey
Qaty
Quiy
QPay
Quiyy 77
Qutyy
Qaty 93
Qety
(a) 2 Layers with Activations
(b) 3 Layers with Activations
Overy
Gary
Seny 1t1
77
Ovty 122
Sway it3
Qury
Query
Quary 55
(c) Layers with Activations
(d) ProbeGen (5 Layers) with Activations
Figure 10: Learned Queries by Non-Linear Generators. Queries learned by non-linear generators for the CIFARIO Wild Park dataset.
17

---

Published as a conference paper at ICLR 2025
Quar
Gurn 230
Colry 126
Quar
Seary
(a) 2 Fully-Coonn Connected Layers
(b) 3 Fully- Connected Layers
(c) 4 Fully-Connected Layers
(d) 5 Fully-Conneeted Layers
Figure 11: Learned Queries by Tally-Connected Linear Generators. Queries learned by fully- connected linear generators for the CIFAR10 Wild Park dataset.
18


Reviews:
[{"review_id": "MnJH742B3K", "reviewer": "Reviewer_wSNk", "summary": "In this paper, the authors introduce ProbeGen, a deep linear method designed for probing model data in weight space learning. They begin by demonstrating empirically that probing is an effective approach for weight space learning, although it has limitations and room for improvement. Building on these findings, they propose ProbeGen as an enhancement and subsequently validate its effectiveness through a series of experiments.", "questions": "- Do the authors have a hypothesis as to why non-linear activation functions lead to poorer probe representations? Have the authors experimented with other non-linear activations besides ReLU activations discussed in lines 255-261?\n\n- It would be interesting to try and enrich the probes by applying augmentations to the input like presented in [1,4].", "limitations": "", "rating": 6, "confidence": 5, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- This work addresses the novel and intriguing field of weight space learning, aiming to tackle the significant challenge of simplifying optimization within this domain.\n- The proposed method ProbeGen is simple, easy to follow, and yet effective.\n- The authors empirically demonstrate the motivation behind ProbeGen, providing interesting insights.\n- The paper is well-written and easy to follow.", "weaknesses": "- The experimental part is limited. From an empirical-oriented paper, I expect the experimental section to be more comprehensive. For instance, the effectiveness of ProbeGen is shown only on INRs overlooking equivariant tasks in weight space learning like domain adaptation [1], alignment [2], editing [3], and more.\n- The authors mentioned that ProbeGen could be used for processing high-dimensional inputs, however, the experimental section focuses on small-scale INRs instead of modern architectures. \n- It would be interesting to see what is the upper bound of ProbeGen w.r.t number of probes. The authors used a max number of 128 probes what will be the performance if we will use 1024 probes for example? Can we match the image space performance?\n- The authors failed to cite relevant works in the field of weight space learning [2,4,5,6].\n- Not sure if it is intended but the supplementary material is missing.\n\n------------\n\n[1] Equivariant Architectures for Learning in Deep Weight Spaces, Navon et al.\n\n[2] Equivariant Deep Weight Space Alignment, Navon et al.\n\n[3] Permutation equivariant neural functionals, Zhou et al.\n\n[4] Improved Generalization of Weight Space Networks via Augmentations, Shamsian et al.\n\n[5] Equivariant Neural Functional Networks for Transformers, Tran et al.\n\n[6] Learning on LoRAs: GL-Equivariant Processing of Low-Rank Weight Spaces for Large Finetuned Models, Putterman et al.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Giqt6mGurv", "reviewer": "Reviewer_PoPY", "summary": "This work addresses the task of learning in weight spaces, and introduces a probing-based method that passes a set of learnable inputs to a model and trains a simple network on top of the concatenated probe outputs. The authors first demonstrate that standalone vanilla probing performs very well, almost on par to state-of-the-art methods. They then extend the method by introducing a shared deep linear network module, and incorporate inductive biases towards structured probes. This module further increases performance by reducing overfitting, while being more efficient than other approaches.", "questions": "[1] It seems intuitive that probing approaches are complementary to mechanistic approaches. Can the proposed method be combined with existing mechanistic approaches (e.g. Neural Graphs)? Does ProbeGen increase the performance of these methods? An ablation study on whether deep linear generators with structured input probes boost mechanistic approaches would be very useful, and would further increase the impact of the proposed method.\n\n[2] The experiments on learned vs unlearned probes are very insightful. It is unclear, though, why the authors did not explore actual unlearned images as probes, potentially in-distribution. For example, a small subset of the training (or validation) set used only for probing. What is the impact of using real images for probes?\n\n[3] Deep linear networks are introduced as a means to introduce regularization and reduce overfitting. Have the authors explored other common regularization methods, e.g. adding weight decay, or dropout? Is the reduced expressivity of the network an advantage here? If so, it is not very clear why that is the case.\n\n[4] The authors often mention the terms \"training dataset classification\", \"extract information about its training dataset\", or variants. These terms seem quite confusing. Do they refer to INR classification tasks? The manuscript could benefit from some rephrasing in these instances.\n\n[5] As the number of probes increases, it approaches the size of the image for INRs, or the test set size for CNN generalization.\nIn that case, probing in ProbeGen becomes almost trivial, since it is nearly identical to feeding the whole image to an MLP, or to running the CNN over the whole test set. In both cases, we can expect the performance monotonically increasing as a function of the number of probes. Is that assumption correct? Is there a sweet spot between efficiency and performance? An ablation study on MNIST/Fashion MNIST with increasing numbers of probes (up to the image size) that shows performance and FLOPS would be very useful here.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper is well written and is quite easy to follow. The connections with binary code analysis and the corresponding insights are very useful towards more scalable weight-space methods. \n\nThe proposed method reaches state-of-the-art performance with high computational efficiency. The authors provide a comprehensive set of ablation studies that offers significant insights on probing-based methods for weight space learning.", "weaknesses": "The proposed method cannot be used in weight-level tasks, such as editing neural network weights, which represents a large class of tasks in learning in weight spaces. The impact of this work would greatly increase if the authors include a more detailed discussion on why that is the case, and potential directions to alleviate this limitation.\n\nIt is unclear if the proposed method can be used to complement existing mechanistic approaches and result in a performance greater than either of the two components. See question 1.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "uTYJzAAWZo", "reviewer": "Reviewer_4rBX", "summary": "This paper advocates for probing methods to solve tasks on the weight space of NNs. It first shows that naive probing performs as well as SoTA methods that operate directly on the weight space in discriminative tasks. Then, by analyzing the learned probs, the authors suggest an improved mechanism for learning the probs using deep linear networks which adds implicit regularization. The authors demonstrate the merits of their approach compared to baseline methods on two tasks, INR classification and image prediction errors. Several ablations are done to inspect different aspects of the method.", "questions": ".", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "* The idea of using a deep linear network to obtain implicit regularization while guarding against potential overfitting is nice and original in that context.\n* The proposed approach is extremely more efficient compared to methods that learn directly on the weights.\n* The paper is written clearly and adequately. The authors empirically justify their design choices every step of the way and show other alternatives (such as learned vs unlearned probs, and the architecture decisions about the generator).\n* The ablation study clears several ambiguities that arise while reading the paper, such as using non-linear generators and the importance of the inductive biases introduced into it.", "weaknesses": "* The authors discussed frankly about the main limitations of their approach, which I agree with. In particular, there are some tasks on which it is not immediately apparent how to use linear probing, such as generative tasks. It may be valuable if the authors could come up with possible remedies for that, but I understand if they don't since it is beyond the scope of this work and maybe non-trivial.\n\n* Experiments:\n  * Section 5.2 demonstrates the importance of deep linear classifiers compared to non-linear ones. The main issue, as I understand it, is overfitting. But, I wonder why didn't you try explicit regularization schemes to benefit from a more expressive network? I believe that an experiment of that flavor is missing.\n  * Given the main advantage of this approach (i.e., agnostic to symmetries and fast) and given its limited applicability, I believe the submission can be made much stronger by evaluating ProbeGen on complex data, and on bigger networks. Both are limitations of current models that operate on the weights directly. For instance, INR classification or generalization error of images from ImageNet.\n  * Related to my previous comment - I can see why on MNIST and even CIFAR small number of probs is sufficient to get good results. However, I speculate that it will not be the case for more challenging datasets (e.g., ImageNet or datasets of fine-grained classification such as CUB). Perhaps additional experiments can clarify that ambiguity.\n  * Exact experimental details are missing to reproduce the results, e.g., what optimizers were used? Did you apply regularization? What is the architecture of the classifier? Also, it is not clear whether data augmentation was used for baseline methods (as it tends to improve the results).\n\n* Some parts of the submission can be improved. Specifically, \n  * There should be some background on INRs and the tasks involving them.\n  * It is not clear how exactly Kendall’s $\\tau$ was used to evaluate the generalization error.\n  * It is not clear how exactly to run the code in the submission (there are no explanations about how to obtain the data, some files seem to be missing, etc.)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "0bxT5snAxL", "reviewer": "Reviewer_DBUB", "summary": "The paper considers the problem of weight space learning, in which a learner is tasked with predicting the properties (e.g., accuracy) of an input (trained) neural net. The paper makes the distinction between two approaches for weight space learning: (1) Mechanistic approaches, which operate on the input weights (without running the model), and (2) Probing approaches, which represent models by their responses to a set of inputs. The paper champions probing approaches and proposes a simple yet effective approach in which the probes (model inputs) are learned using a deep linear generator. On a set of small-scale benchmarks, the authors show the proposed method outperforms both probe-based and mechanistic baselines, with improved computational efficiency.", "questions": "- Have you tried comparing to using real images from the training dataset (for example, CIFAR10 images) as probes?\n- It is not clear if you also used Conv layers for the INR tasks. If so, why does this make sense?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- The paper is well-structured and easy to follow.\n- The method is well motivated through a set of empirical observations.\n- The proposed approach is simple yet effective.", "weaknesses": "My main concern lies with the experimental setup, which is limited to small-scale datasets and includes only two baselines. In addition, all experiments are performed on images with no additional modalities. This restricts the generalizability of the findings. I would expect a broader and more diverse set of experiments (in terms of scale and modalities) from a purely empirical paper. Please consider adding additional experiments with larger models, higher-dimensional inputs, or different modalities.\n\nAdditional limitations:\n1. The approach is not applicable for an important and large set of weight space tasks, namely weights-to-weights tasks, e.g., INR editing, domain adaptation, sparsification, etc.\n1. The provided empirical evaluation is limited to fairly low-dimensional input spaces. I assume probing approaches will be more efficient in low-dim spaces, so it will be beneficial to provide results for more challenging setups with higher-dimensional input spaces (e.g., point clouds) or input spaces with more challenging structures (e.g., graphs).\n1. Similarly, results for additional data modalities will be beneficial and will strengthen the empirical soundness.\n1. Also, including additional baselines like [1, 2] would be beneficial.\n1. While the authors analyze the computational cost, showing clear benefits compared to the GNN-based approach, it is not clear that the results will hold for larger models, as the approach requires forward and backward steps through the models.\n1. The Appendix is completely missing, along with additional important information regarding implementation details, experimental setup, HPs, etc.\n\nReferences:\n\n[1] Kalogeropoulos et al. \"Scale Equivariant Graph Metanetworks,\" NeurIPS 2024.\n\n[2] Schürholt et al. \"Towards Scalable and Versatile Weight Space Learning.\" ICML 2024.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)