MARS – 字节推出优化大模型训练效率的框架 | AI工具集

MARS是什么

MARS（Make vAriance Reduction Shine）是字节跳动推出的创新的优化框架，提升大型模型训练的效率。MARS融合预条件梯度方法与方差减少技术，基于缩放随机递归动量技术优化梯度估计。MARS框架灵活，支持全矩阵或对角Hessian近似，衍生出基于AdamW、Lion和Shampoo的三种优化算法实例。实验结果表明，MARS在训练GPT-2模型时，相较传统的AdamW优化器，展现出卓越的性能。

MARS的主要功能

• 提高训练效率：MARS结合预条件梯度方法和方差减少技术，提高大型模型训练的效率，尤其是在训练深度神经网络和大型语言模型时。
• 统一优化框架：提供适应不同Hessian近似方法的统一框架，包括全矩阵和对角矩阵近似。
• 算法实例化：在MARS框架下，实现三种具体的优化算法实例：MARS-AdamW、MARS-Lion和MARS-Shampoo，分别基于不同的预条件梯度更新策略。
• 方差减少：引入缩放随机递归动量技术，有效减少训练过程中的梯度方差，加速模型收敛。

MARS的技术原理

• 预条件梯度方法：基于预条件梯度方法调整学习率，让每个参数或参数组有定制的学习率，适应其局部曲率。
• 方差减少技术：引入方差减少技术，如STORM（Stochastic Recursive Momentum），减少随机梯度的方差，加速优化过程。
• 缩放随机递归动量：在STORM的基础上引入缩放参数，调整方差减少的强度，定义新的梯度估计器。
• 梯度裁剪和指数移动平均：为优化训练稳定性，MARS在梯度估计器中应用梯度裁剪，用指数移动平均（EMA）计算递归动量。

MARS的项目地址

MARS的应用场景

• 深度学习模型训练：训练深度神经网络，尤其是参数众多的复杂模型。
• 大规模语言模型：优化大型语言模型的训练过程，如GPT系列模型，提高训练效率和模型性能。
• 计算机视觉任务：在图像分类、目标检测等计算机视觉领域中，加速模型的训练和提高模型的泛化能力。
• 强化学习算法：在强化学习中，优化策略网络或价值函数的参数，尤其是在面对高方差梯度时。
• 推荐系统模型：在构建推荐系统时，优化模型参数，更好地处理大规模用户和物品特征。