WeightCLIP: Aligning Datasets and Models for Weight Space Learning
2026-07-03 • Machine Learning
Machine Learning
AI summaryⓘ
The authors present WeightCLIP, a new method to better understand and use the internal settings (weights) of neural networks by linking these settings to the datasets the networks were trained on. They do this by encoding both the neural network weights and dataset examples into a shared space where related pairs are matched closely. This helps improve tasks like finding good network settings from dataset details and refining models to perform better than usual fine-tuning. Their approach shows that adding dataset information directly into weight representations can enhance several applications involving neural networks.
neural networksweight spacelatent representationautoencodercontrastive learningdataset encodingmodel fine-tuningrepresentation learning
Authors
Aron Asefaw, Konstantinos Tzevelekakis, Damian Falk, Léo Meynent, Damian Borth
Abstract
Weight space learning aims to learn representations of neural network (NN) weights, enabling different downstream tasks. Existing approaches show promising performance, but lacking a way to shape these weight-space representations using information about the datasets the models were trained on, thus limiting downstream applications. We propose WeightCLIP, a method for learning a dataset-aligned latent space for neural networks, where datasets information is induced during training. The NNs are encoded as latent representations using an autoencoder, while dataset samples are encoded using a dataset encoder. The two representations are aligned using a contrastive objective, effectively reshaping the weight-space representations according to the datasets. We demonstrate that such representations can be used for different downstream tasks, including mapping dataset information to a weight-space representation that decode to strong models. In addition, we introduce a latent refinement process for generating models that outperforms standard fine-tuning. Overall, our results demonstrate that explicitly incorporating dataset information improves what can be achieved with weight-space representations across retrieval, generation, and refinement. Code will be available at https://github.com/HSG-AIML/WeightCLIP.