Research map / Architectures and efficiency / RNNs and normalization
Deep Linear Probe Generators for Weight Space Learning
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes ProbeGen, which factorizes learned probes into per-probe latent codes passed through a shared deep linear generator, improving probing-based weight space learning while using far fewer FLOPs than graph-based methods.
Problem
Weight space learning: predicting properties of neural networks (training dataset, generalization error) directly from their weights; existing learned probes are no better than random synthetic probes and overfit.
Approach
Replace per-probe latent optimization with probes generated by a shared deep linear (no activation) generator, with convolutional structure for images, giving implicit regularization and an inductive bias toward structured probes.
Key results
- ProbeGen with 128 probes reaches 0.877 accuracy on FMNIST INRs, over 6 points above prior methods, and 0.933 Kendall's tau on CIFAR10 Wild Park (up from 0.889).
- ProbeGen requires roughly 30 to 1000 times fewer FLOPs than Neural Graphs (e.g., 0.02 vs 63.40 billion FLOPs on MNIST INRs).
- On the ResNet18 Model Zoo, ProbeGen achieves Kendall's tau 0.910 versus 0.856 for StatNN and 0.843 for vanilla probing.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace latent-optimized probes with a shared deep linear generator | ProbeGen · MNIST | accuracy: 0.955 -> 0.984 | Helped |
| Replace latent-optimized probes with a shared deep linear generator | ProbeGen · FMNIST | accuracy: 0.808 -> 0.877 | Helped |
| Replace latent-optimized probes with a shared deep linear generator | ProbeGen · CIFAR10-GS | Kendall's tau: 0.936 -> 0.957 | Helped |
| Replace latent-optimized probes with a shared deep linear generator | ProbeGen · CIFAR10 Wild Park | Kendall's tau: 0.889 -> 0.932 | Helped |
| Add a deep linear generator instead of vanilla probing | ProbeGen · ResNet18 Zoo | Kendall's tau: 0.843 -> 0.910 | Helped |
| Add a convolutional local inductive bias rather than a fully connected generator (ablation) | ProbeGen image generator · Image benchmarks | performance | Helped |
| Add non-linear activations between generator layers (ablation) | ProbeGen generator · Weight-space learning tasks | generalization | Hurt |
| Remove the generator and optimize probes directly (ablation) | ProbeGen · Weight-space learning tasks | generalization | Hurt |
| Replace the deep linear generator with non-linear generators using L2 weight decay (ablation) | ProbeGen · MNIST | accuracy: 0.984 vs. 0.982 | Hurt |
| Replace the deep linear generator with non-linear generators using L2 weight decay (ablation) | ProbeGen · FMNIST | accuracy: 0.877 vs. 0.870 | Hurt |
| Replace the deep linear generator with non-linear generators using L2 weight decay (ablation) | ProbeGen · CIFAR10-GS | Kendall's tau: 0.957 vs. 0.948 | Helped |
| Replace the deep linear generator with non-linear generators using L2 weight decay (ablation) | ProbeGen · CIFAR10 Wild Park | Kendall's tau: 0.932 vs. 0.928 | Helped |
Surprising findings
- Learned latent-optimized probes perform comparably to random synthetic or in-distribution probes, suggesting current probe learning is suboptimal.
- Adding non-linear activations to the generator degrades probe quality and increases overfitting, the opposite of the usual expressivity benefit.
Limitations
- Probing requires input and output dimensions with consistent meaning across models, so it fails for models with differing output class orders.
- Probing cannot be used for layer- or weight-level predictions such as weight generation or editing.
- Still requires forwarding and backpropagating through entire models, which may be costly for large models like CLIP or Stable Diffusion.
Experimental setup
Datasets: MNIST INRs FMNIST INRs CIFAR10-GS CIFAR10 Wild Park Tiny-ImageNet (ResNet18 Model Zoo) Neural-Field-Arena ShapeNet point cloud INRs Dead Leaves
Benchmarks: MNIST INRs FMNIST INRs CIFAR10-GS CIFAR10 Wild Park ResNet18 Model Zoo Neural-Field-Arena point clouds
Metrics: accuracy Kendall's tau FLOPs
Baselines: StatNN DWS NFN ScaleGMN Neural Graphs Vanilla Probing
Builds on: Neural Graphs StatNN DCGAN latent optimization
Closest papers by meaning
- Learning on Model Weights using Tree Experts (CVPR 2025 · similarity 0.84)
- Improved Generalization of Weight Space Networks via Augmentations (ICML 2024 · similarity 0.84)
- Scaling Effects on Latent Representation Edits in GPT Models (Student Abstract) (AAAI 2025 · similarity 0.83)
- Deep Minds and Shallow Probes (NeurIPS 2026 · similarity 0.82)
- RAPTOR: Ridge-Adaptive Logistic Probes (NeurIPS 2026 · similarity 0.82)
- Probing Neural Combinatorial Optimization Models (NeurIPS 2025 · similarity 0.81)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.