训练时保存梯度信息,直接用于模型合并,省时高效。
Bridging Training and Merging Through Momentum-Aware Optimization
- 训练中积累动量与曲率统计,为合并提供重要性评估
- 在语言理解任务上,合并后性能比基线高1.6%
- 无需额外计算,可直接生成可合并模型,适合多任务场景
训练大型神经网络与合并特定任务模型均依赖低秩结构和参数重要性估计,但两者长期孤立研究。现有流程在训练中计算曲率信息后丢弃,合并时又重新计算——浪费算力并丢失轨迹数据。本文提出统一框架,在训练中持续维护因子化动量与曲率统计,并复用该信息实现几何感知的模型融合。方法仅增加约30%内存开销(相比AdamW),通过优化过程自然产出任务显著性分数,其重要性估计效果接近事后Fisher计算,且可直接生成适配合并的模型。理论证明非凸目标下收敛性,近似误差受梯度奇异值衰减约束。在自然语言理解基准测试中,曲率感知参数选择在所有稀疏度下均优于仅基于幅度的基线,多任务合并性能较强基线提升1.6%。该框架具有秩不变收敛性与更强超参数鲁棒性,相较于现有低秩优化器表现更优。通过将优化轨迹视为可复用资产,表明训练期曲率信息足以支撑有效模型组合,实现训练-合并一体化流程。
原文摘要 · Abstract (English)
Training large neural networks and merging task-specific models both exploit low-rank structure and require parameter importance estimation, yet these challenges have been pursued in isolation. Current workflows compute curvature information during training, discard it, then recompute similar information for merging--wasting computation and discarding valuable trajectory data. We introduce a unified framework that maintains factorized momentum and curvature statistics during training, then reuses this information for geometry-aware model composition. The proposed method incurs modest memory overhead (approximately 30% over AdamW) to accumulate task saliency scores that enable curvature-aware merging. These scores, computed as a byproduct of optimization, provide importance estimates comparable to post-hoc Fisher computation while producing merge-ready models directly from training. We establish convergence guarantees for non-convex objectives with approximation error bounded by gradient singular value decay. On natural language understanding benchmarks, curvature-aware parameter selection outperforms magnitude-only baselines across all sparsity levels, with multi-task merging improving 1.6% over strong baselines. The proposed framework exhibits rank-invariant convergence and superior hyperparameter robustness compared to existing low-rank optimizers. By treating the optimization trajectory as a reusable asset rather than discarding it, our approach demonstrates that training-time curvature information suffices for effective model composition, enabling a unified training-merging pipeline.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。