Flows and Diffusions on the Neural Manifold
This work addresses the problem of weight generation and optimization for machine learning practitioners, offering incremental improvements by applying existing generative techniques to a new domain with specific algorithmic enhancements.
The paper tackles the problem of generating neural network weights by extending diffusion and flow-based generative models to weight space learning, modeling gradient descent trajectories as inductive bias. The method matches or surpasses baselines in generating in-distribution weights, improves initialization for downstream training, and outperforms baselines in detecting harmful covariate shifts in safety-critical systems.
Diffusion and flow-based generative models have achieved remarkable success in domains such as image synthesis, video generation, and natural language modeling. In this work, we extend these advances to weight space learning by leveraging recent techniques to incorporate structural priors derived from optimization dynamics. Central to our approach is modeling the trajectory induced by gradient descent as a trajectory inference problem. We unify several trajectory inference techniques towards matching a gradient flow, providing a theoretical framework for treating optimization paths as inductive bias. We further explore architectural and algorithmic choices, including reward fine-tuning by adjoint matching, the use of autoencoders for latent weight representation, conditioning on task-specific context data, and adopting informative source distributions such as Kaiming uniform. Experiments demonstrate that our method matches or surpasses baselines in generating in-distribution weights, improves initialization for downstream training, and supports fine-tuning to enhance performance. Finally, we illustrate a practical application in safety-critical systems: detecting harmful covariate shifts, where our method outperforms the closest comparable baseline.