将扩散模型用于神经网络权重空间,实现更优初始化与安全检测。
Flows and Diffusions on the Neural Manifold
- 把梯度下降路径建模为轨迹推断问题,引入优化动态作为先验。
- 生成的权重在下游任务上表现优于基线,初始化效果显著提升。
- 适合需要可靠初始化或异常检测的高风险应用场景。
扩散和流模型在图像合成、视频生成与自然语言建模等领域取得显著成功。本文将这些进展拓展至权重空间学习,利用近期技术引入源于优化动力学的结构先验。核心思想是将梯度下降产生的轨迹视为轨迹推断问题。我们统一多种轨迹推断方法以匹配梯度流,为将优化路径作为归纳偏置提供理论框架。进一步探索了架构与算法选择,包括通过伴随匹配进行奖励微调、使用自编码器进行隐式权重表示、基于任务特定上下文数据的条件生成,以及采用如Kaiming均匀等信息丰富的源分布。实验表明,该方法在生成分布内权重方面达到或超越基线,改善了下游训练的初始化,并支持微调以提升性能。最后,我们在安全关键系统中演示了实际应用:检测有害的协变量偏移,结果优于最接近的可比基线。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。