通过选择性回传部分输入标记,显著降低大模型微调内存占用。
Memory-Efficient Fine-Tuning of Transformers via Token Selection
- 仅缓存部分输入标记的中间激活值,减少前向传播内存开销。
- 在少样本任务中性能接近全量微调,内存消耗降低超50%。
- 可与LoRA等方法结合,适合资源受限场景下的大模型定制化训练。
微调是使预训练模型适应下游任务的有效手段,但对大型Transformer模型(如大语言模型)常带来高昂的内存开销。现有方法虽能减少部分内存需求,但仍需缓存前向传播中所有中间激活值以进行反向传播更新。本文提出TokenTune,通过在反向传播中仅对输入标记的子集进行梯度反传,大幅降低中间激活值的缓存需求。实验表明,该方法在包含数十亿参数的预训练模型上,于文本分类与问答等少样本任务中表现与全量微调或主流高效微调方法相当,同时显著降低内存占用,尤其在与LoRA等互补机制结合时效果更优。代码已开源。本工作旨在推动大模型在特定领域定制或系统级联合训练中的应用。
原文摘要 · Abstract (English)
Fine-tuning provides an effective means to specialize pre-trained models for various downstream tasks. However, fine-tuning often incurs high memory overhead, especially for large transformer-based models, such as LLMs. While existing methods may reduce certain parts of the memory required for fine-tuning, they still require caching all intermediate activations computed in the forward pass to update weights during the backward pass. In this work, we develop TokenTune, a method to reduce memory usage, specifically the memory to store intermediate activations, in the fine-tuning of transformer-based models. During the backward pass, TokenTune approximates the gradient computation by backpropagating through just a subset of input tokens. Thus, with TokenTune, only a subset of intermediate activations are cached during the forward pass. Also, TokenTune can be easily combined with existing methods like LoRA, further reducing the memory cost. We evaluate our approach on pre-trained transformer models with up to billions of parameters, considering the performance on multiple downstream tasks such as text classification and question answering in a few-shot learning setup. Overall, TokenTune achieves performance on par with full fine-tuning or representative memory-efficient fine-tuning methods, while greatly reducing the memory footprint, especially when combined with other methods with complementary memory reduction mechanisms. We hope that our approach will facilitate the fine-tuning of large transformers, in specializing them for specific domains or co-training them with other neural components from a larger system. Our code is available at https://github.com/facebookresearch/tokentune.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。