Reversing Large Language Models for Efficient Training and Fine-Tuning
This addresses the memory bottleneck for researchers and practitioners training or fine-tuning LLMs, offering a scalable and efficient path, though it builds on existing reversible network concepts.
The authors tackled the problem of high memory consumption in training and fine-tuning large language models by introducing reversible architectures that retrieve hidden states during backpropagation instead of storing activations, resulting in a drastic reduction in memory usage and improved throughput while maintaining comparable or improved performance on several benchmarks.
Large Language Models (LLMs) are known for their expensive and time-consuming training. Thus, oftentimes, LLMs are fine-tuned to address a specific task, given the pretrained weights of a pre-trained LLM considered a foundation model. In this work, we introduce memory-efficient, reversible architectures for LLMs, inspired by symmetric and symplectic differential equations, and investigate their theoretical properties. Different from standard, baseline architectures that store all intermediate activations, the proposed models use time-reversible dynamics to retrieve hidden states during backpropagation, relieving the need to store activations. This property allows for a drastic reduction in memory consumption, allowing for the processing of larger batch sizes for the same available memory, thereby offering improved throughput. In addition, we propose an efficient method for converting existing, non-reversible LLMs into reversible architectures through fine-tuning, rendering our approach practical for exploiting existing pre-trained models. Our results show comparable or improved performance on several datasets and benchmarks, on several LLMs, building a scalable and efficient path towards reducing the memory and computational costs associated with both training from scratch and fine-tuning of LLMs.