arXiv:2502.10119cs.LG2025-02被引 3

自动挑选训练末期关键参数点,提升模型泛化与收敛速度

SeWA: Selective Weight Average via Probabilistic Masking

  • 用概率掩码自动筛选训练末期的参数快照进行平均
  • 仅需少数几个快照就实现更好泛化和更快收敛
  • 理论证明比SGD更优,适用于多种任务场景

权重平均已成为提升模型性能的标准方法。然而,如随机权重平均(SWA)和最新权重平均(LAWA)等方法通常需要人工设计从训练轨迹中采样策略,且结果高度依赖超参数调优。为减少人工干预,本文提出一种简单高效的算法——选择性权重平均(SeWA),该方法在训练末期自适应地选择用于平均的检查点。基于SeWA,我们发现仅需少量关键点即可实现更优的泛化能力和更快的收敛速度。由于离散子集选择问题本身具有挑战性,我们将其转化为连续的概率优化框架,并采用Gumbel-Softmax估计器学习每个检查点的不可导掩码。进一步地,我们理论上推导了基于稳定性的泛化界,其在凸与非凸假设下均优于SGD。最后,在行为克隆、图像分类和文本分类等多个领域开展的充分实验,进一步验证了本方法的有效性。

原文摘要 · Abstract (English)

Weight averaging has become a standard technique for enhancing model performance. However, methods such as Stochastic Weight Averaging (SWA) and Latest Weight Averaging (LAWA) often require manually designed procedures to sample from the training trajectory, and the results depend heavily on hyperparameter tuning. To minimize human effort, this paper proposes a simple yet efficient algorithm called Selective Weight Averaging (SeWA), which adaptively selects checkpoints during the final stages of training for averaging. Based on SeWA, we show that only a few points are needed to achieve better generalization and faster convergence. Theoretically, solving the discrete subset selection problem is inherently challenging. To address this, we transform it into a continuous probabilistic optimization framework and employ the Gumbel-Softmax estimator to learn the non-differentiable mask for each checkpoint. Further, we theoretically derive the SeWA's stability-based generalization bounds, which are sharper than that of SGD under both convex and non-convex assumptions. Finally, solid extended experiments in various domains, including behavior cloning, image classification, and text classification, further validate the effectiveness of our approach.

权重平均自适应优化泛化能力训练效率

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。