OASIS: Online Activation Subspace Learning for Memory-Efficient Training

📄 arXiv: 2604.09406v1 📥 PDF

作者: Sakshi Choudhary, Utkarsh Saxena, Kaushik Roy

分类: cs.LG

发布日期: 2026-04-10

💡 一句话要点

OASIS：在线激活子空间学习，用于内存高效的大模型训练

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

关键词: 大语言模型 内存优化 激活子空间 在线学习 低秩分解

📋 核心要点

  1. 大型语言模型训练面临激活内存瓶颈，现有方法如低秩分解或架构修改存在局限性。
  2. OASIS算法通过在线学习并更新激活子空间，将激活投影到低维空间，从而减少内存占用。
  3. 实验表明，OASIS在微调和预训练任务中，峰值内存降低高达2倍，性能与全量微调相当。

📝 摘要（中文）

训练大型语言模型(LLMs)受到内存需求的限制，其中激活占据了总内存占用的大部分。现有方法通过低秩权重参数化或优化器状态的低秩梯度子空间来减少内存，而激活内存则通过架构修改或基于周期性更新投影的压缩方案来解决。我们提出了OASIS，一种在线激活子空间学习算法，用于内存高效的训练，该算法在训练期间跟踪并持续更新低维激活子空间。中间激活被投影到这个不断演化的子空间上，减少了内存，而无需修改前向传递计算。演化的激活子空间诱导了低秩梯度表示，使得梯度和优化器状态都可以直接在这个子空间中维护，同时一个感知投影的优化器在子空间更新中一致地传输优化器状态，以实现稳定的训练。在各种微调和预训练任务中，OASIS实现了比完整微调低至2倍的峰值内存，同时匹配其性能并优于先前的低秩方法。

🔬 方法详解

问题定义：大型语言模型训练过程中，激活值占用了大量的内存，成为训练的瓶颈。现有方法，例如低秩权重参数化、梯度子空间方法、架构修改和基于周期性更新投影的压缩方案，虽然在一定程度上缓解了内存问题，但或者会影响模型性能，或者需要复杂的工程实现。因此，需要一种更高效、更通用的方法来减少激活内存占用，同时保持模型性能。

核心思路：OASIS的核心思想是在训练过程中动态地学习并维护一个低维的激活子空间。通过将中间激活投影到这个子空间，可以显著减少内存占用，而无需修改前向计算过程。此外，该方法还利用这个低秩的激活子空间来表示梯度和优化器状态，进一步降低内存需求。

技术框架：OASIS算法主要包含以下几个关键模块：1) 在线激活子空间学习模块：该模块负责在训练过程中动态地学习和更新激活子空间。具体来说，它会跟踪中间激活，并使用例如奇异值分解（SVD）等方法来提取主要的激活模式，从而构建低维子空间。2) 激活投影模块：该模块将中间激活投影到学习到的低维子空间中，从而减少内存占用。3) 投影感知优化器：该优化器负责在子空间更新时，将优化器状态正确地传输到新的子空间中，以保证训练的稳定性。

关键创新：OASIS的关键创新在于其在线学习激活子空间的能力。与需要预先确定子空间的静态方法不同，OASIS可以根据训练过程中激活的实际分布动态地调整子空间，从而更好地适应模型的学习过程。此外，OASIS还提出了一种投影感知优化器，解决了子空间更新带来的优化器状态不一致问题，保证了训练的稳定性。

关键设计：OASIS的关键设计包括：1) 子空间更新频率：需要根据实际情况调整子空间更新的频率，以平衡计算开销和子空间的适应性。2) 子空间维度选择：子空间的维度需要根据内存限制和模型性能要求进行权衡。维度太小可能会导致信息损失，影响模型性能；维度太大则无法有效降低内存占用。3) 投影感知优化器的实现：需要仔细设计优化器的更新规则，以确保在子空间更新时，优化器状态能够正确地传输到新的子空间中。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

OASIS在各种微调和预训练任务中表现出色，实现了比完整微调低至2倍的峰值内存占用，同时保持了与全量微调相当的性能。此外，OASIS的性能优于先前的低秩方法，证明了其在线学习激活子空间的有效性。这些实验结果表明，OASIS是一种非常有前景的内存高效训练方法。

🎯 应用场景

OASIS算法可广泛应用于各种需要训练大型语言模型的场景，例如自然语言处理、机器翻译、文本生成等。该方法能够显著降低训练过程中的内存需求，使得在资源受限的设备上训练大型模型成为可能。此外，OASIS还可以加速模型训练，提高开发效率，并促进大型语言模型在更多领域的应用。

📄 摘要（原文）

Training large language models (LLMs) is constrained by memory requirements, with activations accounting for a substantial fraction of the total footprint. Existing approaches reduce memory using low-rank weight parameterizations or low-rank gradient subspaces for optimizer states, while activation memory is addressed through architectural modifications or compression schemes based on periodically updated projections. We propose OASIS, an online activation subspace learning algorithm for memory-efficient training that tracks and continuously updates a low-dimensional activation subspace during training. Intermediate activations are projected onto this evolving subspace, reducing memory without modifying forward-pass computations. The evolving activation subspace induces low-rank gradient representations, enabling both gradients and optimizer states to be maintained directly in this subspace, while a projection-aware optimizer consistently transports optimizer states across subspace updates for stable training. Across various finetuning and pretraining tasks, OASIS achieves up to $2\times$ lower peak memory than full fine-tuning while matching its performance and outperforming prior low-rank methods.