Abstract
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 under the framework of gradient flow matching, 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.
Abstract (translated)
扩散模型和基于流的生成模型在图像合成、视频生成和自然语言建模等领域取得了显著的成功。在这项工作中，我们通过利用最近的技术将这些进展扩展到权重空间学习中，这些技术从优化动力学中推导出结构先验并将其纳入考虑。我们的方法的核心是将梯度下降诱导的轨迹作为轨迹推理问题进行建模。我们将几种轨迹推理技术统一在梯度流匹配框架下，提供了一种理论框架，用于处理优化路径作为归纳偏置。 此外，我们进一步探索了架构和算法选择，包括通过伴随匹配进行奖励微调、使用自编码器表示潜在权重、根据特定任务的上下文数据进行条件设置以及采用如Kaiming均匀分布等信息源分布。实验结果表明，我们的方法在生成符合分布的权重方面与基线相匹敌或超越，并且可以改善下游训练的初始化，支持微调以增强性能。 最后，我们展示了这种方法在一个实际应用中的优势：安全关键系统的有害协变量偏移检测，在此领域中，我们的方法优于最接近的同类基准。
URL
https://arxiv.org/abs/2507.10623
