Lorenzo Sciarretta

Neural network optimization / Weight-space learning

Can AI help us train
better neural networks?

From a new dataset to useful model weights.
Generate a starting model. Then fine-tune.

Suppose you have labeled images for a new defect-detection task and have chosen a CNN architecture. You could adapt a pretrained checkpoint or train from scratch with SGD or Adam. The goal is to generate useful starting weights from your dataset, then fine-tune them. That would let a new task benefit from patterns learned across existing networks.

Two approaches work directly in weight space, the space of network parameters. SANE (2024) learns representations of trained networks and decodes sampled embeddings into new weights. Weight-Space Programming (2026) explores general-purpose LLMs reading, editing, and writing network weights through reasoning.

Our method. SANE-MoE adds expert routing to the SANE autoencoder. Evaluated on mixed collections of small CNNs trained on CIFAR-10 and STL-10, it improves both weight reconstruction and generated-model accuracy over the dense baseline.

Lorenzo Sciarretta & Zakaria Omarar · AI/ML Lab, University of St. Gallen · IMP 2025 project

Read our report See how it works

Read further

The report, its foundations, and two complementary directions.

Our method · 2025 Sciarretta & OmararWeight Space Learning with Mixture of Experts (MoE) Adds expert routing to SANE; improves reconstruction and generation in mixed CNN zoos. Foundation · ICML 2024 Schürholt, Mahoney & BorthTowards Scalable and Versatile Weight Space Learning Introduces SANE: tokenizes network weights for scalable encoding, analysis, and model sampling. Our backbone. Model-zoo data · NeurIPS 2022 Schürholt et al.Model Zoos: A Dataset of Diverse Populations of Neural Network Models Provides standardized checkpoint collections and benchmarks for learning from populations of networks. Dataset conditioning · ICML 2026 Asefaw et al.WeightCLIP: Aligning Datasets and Models for Weight Space Learning Aligns dataset and weight embeddings for dataset-guided model retrieval, generation, and refinement. Related perspective · 2026 Takehi et al.Sparks of Weight-Space Programming Explores general-purpose LLMs reading and editing weights directly. A complementary approach to AI-assisted model optimization.
01 / Initialize

A learned head start

Sample candidate weights, evaluate them, then fine-tune with SGD or Adam on your target task.

02 / Generate

New model checkpoints

Decode samples from the learned embedding distribution into complete network weights.

03 / Understand

Read a network’s weights

Use embeddings as features to predict model properties, such as accuracy or generalization gap.

The idea in one picture

Trained models become the training data.

A model zoo is a collection of trained checkpoints. SANE learns to encode and reconstruct their weights, processing windows of weight tokens instead of a whole network at once.

01 Collect experience

Model zooWeight tokens

Networks → sequences

Standardize and align weights; split them into tokens with layer and position information.

02 Learn a representation

EDEncoderEmbeddingsDecoder

Encode → reconstruct

Train an autoencoder on token windows. Each token gets an embedding; a model has a sequence of them.

03 Put it to work

Generate weightsSample embeddings → decode
→ calibrate → evaluate → fine-tune
Predict propertiesEncode a checkpoint → fit a probe
→ estimate model characteristics

Generate or analyze

Reuse the decoder to produce networks, or the encoder to study existing ones.

Conceptual redesign of the SANE pipeline, after Schürholt et al. (2024), Fig. 2. The latent drawing is schematic; each network is represented by a sequence of token embeddings.

Where the optimizer fits. Training SANE still uses gradient optimization. To generate a network, fit a sampling distribution to checkpoint embeddings, decode new weights, calibrate batch-normalization statistics, and evaluate the candidates. Selected weights can then be fine-tuned with SGD or Adam.

Our contribution / SANE-MoE

Routing weight tokens to experts.

Weights from different layers and datasets need not follow the same patterns. We replace feed-forward sublayers in selected SANE Transformer blocks with a router and multiple expert networks. The router can use token content, position, and layer information to choose which experts process each weight token.

Inside an MoE sublayer

Sparse expert routing

The router selects experts for each token. Their outputs are combined using normalized routing weights.

SANE-MoE architecture, adapted from our report, Section 3 and Fig. 1. MoE replaces feed-forward sublayers in selected encoder and decoder blocks; residual connections and normalization are omitted. The zoom illustrates top-2 routing among four experts; active choices vary by token.

We compare top-k, position-aware, and hierarchical routing. Sparse top-k routing increases total capacity while activating only a subset of experts for each token; computation depends on the active expert count and routing overhead.

What we measured

Better generated checkpoints in a mixed zoo.

Trained on a mixture of small CNN zoos from CIFAR-10 and STL-10, every evaluated MoE routing variant improves mean generated-model accuracy over dense SANE.

+5.14 ppCIFAR-10 generationTop-k: 38.74% → 43.88%
+3.98 ppSTL-10 generationHierarchical: 28.90% → 32.88%
0.810 → 0.921Reconstruction R²Hierarchical routing · validation
Dense SANE SANE-MoE variants Original CNN mean
CIFAR-10: generated-model accuracySANE baseline 38.74 percent; top-k 43.88; position-aware 43.07; hierarchical 42.41. Whiskers show one standard deviation. Original CNN reference 45.81 percent. image/svg+xml Matplotlib v3.10.8, https://matplotlib.org/ 0 10 20 30 40 50 Generated-model test accuracy (%) SANE baseline MoE · top-k MoE · position MoE · hierarchical 38.74% 43.88% 43.07% 42.41% CIFAR-10 Original CNN mean: 45.81%
STL-10: generated-model accuracySANE baseline 28.9 percent; top-k 32.13; position-aware 30.44; hierarchical 32.88. Whiskers show one standard deviation. Original CNN reference 34.93 percent. image/svg+xml Matplotlib v3.10.8, https://matplotlib.org/ 0 10 20 30 40 50 Generated-model test accuracy (%) SANE baseline MoE · top-k MoE · position MoE · hierarchical 28.90% 32.13% 30.44% 32.88% STL-10 Original CNN mean: 34.93%
Generated-model test accuracy after batch-normalization conditioning and before fine-tuning (report, Table 4). Bars show means; whiskers show ±1 standard deviation over the reported 150 sampled models. The dashed line marks original CNN accuracy. The gain summaries compare means; the best router differs by dataset.
Exact values from the report
Generated-model test accuracy (%) · Report, Table 4
MethodCIFAR-10STL-10
SANE baseline38.74 ± 4.1628.90 ± 5.43
MoE · top-k43.88 ± 3.0132.13 ± 4.38
MoE · position-aware43.07 ± 3.2030.44 ± 5.31
MoE · hierarchical42.41 ± 2.1932.88 ± 4.01
Original CNNs (reference)45.8134.93

Conclusion / Toward scalable weight-space learning

What improves, and what comes next.

Expert routing improves the SANE backbone in this mixed-CNN setting: reconstructed weights are more faithful to the originals, and generated models are more accurate. The next question is whether these gains hold for larger networks and broader model zoos.

Future direction · conceptual pipeline

Dataset + compatible architecturedataset-aligned weight spacegenerated NN weightstask-specific fine-tuning

WeightCLIP (2026) aligns dataset information with weight embeddings. Combining that conditioning with an expert-based backbone could support scalable generation of weights for a new dataset and compatible architecture. Our report does not test this combination or measure fine-tuning speed; both require further experiments.

← Back to selected work