Wavelet Meets Adam: Compressing Gradients for Memory-Efficient Training

📄 arXiv: 2501.07237v3 📥 PDF

作者: Ziqing Wen, Ping Luo, Jiahuan Wang, Xiaoge Deng, Jinping Zou, Kun Yuan, Tao Sun, Dongsheng Li

分类: cs.LG, cs.AI

发布日期: 2025-01-13 (更新: 2025-07-30)

💡 一句话要点

提出梯度小波变换（GWT）以压缩Adam优化器状态，实现大模型高效训练。

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

关键词: 梯度压缩 小波变换 内存优化 大型语言模型 Adam优化器

📋 核心要点

  1. 大型语言模型训练中，Adam等优化器状态占用大量内存，现有低秩分解等方法效果欠佳。
  2. 提出梯度小波变换（GWT），通过小波变换压缩梯度信息，降低优化器状态的内存占用。
  3. 实验表明，GWT在预训练和微调任务上，内存效率和训练性能均优于现有内存优化方法。

📝 摘要（中文）

大型语言模型（LLMs）在各种自然语言处理任务中表现出令人印象深刻的性能。然而，它们庞大的参数量在训练期间带来了巨大的内存挑战，尤其是在使用像Adam这样内存密集型优化器时。现有的内存高效算法通常依赖于奇异值分解投影或权重冻结等技术。虽然这些方法有助于缓解内存限制，但与全秩更新相比，它们通常会产生次优结果。在本文中，我们研究了低秩训练之外的内存高效方法，提出了一种名为梯度小波变换（GWT）的新颖解决方案，该方法将小波变换应用于梯度，从而显著降低了维护优化器状态的内存需求。我们证明了GWT可以与内存密集型优化器无缝集成，从而实现高效训练而不会牺牲性能。通过在预训练和微调任务上的大量实验，我们表明，在内存使用和训练性能方面，GWT与先进的内存高效优化器和全秩方法相比，实现了最先进的性能。

🔬 方法详解

问题定义：论文旨在解决大型语言模型训练过程中，使用Adam等内存密集型优化器时，优化器状态占用大量内存的问题。现有方法如低秩分解或权重冻结虽然能降低内存占用，但会牺牲模型性能，无法达到全秩更新的效果。

核心思路：论文的核心思路是利用小波变换对梯度进行压缩，从而减少优化器状态的内存占用。小波变换具有良好的能量集中特性，可以将梯度中的重要信息集中在少数几个系数上，从而实现高效压缩。

技术框架：GWT可以无缝集成到现有的训练流程中。具体来说，在每次梯度更新时，首先对梯度进行小波变换，然后只保留重要的（例如，绝对值最大的）小波系数，其余系数置零。压缩后的梯度用于更新优化器状态。在需要时，可以通过逆小波变换恢复近似的原始梯度。

关键创新：GWT的关键创新在于将小波变换应用于梯度压缩，从而在不显著损失模型性能的前提下，大幅降低了优化器状态的内存占用。与传统的低秩分解方法相比，GWT能够更好地保留梯度中的重要信息，从而实现更好的训练效果。

关键设计：GWT的关键设计包括小波基的选择、小波分解的层数、以及保留的小波系数的比例。论文中可能探讨了不同小波基（如Haar小波、Daubechies小波等）对压缩效果的影响。此外，保留的小波系数的比例需要根据具体的任务和模型进行调整，以在内存占用和模型性能之间取得平衡。论文可能还涉及了如何有效地实现小波变换和逆变换，以减少计算开销。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

论文通过在预训练和微调任务上的大量实验，证明了GWT的有效性。实验结果表明，GWT在降低内存占用的同时，能够保持甚至超过现有内存优化方法的性能。具体来说，GWT在某些任务上达到了与全秩更新相当的性能，同时显著降低了内存需求。与先进的内存高效优化器相比，GWT在内存使用和训练性能方面均取得了显著的提升。

🎯 应用场景

该研究成果可广泛应用于大型语言模型的训练和微调，尤其是在资源受限的环境下，例如移动设备或边缘计算平台。通过降低内存需求，GWT可以使更大规模的模型能够在更小的硬件上进行训练，从而加速AI技术的普及和应用。此外，该方法也可以应用于其他需要大量内存的机器学习任务，例如图像生成和强化学习。

📄 摘要（原文）

Large language models (LLMs) have shown impressive performance across a range of natural language processing tasks. However, their vast number of parameters introduces significant memory challenges during training, particularly when using memory-intensive optimizers like Adam. Existing memory-efficient algorithms often rely on techniques such as singular value decomposition projection or weight freezing. While these approaches help alleviate memory constraints, they generally produce suboptimal results compared to full-rank updates. In this paper, we investigate the memory-efficient method beyond low-rank training, proposing a novel solution called Gradient Wavelet Transform (GWT), which applies wavelet transforms to gradients in order to significantly reduce the memory requirements for maintaining optimizer states. We demonstrate that GWT can be seamlessly integrated with memory-intensive optimizers, enabling efficient training without sacrificing performance. Through extensive experiments on both pre-training and fine-tuning tasks, we show that GWT achieves state-of-the-art performance compared with advanced memory-efficient optimizers and full-rank approaches in terms of both memory usage and training performance.