WeightCLIP: Aligning Datasets and Models for Weight Space Learning

2026-07-03

AI Summary

Existing weight-space learning approaches struggle to align model weights with the information contained in training datasets, limiting their effectiveness in downstream tasks. This work proposes WeightCLIP, the first method to explicitly incorporate dataset-level information into weight-space learning. It employs an autoencoder to encode model weights and a dedicated dataset encoder to extract sample features, then leverages contrastive learning to align these two representations. Furthermore, a latent-space refinement procedure is introduced to enhance the quality of generated weights. Evaluated across model retrieval, generation, and refinement tasks, WeightCLIP significantly outperforms standard fine-tuning and enables the generation of models with strong generalization capabilities conditioned on dataset information.

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.

Problem

• weight space learning
• dataset alignment
• neural network representations
• downstream tasks
• latent space

Innovation

• weight space learning
• dataset alignment
• contrastive learning
• latent refinement
• neural network representation