MARS框架三大实例：MARS-AdamW、MARS-Lion与MARS-Shampoo对比分析

【免费下载链接】MARSThe official implementation of MARS: Unleashing the Power of Variance Reduction for Training Large Models项目地址: https://gitcode.com/gh_mirrors/mars11/MARS

MARS（Make variance Reduction Shine）是一个统一的优化框架，旨在解决训练大型模型时的固有挑战。它结合了预条件梯度方法和方差减少技术，通过缩放随机递归动量和预条件更新，加速优化过程中临界点的搜索。本文将深入对比MARS框架下的三大实例——MARS-AdamW、MARS-Lion和MARS-Shampoo，帮助开发者选择最适合自己需求的优化器。

MARS框架概述

MARS框架的核心在于将方差减少技术与预条件梯度方法相结合，其数学形式包含四个关键步骤：缩放梯度校正、梯度裁剪、动量更新和预条件优化。通过这种设计，MARS能够同时实现更好的梯度复杂度和每步迭代复杂度，为大型模型训练提供强大支持。

在MARS框架下，有三个主要实例，它们基于不同的Hessian矩阵近似方法：

• MARS-AdamW：使用对角矩阵近似Hessian
• MARS-Lion：基于动量的Hessian近似
• MARS-Shampoo：采用矩阵分解技术近似Hessian

需要注意的是，框架中的超参数默认针对MARS-AdamW进行了调优。使用其他实例时，特别是学习率等关键参数，需要重新调整以获得最佳性能。

MARS-AdamW：对角矩阵近似的高效实现

MARS-AdamW是MARS框架中最成熟的实例，通过对角矩阵近似Hessian矩阵。它的实现位于MARS/mars.py，通过设置mars_type="mars-adamw"启用。

核心公式

MARS-AdamW的Hessian矩阵近似定义为：

v_t = β₂v_{t-1} + (1-β₂)(∇f(x_t, ξ_t))² H_t := √(diag(v_t)) · (1 - β₁ᵗ)/√(1 - β₂ᵗ)

性能表现

MARS-AdamW在多个任务上表现出色。在FineWeb-Edu数据集上，GPT-2 Small模型使用MARS-AdamW达到了45.93的平均分数，显著优于AdamW和OpenAI基线。

图1：MARS-AdamW在GPT-2 Small模型上的训练损失曲线，显示出快速收敛特性

在GPT-2 XL模型上，MARS-AdamW更是实现了56.52的HellaSwag准确率，证明了其在大型模型上的优势。

图2：MARS-AdamW在GPT-2 XL模型上的验证损失曲线，展示了持续的低损失表现

适用场景

• 需要快速收敛的大型语言模型训练
• 资源有限但追求高精度的场景
• 作为其他MARS实例的性能基准

MARS-Lion：基于动量的轻量级方案

MARS-Lion是MARS框架的轻量级实例，通过动量的平方对角矩阵近似Hessian。它的实现同样位于MARS/mars.py，通过设置mars_type="mars-lion"启用。

核心公式

MARS-Lion的Hessian矩阵近似定义为：

H_t := √(diag(m_t²))

其中m_t是动量项，这种设计使得MARS-Lion的计算复杂度低于MARS-AdamW。

特点与优势

• 计算效率高：无需维护二阶矩估计，内存占用更小
• 收敛稳定：基于动量的更新有助于避免局部最优
• 调参简单：相比MARS-AdamW，需要调整的超参数更少

适用场景

• 内存受限的大型模型训练
• 需要简化调参流程的场景
• 对训练速度要求较高的应用

MARS-Shampoo：矩阵分解的高阶近似

MARS-Shampoo是MARS框架中最复杂的实例，采用矩阵分解技术近似Hessian矩阵。它的实现位于MARS/mars.py，通过设置mars_type="mars-shampoo"启用。

核心公式

MARS-Shampoo的预条件器基于SVD分解：

U_t, Σ_t, V_t = SVD(G_t) x_{t+1} = x_t - η_t U_t V_t^⊤

在实践中，MARS-Shampoo使用Newton-Schulz迭代加速SVD问题的求解，平衡了计算复杂度和近似精度。

特点与优势

• 高阶近似：能够捕捉参数间的相关性
• 理论最优：在某些条件下可达到二阶方法的收敛速度
• 泛化能力强：在小样本数据上可能表现更好

适用场景

• 数据量有限但模型复杂的任务
• 追求理论最优解的研究场景
• 特征维度高且存在相关性的应用

三大实例的性能对比

虽然MARS-AdamW是目前文档中唯一提供详细实验数据的实例，但我们可以根据三种方法的特性进行理论对比：

计算复杂度

• MARS-AdamW：O(d)，d为参数维度
• MARS-Lion：O(d)，但常数因子更小
• MARS-Shampoo：O(d^(3/2))，但可通过低秩近似优化

内存占用

• MARS-AdamW：高（需要存储二阶矩）
• MARS-Lion：低（仅需存储动量）
• MARS-Shampoo：中到高（取决于矩阵分解策略）

收敛特性

• MARS-AdamW：均衡的收敛速度和稳定性
• MARS-Lion：初期收敛快，后期可能震荡
• MARS-Shampoo：理论收敛快，但实际可能受近似质量影响

实际应用建议

对于大多数用户，我们建议从MARS-AdamW开始，因为它在各种任务上都表现出稳定的高性能。如果你面临内存限制，可以尝试MARS-Lion；如果追求理论上的最优解，且能接受更高的计算成本，可以尝试MARS-Shampoo。

图3：MARS实例在不同规模GPT模型上的训练效率对比，展示了MARS框架的整体优势

快速开始使用MARS实例

要在你的项目中使用MARS框架的不同实例，只需在初始化优化器时指定mars_type参数：

# 导入MARS优化器 from mars import MARS # 使用MARS-AdamW optimizer = MARS(model.parameters(), lr=1e-3, mars_type="mars-adamw") # 使用MARS-Lion optimizer = MARS(model.parameters(), lr=1e-3, mars_type="mars-lion") # 使用MARS-Shampoo optimizer = MARS(model.parameters(), lr=1e-3, mars_type="mars-shampoo")

项目提供了多种配置文件和脚本，方便快速启动不同模型的训练：

• 配置文件：config/目录下包含各种模型的训练配置
• 脚本文件：scripts/目录下提供了一键启动脚本

例如，要使用MARS-AdamW训练GPT-2 Small模型，可以运行：

$ bash scripts/run_mars_small.sh

总结

MARS框架通过统一的设计理念，提供了三种各具特色的优化器实例，满足不同场景下的训练需求。MARS-AdamW作为基准实例，在性能和稳定性之间取得了平衡；MARS-Lion以其轻量级设计提供了高效的训练方案；MARS-Shampoo则通过高阶近似追求理论上的最优解。

无论你是训练大型语言模型还是计算机视觉模型，MARS框架都能提供强大的优化支持。建议根据你的具体任务需求、资源限制和精度要求，选择最适合的MARS实例，并通过调整超参数进一步优化性能。

图4：MARS在CIFAR-10数据集上的测试准确率，展示了其在计算机视觉任务上的优势

通过合理选择和配置MARS框架的优化器实例，你可以充分释放大型模型的训练潜力，加速模型收敛并提高最终性能。

【免费下载链接】MARSThe official implementation of MARS: Unleashing the Power of Variance Reduction for Training Large Models项目地址: https://gitcode.com/gh_mirrors/mars11/MARS