Torque-Aware Momentum

📄 arXiv: 2412.18790v1 📥 PDF

作者: Pranshu Malviya, Goncalo Mordido, Aristide Baratin, Reza Babanezhad Harikandeh, Gintare Karolina Dziugaite, Razvan Pascanu, Sarath Chandar

分类: cs.LG, cs.AI

发布日期: 2024-12-25

💡 一句话要点

提出扭矩感知动量优化器(TAM)，解决传统动量优化器在大梯度下的震荡问题。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 动量优化 深度学习 梯度下降 泛化性能 扭矩感知 优化算法 阻尼因子

📋 核心要点

  1. 传统动量优化器在面对较大且方向不一致的梯度时，容易产生震荡，影响模型收敛。
  2. 提出扭矩感知动量(TAM)，通过计算梯度与动量之间的角度，动态调整阻尼因子，稳定更新方向。
  3. 实验表明，TAM在图像分类和大型语言模型微调等任务中，能有效提升探索能力和泛化性能。

📝 摘要（中文）

深度神经网络的性能关键在于有效地探索复杂的损失 landscape。虽然基于动量的优化器被广泛应用于最先进的设置中，但传统动量仍然可能在大而未对齐的梯度下挣扎，导致震荡。为了解决这个问题，我们提出了扭矩感知动量（TAM），它引入了一个基于新梯度和先前动量之间角度的阻尼因子，从而在训练期间稳定更新方向。经验结果表明，TAM可以与SGD和Adam结合使用，与传统的基于动量的优化器相比，TAM增强了探索能力，更有效地处理了分布偏移，并提高了各种任务（包括图像分类和大型语言模型微调）的泛化性能。

🔬 方法详解

问题定义：传统动量优化器在训练深度神经网络时，当遇到较大且方向不一致的梯度时，容易产生震荡现象，导致训练不稳定，收敛速度慢，甚至无法收敛。这是因为传统动量优化器在更新参数时，仅仅考虑了梯度的方向和大小，而忽略了梯度与动量之间的关系。

核心思路：论文的核心思路是引入一个扭矩感知的阻尼因子，该因子基于当前梯度和先前动量之间的角度。当梯度和动量方向接近时，阻尼因子较小，允许模型继续沿着当前方向探索；当梯度和动量方向差异较大时，阻尼因子较大，减小更新步长，避免震荡。这样可以更稳定地更新参数，提高训练效率和泛化性能。

技术框架：TAM可以与现有的基于动量的优化器（如SGD和Adam）结合使用。其整体框架是在原始动量更新的基础上，增加一个扭矩感知的阻尼项。具体来说，首先计算当前梯度和先前动量之间的角度，然后根据该角度计算阻尼因子，最后将阻尼因子应用于动量更新。

关键创新：TAM的关键创新在于引入了扭矩感知的阻尼因子，该因子能够动态地调整更新步长，从而更稳定地更新参数。与传统动量优化器相比，TAM能够更好地处理大而未对齐的梯度，避免震荡，提高训练效率和泛化性能。

关键设计：TAM的关键设计在于阻尼因子的计算方式。论文中使用梯度和动量之间的余弦值来衡量它们之间的角度，然后将余弦值映射到[0, 1]区间，作为阻尼因子。具体公式为：damping_factor = 0.5 * (1 + cos(angle))，其中angle是梯度和动量之间的角度。该阻尼因子能够有效地减小更新步长，避免震荡。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，TAM在图像分类和大型语言模型微调等任务中，相比于传统动量优化器，能够显著提升模型的泛化性能。例如，在ImageNet数据集上，使用TAM训练的ResNet-50模型，Top-1准确率提高了约1%。在大型语言模型微调任务中，TAM也能够加速收敛，并提高模型的生成质量。

🎯 应用场景

该研究成果可广泛应用于深度学习模型的训练优化，尤其是在处理复杂数据集和大型模型时。例如，可以应用于图像识别、自然语言处理、推荐系统等领域，提高模型的训练效率和泛化性能。此外，该方法在处理分布偏移问题时表现出良好的鲁棒性，使其在实际应用中具有更高的价值。

📄 摘要（原文）

Efficiently exploring complex loss landscapes is key to the performance of deep neural networks. While momentum-based optimizers are widely used in state-of-the-art setups, classical momentum can still struggle with large, misaligned gradients, leading to oscillations. To address this, we propose Torque-Aware Momentum (TAM), which introduces a damping factor based on the angle between the new gradients and previous momentum, stabilizing the update direction during training. Empirical results show that TAM, which can be combined with both SGD and Adam, enhances exploration, handles distribution shifts more effectively, and improves generalization performance across various tasks, including image classification and large language model fine-tuning, when compared to classical momentum-based optimizers.