Abstract
大型语言模型（LLMs）在多种自然语言处理任务中展现出了卓越的性能。然而，其庞大的参数规模在训练过程中带来了显著的内存挑战，尤其是在使用像Adam这样内存消耗较高的优化器时 。现有的内存高效算法通常依赖于奇异值分解投影或权重冻结等技术 。尽管这些方法有助于缓解内存限制，但与全秩更新相比，它们通常会产生次优的结果。在本文中，我们探索了超越低秩训练的内存高效方法，提出了一种名为梯度小波变换（GWT）的创新解决方案。该方法通过将小波变换应用于梯度，显著降低了维护优化器状态所需的内存需求。我们证明了GWT能够与高内存消耗的优化器无缝集成，从而在不牺牲性能的前提下实现高效训练。通过在预训练和微调任务上的广泛实验，我们展示了GWT在内存使用和训练性能方面均达到了与先进内存高效优化器及全秩方法相当的效果。
1 Introduction
自ChatGPT问世以来[2]，大型语言模型（LLMs）取得了突破性进展。其卓越性能与可扩展性推动了该技术在多个跨学科领域的快速应用，涵盖科学领域[1,27,28]、医学领域[13,29,33,49]以及生物学领域的政策制定[20]。LLMs的卓越能力源于其庞大的参数量和海量训练数据集。然而模型训练对优化器提出了严苛要求，其中Adam优化器[22]凭借快速收敛特性、稳健性及丰富的实验数据，已成为业界最常用的优化器。
尽管具有这些优势，Adam优化器的内存需求仍存在显著缺陷。具体而言，其内存占用量是模型本身的两倍（如图1所示），这使得内存使用成为大语言模型训练中的关键瓶颈。例如，预训练LLaMA 7B[43]模型需要约58GB显存：其中14GB用于模型权重、14GB用于梯度计算、2GB用于激活函数处理，以及在BF16浮点精度下运行Adam优化器状态所需的28GB内存[54]。当模型参数量超过1000亿时（如GPT-3[2]的1750亿参数），Adam优化器带来的额外内存开销可超过700GB，此时至少需要额外配置9块 NVIDIA A100 80GB显卡。
鉴于训练大语言模型的高昂成本，那些通过减小批量大小或使用更大GPU的方法很快就会变得不可持续。为了缓解这些问题，许多研究聚焦于优化大语言模型训练中的内存使用。一个颇具前景的解决方案是低秩训练 [21, 54]，该方法在缓解内存限制方面取得了显著进展。这些内存高效的方法大致可以分为两类：基于权重的方法和基于梯度的方法 [6]。其中代表性的工作包括低秩适应 (LoRA) [21] 和梯度低秩投影 (GaLore) [54]。
LoRA通过冻结预训练模型权重，并采用可训练的低秩分解矩阵对权重更新进行重新参数化，充分利用了模型参数更新的低秩特性。大量实验表明，LoRA在微调任务中效果显著，但在预训练任务中的表现仍存在局限性。此外，LoRA所依赖的强低秩假设往往导致次优结果[50,52]。相比之下，FLoRA[16]通过采用重采样投影矩阵来解决LoRA的低秩约束问题。
与通过分解更新矩阵实现优化的LoRA不同，GaLore通过在梯度上执行奇异值分解（SVD）来降低优化器状态内存占用量。这是因为优化器状态所占用的内存直接与梯度维度相关（如图1所示）。GaLore在预训练和微调任务中均展现出优异性能，但其表现仍不及满秩方法，且仅能捕捉投影子空间内的信息，会丢弃子空间外的梯度信息。当使用低秩矩阵或梯度存在噪声时，这一局限性会变得尤为突出[6,17]。
为了应对这些挑战，Fira [6] 引入了全秩梯度与其子空间投影之间的误差残差，而 GoLore [17] 则通过利用随机投影来增强 GaLore，从而在噪声占主导时更好地捕捉梯度信息。然而，这些方法仍然依赖奇异值分解（SVD）来更新子空间，而 SVD 的高计算成本（对于大小为 \( m \times n \) 的矩阵，复杂度为 \( O(m \times n^2) \)）可能会显著拖慢训练速度，尤其是在大模型上。为了缓解这一问题，APOLLO [56] 提出了一种无 SVD 的 GaLore 变体，用随机投影替代了 SVD。这一改进不仅提升了训练速度，还带来了更优的经验性能。
虽然人们直观认为低秩近似中的信息丢失会导致性能下降，但大量实验结果却表明事实恰恰相反：Fira和Apollo算法在多种场景下均展现出超越全秩训练的惊人性能。这些发现凸显了内存高效优化算法作为全秩方法替代方案的潜力，使得高效训练不仅成为必要条件，更在实际应用中具有显著优势。
上述方法均基于低秩分解技术。除低秩分解外，小波变换（WT）在压缩领域也具有广泛的应用价值。相较于低秩分解，小波变换具有更低的计算复杂度（O(m×n)）、更强的局部信息特征捕捉能力以及更优的高频噪声抑制效果[37]。此外，由于其对噪声具有鲁棒性，特别适用于包含噪声的随机梯度下降训练模型方法。因此，我们自然会提出以下问题：
小波分解（如奇异值分解SVD）能否像其他方法一样，对大型语言模型的训练产生积极影响，从而降低内存消耗并提升模型性能？
为了解决现有低秩训练方法的局限性，我们提出了梯度小波变换 （Gradient Wavelet Transform, GWT），一个简单而有效的即插即用框架，它在训练过程中对小波变换应用于梯度。通过利用小波天然的稀疏性和局部化特性，GWT 压缩了梯度表示，并显著减少了存储优化器状态所需的内存，且不依赖于低秩假设。与传统的基于矩阵的压缩技术不同，GWT 是首个从**频域视角**进行梯度压缩的方法。它通过选择性地控制低频和高频系数来实现内存效率。当应用于大规模语言模型预训练时，例如在 C4 数据集上训练 LLaMA 模型，GWT 在 3B 规模模型上实现了高达 **71% 的优化器内存减少**和 **1.9 倍的训练加速**，同时保持甚至提升了模型质量。此外，GWT 能够与 Adam 之外的优化器无缝集成，展现了其在不同训练场景下的通用性。这些结果表明，GWT 为内存受限的深度学习提供了一种**实用、优化器无关且高效**的解决方案。
2 Related Works
**Memory-Efficient Optimizers.**鉴于 Adam 优化器 [22] 的高内存需求，许多方法致力于改进该优化器或创造新的优化方法，以降低大语言模型的训练成本。低秩适应（Low-Rank Adaptation, LoRA） [21] 通过冻结预训练模型的参数，并引入可训练的低秩矩阵来分解预训练权重，从而解决了这一问题。LoRA 已被证明是优化大语言模型训练的有效策略，后续多项研究在此基础上进一步提升了优化性能 [16, 23, 50, 52]。梯度低秩投影（Gradient Low-Rank Projection, GaLore）[54] 通过将梯度投影到低维子空间中，降低了优化器状态的内存成本。该方法不仅支持大语言模型的全参数训练，还在预训练和微调任务中均展现出强劲的性能。近期基于 GaLore 的扩展方法已在实验结果表明其性能优于全秩方法 [6, 56]。
**System Memory Efficiency.**除了在优化器层面优化内存使用外，一些策略还致力于降低整个系统的内存消耗。**梯度检查点（Gradient Checkpointing）** [5] 就是其中一种方法，它通过不存储激活值来节省内存。**内存卸载（Memory Offloading）** [35] 通过利用非 GPU 内存资源来提升训练过程中的内存可用性。**量化技术（Quantization Techniques）** [8, 9] 则通过采用更低比特的数据表示来进一步降低内存开销。这些方法极大地加速并优化了大语言模型的训练过程。我们的 **GWT 方法**相较于上述技术提供了一种独特的优势，它是一种内存高效的替代方案。此外，我们的方法设计为与现有的系统级优化策略兼容，能够通过协同集成，实现更快、更高效的大语言模型训练。
3 Methods
本节首先概述小波变换的基本理论，随后详细阐述所提出的梯度小波变换（GWT）方法如何集成到Adam优化器[22]中。
3.1 Discrete Haar Wavelet Transform
小波变换可分为连续型和离散型两类。在本文中，我们主要关注**离散 Haar 小波**，原因在于其基函数简单易懂、计算效率高、在压缩任务中适用性广 [32]，并且分解后的数据总量与原始数据相等，不会引入额外信息。它将信号（或图像）分解为两个主要分量：(a) 近似系数（低频），代表信号/图像中平滑或平均的部分；(b) 细节系数（高频），捕捉信号/图像中更精细的细节或差异。
因此，我们得到了二级 DHT 的简单实现，其对应的重构过程与一级类似。DHT 将一个包含 8 个元素的向量转换为两个各含 4 个元素的向量 (A1,D1)，或三个向量 (A2,D2,D1)。仅存储近似系数将使我们的内存成本减少 50% 或 75%。这个过程可以迭代重复，将信号逐级分解为更细粒度的近似和细节层次，从而便于高效的压缩和存储。
3.2 Adam with GWT
总之，我们将提出的 GWT 方法整合到Adam优化器中，发现该方法在理论上可将优化器状态的内存占用量降低高达75%。 GWT 通过将梯度投影到低维子空间，并在该子空间内更新优化器状态来实现这一目标。这种内存高效策略显著降低了存储需求，可广泛应用于需要维护大量状态信息的各类优化器。值得注意的是， GWT 的应用范围不仅限于Adam优化器。在附录F中，我们提出了将 GWT 扩展至其他内存密集型优化器的统一框架，并通过实证结果表明：当与Adam-mini[53]、MUON[24]等优化器结合使用时， GWT 能达到与这些优化器相当的性能表现。
4 Experiments
在本节中，我们证明**梯度小波变换（GWT）** 作为一种内存高效的优化技术，能够达到与全秩优化器相当甚至更优的性能，同时显著降低内存使用并提高训练吞吐量。具体而言，我们在两个场景下评估 GWT：在 Colossal Clean Crawled Corpus (C4) [34] 英文基准上预训练 LLaMA [43] 模型，以及在多任务语言理解（MMLU）[18] 基准上微调各种预训练模型，以及在通用语言理解评估（GLUE）[45] 基准上微调 RoBERTa-base 模型。我们在所有实验中采用 BF-16 格式来降低内存使用，并在附录 G 中提供了这些数据集的详细描述。
对于 GWT，我们默认使用离散 Haar 小波作为滤波器。此外，我们在附录 A.2 中呈现了一项消融研究，以考察关键超参数的影响，包括缩放因子（α）、初始学习率（lr）、GWT 级数（l）以及其他小波基。实验设置的全面描述见附录 A。
4.1 Memory-Efficient Pre-Training
我们将提出的 GWT 方法集成到 Adam 优化器 [22] 中，使用默认超参数（\(\beta_1 = 0.9, \beta_2 = 0.999, \epsilon = 10^{-6}\)），并在 C4 预训练任务上评估其在 LLaMA 模型上的性能。为了提供全面的比较，我们纳入了全秩 Adam [22]、GaLore [54] 和 APOLLO [56] 的复现结果，以及 Fira [6] 和 LoRA [54] 报告的参考结果。所有实验采用一致的配置：对于 GaLore 和 GWT，我们使用 GaLore 的默认超参数（\(lr = 0.01, \alpha = 0.25\)）；对于 APOLLO，使用 \(\alpha = 1.0\)。学习率采用余弦退火调度。除非另有说明，我们使用序列长度 256 和批量大小 512，每批次总计 131K 个 token。
对于 GaLore-r 和 APOLLO-r，我们实验了与模型秩的 1/4 和 1/8 对应的压缩比，分别记为 GaLore-1/4 和 GaLore-1/8。这些设置与使用 2 级和 3 级分解的 GWT 配置相对应。遵循先前研究 [54]，我们采用相同的网络初始化策略，并在 Transformer 的多层感知机（MLP）和注意力模块中都启用了 GWT。
**GWT 在预训练中是有效的。**所有方法的最终验证困惑度（PPL）和估计内存使用量汇总于表 1。实验结果表明，与其他内存高效的基线方法相比，GWT 在降低内存使用的同时， consistently 实现了更低的验证困惑度。在所有预训练实验中，GWT 均优于全秩优化器。具体而言，在 LLaMA 1B 预训练实验中，我们的 2 级 GWT 方法将总内存成本降低了 51%，优化器内存成本降低了 76%（在 LLaMA 1B 上从 7.80GB 降至 3.81GB）。即使使用 2 级 GWT，我们的方法仍然取得了超越全秩的实验结果，同时将总内存消耗降低了高达 43%。在同一场景下，GaLore-1/8 与全秩之间的差距进一步扩大。内存估计的详细信息见附录 A.4。
**GWT 能够实现更快的收敛。**我们在图 4 中展示了 LLaMA 60M-350M 预训练任务的学习曲线。这些曲线表明，在不同级别的小波变换下，GWT 始终优于全秩 Adam 和 GaLore。在大多数情况下，GaLore 与全秩 Adam 相比始终显示出显著的差距，而所有 GWT 方法都保持了超越全秩的结果。此外，与基线方法相比，GWT 不仅实现了更低的最终验证困惑度，还加速了收敛。
**GWT 与优化器无关。**我们将方法扩展到了 Adam-mini [53] 和 MUON [24] 优化器，并在图 4 中展示了相应的实验结果。与在 Adam 上的表现类似，GWT 继续取得了与全秩方法相当或更优的结果。这突显了我们的方法在与其他优化器集成时的通用性和有效性，展示了其在内存高效优化中的广泛适用性。
GWT 实现了更高的训练吞吐量。 除了评估最终验证困惑度和内存消耗外，我们还测量了在使用 GWT-2、GaLore-1/4、APOLLO-1/4 和 8-bit Adam 预训练 LLaMA-3B 模型时的 token 吞吐量。结果如表 2 所示。我们的方法 GWT-2 实现了每 GPU 每秒 0.532K token 的训练吞吐量，与 8-bit Adam [8] 相比实现了超过 **1.9 倍的加速**，并且比 GaLore 提升了约 6%。此外，GWT-2 实现了与无 SVD 的 APOLLO 方法相当的吞吐量，凸显了其在大规模模型训练中的高效性。
**GWT 对长序列长度具有鲁棒性。**我们在保持每批次总训练 token 数不变的情况下，评估了 GWT 在处理更长序列（512、1024）时的表现。表 3 报告了不同方法下的最终验证困惑度。与 GaLore（随着序列长度增加表现出明显的性能下降）相比，GWT 保持了稳定的性能，仅显示出轻微的 PPL 增加。这种鲁棒性凸显了 GWT 在部分大语言模型预训练场景（即长上下文窗口和数万亿训练 token）中是一个有前景的候选方法。
**在SGD级别的内存限制下， GWT 仍保持高效性能。**如图5所示，我们评估了 GWT 在更高分解层级下的表现。值得注意的是，即使在极端层级（如l=7、8、9）下------此时优化器的内存占用接近SGD水平------ GWT 在验证PPL指标上始终优于满秩Adam，这充分证明了其在低内存训练环境中的稳健性和有效性。
4.2 Memory-Efficient Fine-Tuning
在本节中，我们通过微调 LLaMA-3.2-3B [15]、Gemma3-4B [41] 和 Qwen2.5-7B [51] 模型，进一步评估了我们方法的有效性，并报告了它们在 MMLU 基准 [18] 上的性能。为了公平比较，LoRA、GaLore 和 APOLLO 均配置为秩 8，而 GWT 设置为分解级数 8，以确保各方法之间的内存对齐。我们通过在 \([1.0e-6, 2.5e-6, 5.0e-6, 1.0e-5, 2.5e-5, 5.0e-5, 1.0e-4]\) 范围内扫描学习率来报告最佳准确率，该范围覆盖了足够宽的区间，确保对学习率敏感性不同的优化器都能得到公平评估。此外，为了提供更全面的评估，我们在附录 B 中展示了在 GLUE 基准上微调 RoBERTa-base 模型 [25] 的结果。
**GWT 在微调中同样有效。**我们在表 4 中报告了 MMLU 结果，在表 10 中报告了 GLUE 结果。如图所示，GWT 持续达到或超越其他内存高效基线的性能。这些结果表明，GWT 不仅在预训练中有效，在微调场景中同样表现出色，能够在保持内存效率的同时，在多样化的任务上取得具有竞争力的性能。
5 Conclusion
在本文中，我们探索了超越传统低秩分解和系统级技术的内存高效优化算法。受当前优化方法的局限性以及小波变换在图像和信号处理中已证实的成功的启发，我们提出了一种新颖的优化器状态内存压缩方法，将小波变换集成到梯度中。该方法弥合了梯度压缩与小波变换之间的鸿沟，通过减少内存使用和加速训练速度，实现了更高效的大语言模型训练。
我们的实验结果表明，GWT 在不同的大语言模型训练阶段中，其性能均能与先进的内存高效优化器及全秩方法相媲美。这种梯度压缩技术不仅提高了最终验证分数，还加速了训练收敛并提高了吞吐量。此外，该方法适用于 Adam 之外现有的内存密集型优化器。总之，我们的方法验证了应用小波变换来提升大语言模型训练效率的可行性，为设计能够有效解决大语言模型训练中内存瓶颈的内存高效算法提供了新的视角。