Flow to Learn: Flow Matching on Neural Network Parameters

TL;DR

This work addresses the challenge of generating effective neural network weights for image tasks by applying flow matching to latent weight representations. It introduces FLoWN, a flow-to-learn framework that uses a conditional vector field $v_ heta({f x},t;{f y})$ to steer latent weights from a prior $p_0$ toward a target $p_1$ conditioned on context ${f y}$, trained via $L_{ ext{cfm}}$. The method combines a weight encoder (VAE, graph encoder, or LoRA-based) with a conditional flow meta-model and a CAML-based conditioning module to produce task-specific weights, enabling strong in-distribution performance, improved initializations for OOD fine-tuning, and competitive few-shot behavior. Empirically, FLoWN matches conventional training on standard in-distribution tasks, provides faster convergence when reusing retrieved weights, and demonstrates meaningful gains when fine-tuned on OOD data, albeit with some limitations in certain OOD datasets and tasks. Overall, the approach offers a principled, flexible pathway for rapid weight generation in image-domain meta-learning, with potential extensions to broader datasets and applications such as medical imaging.

Abstract

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. Our approach models the flow on latent space, while conditioning the process on context data. 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.

Figures (2)

• Figure 1: A schematic of the training process of FLoWN for few-shot learning. Given a set of pre-trained target weights and a support set, we apply the conditioned flow model to pushforward a sample of the latent prior towards encoded target weights. The decoder is used during inference where we start from a sample of the latent prior and pushforward towards the target distribution with a trained vector field $v_\theta(\cdot, t; {\bm{y}})$ where ${\bm{y}}$ is the support set embedding.
• Figure 2: The training loss curve for the mini-Imagenet run from Table \ref{['tab:model-retrieval']}.