Reversing Large Language Models for Efficient Training and Fine-Tuning

TL;DR

This paper introduces reversible architectures for large language models that significantly reduce memory usage during training and fine-tuning, enabling larger batch processing and improved efficiency without sacrificing performance.

Contribution

The authors propose memory-efficient, reversible LLM architectures inspired by differential equations and a method to convert existing models into reversible ones through fine-tuning.

Findings

Comparable or improved performance on multiple datasets

Significant memory reduction during training

Enables larger batch sizes for the same hardware constraints

Abstract

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.