GaLore：基于梯度低秩投影的大语言模型高效训练方法详解一

news2025/6/8 4:18:41

📘 GaLore：基于梯度低秩投影的大语言模型高效训练方法详解

一、论文背景与动机

随着大语言模型（LLM）参数规模的不断增长，例如 GPT-3（175B）、LLaMA（65B）、Qwen（100B+），在训练过程中所需的显存和计算资源成为了一个重大瓶颈。

传统的训练方式如 AdamW 需要存储每个参数的梯度、动量、以及二阶矩估计值，导致：

显存占用巨大
训练成本高昂
只能在高配 GPU 上运行

为了解决这个问题，来自 UC Berkeley 的团队提出了 GaLore（Gradient Low-Rank Projection），一种基于梯度低秩近似 的新方法，能够在显著降低训练显存消耗的同时，保持模型性能不变甚至提升。

二、GaLore 的核心思想

1. 梯度矩阵的低秩特性观察

作者通过大量实验发现：Transformer 中大多数层的梯度矩阵具有显著的低秩结构，即梯度信息集中在少数几个主成分上。

💡 举例：

在 ViT、BERT、LLM 等模型中，梯度矩阵的奇异值快速衰减。
这意味着我们不需要完整地保留所有梯度维度，只需要保留前 r 个主成分即可近似整个梯度方向。

2. 核心操作流程

GaLore 的训练流程分为以下几个关键步骤：

（1）梯度投影到低秩空间

对于每一层权重 $ W \in \mathbb{R}^{m \times n} $，其梯度 $ \nabla_W L $ 是一个 $ m \times n $ 的矩阵。GaLore 将其进行 SVD 分解，并只保留前 r 个奇异向量：

$\nabla_W L \approx U_r \Sigma_r V_r^T$

其中 $ U_r, V_r $ 是左右奇异向量矩阵，$ \Sigma_r $ 是对角矩阵。

（2）在低秩空间中更新参数

仅对低秩空间中的参数进行更新：

$\tilde{W} = W + \eta \cdot U_r \cdot \text{Adam}(U_r^T \nabla_W L V_r) \cdot V_r^T$

其中 $ \eta $ 是学习率。

（3）反投影回原始空间

更新后的参数 $ \tilde{W} $ 已经是低秩修正后的结果，可以直接应用于下一轮训练。

三、GaLore 的优势与特点

特性	描述
显存节省	相比 AdamW，节省高达 3~5 倍显存
不影响精度	实验显示 GaLore 在多种任务上性能接近甚至超过标准 AdamW
兼容性强	支持所有主流优化器（AdamW、SGD with Momentum、LAMB 等）
多模态支持	可用于视觉、语言、多模态任务（ViT、CLIP、VQA 等）
易于集成	可直接替换现有优化器，无需修改模型结构

四、技术实现细节

1. 梯度压缩过程

GaLore 并不是对参数本身做低秩约束，而是对梯度进行低秩投影，从而减少每次更新所需的信息量。

# 伪代码示意
def galore_update(W, grad, optimizer):
    U, S, V = torch.svd_lowrank(grad, rank=r)
    low_rank_grad = U @ torch.diag(S) @ V.T
    # 使用低秩梯度更新参数
    updated = optimizer.step(low_rank_grad)
    return updated

2. 动量与 Adam 的适配

GaLore 支持动量机制和 AdamW 的变体，只需将动量和方差也投影到低秩空间中：

动量项：$ m_t = \beta_1 m_{t-1} + (1-\beta_1)\nabla $
方差项：$ v_t = \beta_2 v_{t-1} + (1-\beta_2)\nabla^2 $

这些项也被投影到低秩空间，大大减少了内存开销。

3. 层级控制策略

不同层的梯度低秩程度不同，因此 GaLore 提出了层级自适应投影策略（Layer-wise Adaptive Rank Selection），根据每层梯度的奇异值分布自动选择合适的秩 r。

五、实验验证与性能表现

作者在多个任务上进行了广泛的实验验证，包括：

任务类型	数据集	模型
图像分类	ImageNet	ViT-B/16
自然语言处理	GLUE	BERT-base
大语言模型	WikiText-2、OpenWebText	GPT-2 small
多模态理解	VQA、ImageNet	CLIP-ViT-B/16

实验结果亮点：

在 ImageNet 上使用 ViT-B/16，GaLore 节省了 4.8x 显存，准确率仅下降 0.3%
在 GLUE 任务上，BERT 使用 GaLore 后显存减少 4.2x，平均指标损失小于 1.5%
在 GPT-2 small 上，GaLore 达到了与 AdamW 相当的语言建模性能，但显存需求更低
在多模态任务中，GaLore 在 CLIP 上表现出色，尤其在图像检索任务中几乎无损性能

六、与其他参数效率训练方法的对比

方法	显存节省	是否改变模型结构	是否依赖预训练	是否适用于多模态
LoRA	✅	✅	✅	✅
AdaLoRA	✅	❌	✅	✅
BitFit / DiffPruning	✅	❌	❌	❌
GaLore	✅✅✅	❌	❌	✅✅

✅ 表示支持，❌ 表示不支持或受限。

GaLore 的最大优势在于：它不改变模型结构，也不需要额外预训练，就可以实现显存节约和训练加速。

七、GaLore 的适用场景

场景	说明
本地训练 LLaMA 系列模型	如 LLaMA-7B、CodeLlama、TinyLlama 等
教学科研平台	在消费级 GPU（如 RTX 3090、4090）上训练大模型
显存敏感型部署	当显存成为训练瓶颈时，GaLore 可作为首选优化器
多模态模型训练	如 CLIP、Flamingo、BLIP 等，可大幅降低训练成本