WeightCLIP: Aligning Datasets and Models for Weight Space Learning

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.

Training a neural network from scratch for every new dataset can be expensive and time-consuming. Recent research has explored whether we can learn from collections of already trained neural networks by studying their weights. However, existing methods often do not clearly connect a model’s weights to the dataset it was trained on, making it difficult to choose the right model for a new dataset or to generate a new model for a new dataset. We introduce WeightCLIP, a method that links samples from a dataset with the weights of neural networks trained on that dataset. The dataset samples can then act like a prompt: given a few examples from a new dataset WeightCLIP can find relevant models, generate a new model for that dataset, and improve it with a small additional training step. Our results show that this link makes it easier to reuse what previously trained models have learned. In the long run, this could make it easier to generate new models by reusing what has already been learned, rather than starting from scratch each time.