Flow to Learn: Flow Matching on Neural Network Parameters

Foundational language models show a remarkable ability to learn new concepts during inference via context data. However, similar work for images lag behind. To address this challenge, we introduce FLoWN, a flow matching model that learns to generate neural network parameters for different tasks. Experiments verify that FLoWN attains various desiderata for a meta-learning model. In addition, it matches or exceeds baselines on in-distribution tasks, provides better initializations for classifier training, and is performant on out-of-distribution few-shot tasks while having a fine-tuning mechanism to improve performance. However, its application to neural network weights has not been explored. By leveraging the principled, yet versatile training of FM, we aim to generate task-specific weights on novel tasks. Multiple approaches have been tried to generate weights capable of few-shot learning (FSL), motivated by its speed compared to conventional training. For instance, various diffusion-based approaches (Soro et al., 2024; Zhang et al., 2024; Wang et al., 2024) have been used to generate neural network weights. However, flexibility is limited by its restriction to Gaussian processes and a sluggish inference speed.