Abstract
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 this https URL.
Abstract (translated)
微调提供了一种有效的方法，可以将预训练模型专门化为各种下游任务。然而，微调通常会产生较高的内存开销，尤其是对于基于转换器的大型语言模型（如LLM）。尽管现有方法可能减少了某些部分在微调过程中所需的内存，但它们仍然需要缓存前向传播计算的所有中间激活以更新反向传播过程中的权重。在此项工作中，我们开发了TokenTune方法，该方法旨在减少基于转换器模型的微调中用于存储中间激活的内存使用量。在反向传播期间，TokenTune通过仅对输入令牌的一个子集进行反向传播来近似梯度计算。因此，在前向传播过程中，只需要缓存一个部分的中间激活。此外，TokenTune可以轻松地与现有方法（如LoRA）结合使用，进一步减少内存成本。我们在具有数十亿参数的预训练转换器模型上评估了我们的方法，并考虑了在少量样本学习设置下的多个下游任务性能，例如文本分类和问答。总体而言，TokenTune实现了与完整微调或代表性内存高效微调方法相当的性能，同时大大减少了内存占用量，尤其是在与其他具有互补内存减少机制的方法结合使用时更是如此。我们希望我们的方法将有助于大型转换器模型的微调工作，以专门针对特定领域或在更大的系统中与其他神经组件进行联合训练。我们的代码可在以下链接获取：[此URL]。
URL
https://arxiv.org/abs/2501.18824
