VeLoRA: Memory Efficient Training using Rank-1 Sub-Token Projections

📄 arXiv: 2405.17991v2 📥 PDF

作者: Roy Miles, Pradyumna Reddy, Ismail Elezi, Jiankang Deng

分类: cs.CV, cs.AI

发布日期: 2024-05-28 (更新: 2024-10-21)

备注: NeurIPS 2024. Code available at https://github.com/roymiles/VeLoRA

💡 一句话要点

VeLoRA：利用Rank-1子Token投影实现内存高效的LLM训练

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

关键词: 大型语言模型 参数高效微调 低秩近似 内存优化 子Token投影

📋 核心要点

  1. 大型语言模型训练和微调面临计算和内存资源的巨大挑战。
  2. VeLoRA通过将tokens投影到低维空间，显著降低了内存占用，同时保持性能。
  3. 实验表明，VeLoRA在微调和预训练任务上均表现出色，优于或媲美现有方法。

📝 摘要（中文）

大型语言模型（LLMs）已经成为处理多种语言任务的强大工具。然而，训练和微调这些模型在计算和内存方面仍然非常密集。本文旨在识别并描述使用梯度下降进行有效模型收敛所需的重要组成部分。研究发现，在不降低性能的情况下，用于实现反向传播的中间激活可以被过度压缩。基于此，本文提出了一种廉价且内存高效的算法，用于微调和预训练LLMs。该算法在正向传播过程中，简单地将tokens分割成更小的子tokens，然后将它们投影到固定的1维子空间。在反向传播过程中，这些特征被粗略地重建以实现更新规则。实验表明，该算法与VTAB-1k微调基准上的许多最先进的PEFT方法互补。此外，在微调LLaMA时，该算法优于QLoRA，并在大规模C4数据集上展现出与其他内存高效预训练方法相媲美的性能。

🔬 方法详解

问题定义：现有的大型语言模型（LLMs）训练和微调过程需要消耗大量的计算资源和内存。尤其是在反向传播过程中，中间激活的存储成为了一个瓶颈。现有的参数高效微调（PEFT）方法虽然在一定程度上缓解了这个问题，但仍然存在改进空间。

核心思路：VeLoRA的核心思路是在正向传播过程中，将tokens分解为更小的子tokens，然后将这些子tokens投影到一个固定的低秩（Rank-1）子空间。这样可以显著减少需要存储的激活数量，从而降低内存消耗。在反向传播过程中，再对这些低维特征进行粗略重建，用于参数更新。

技术框架：VeLoRA算法主要包含以下几个步骤：1) Token分割：将输入的tokens分割成更小的子tokens。2) 投影：将子tokens投影到预定义的Rank-1子空间。3) 前向传播：利用投影后的低维特征进行前向计算。4) 反向传播：根据低维特征重建梯度，并进行参数更新。

关键创新：VeLoRA的关键创新在于利用Rank-1投影来压缩中间激活，从而在保证模型性能的同时，显著降低了内存需求。与传统的量化方法（如QLoRA）相比，VeLoRA通过子token投影的方式，更好地保留了token之间的关系，从而提高了微调效果。

关键设计：VeLoRA的关键设计包括：1) Rank-1子空间的选取：子空间的选择会影响模型的性能，论文中可能采用了随机初始化或其他策略来选择合适的子空间。2) 子token的大小：子token的大小需要根据具体的任务和模型进行调整，以平衡内存消耗和性能。3) 重建方法：反向传播过程中，如何从低维特征重建梯度，也是一个重要的设计考虑因素。具体的重建方法可能包括简单的线性插值或其他更复杂的上采样技术。

🖼️ 关键图片

fig_0
fig_1

📊 实验亮点

VeLoRA在VTAB-1k微调基准上与许多最先进的PEFT方法互补，表明其具有良好的泛化能力。在微调LLaMA时，VeLoRA优于QLoRA，证明了其在内存效率和性能方面的优势。在大规模C4数据集上的预训练实验中，VeLoRA展现出与其他内存高效预训练方法相媲美的性能，进一步验证了其有效性。

🎯 应用场景

VeLoRA具有广泛的应用前景，尤其是在资源受限的环境下训练和微调大型语言模型。例如，可以在消费级GPU上进行LLaMA等模型的微调，或者在移动设备上部署轻量级的语言模型。此外，VeLoRA还可以应用于其他需要处理大规模数据的机器学习任务，例如图像识别和语音识别。

📄 摘要（原文）

Large language models (LLMs) have recently emerged as powerful tools for tackling many language-processing tasks. Despite their success, training and fine-tuning these models is still far too computationally and memory intensive. In this paper, we identify and characterise the important components needed for effective model convergence using gradient descent. In doing so we find that the intermediate activations used to implement backpropagation can be excessively compressed without incurring any degradation in performance. This result leads us to a cheap and memory-efficient algorithm for both fine-tuning and pre-training LLMs. The proposed algorithm simply divides the tokens up into smaller sub-tokens before projecting them onto a fixed 1-dimensional subspace during the forward pass. These features are then coarsely reconstructed during the backward pass to implement the update rules. We confirm the effectiveness of our algorithm as being complimentary to many state-of-the-art PEFT methods on the VTAB-1k fine-tuning benchmark. Furthermore, we outperform QLoRA for fine-tuning LLaMA and show competitive performance against other memory-efficient pre-training methods on the large-scale C4 dataset.