Stanford’s ReFT fine-tunes LLMs at a fraction of the cost

robot concepts
Image generated with Bing Image Creator

This article is part of our coverage of the latest in AI research.

Researchers at Stanford University have developed a new technique that can fine-tune large language models (LLM) while modifying less than 1% of their representations. 

The technique, called Representation Fine-Tuning (ReFT), seeks out specific parts of the model that are relevant to the task for which it is being fine-tuned. And it turns out that for most downstream tasks, a very small number of weights need to be modified.

Experiments show that ReFT outperforms other parameter-efficient fine-tuning (PEFT) techniques, including the popular low-rank adaptation (LoRA). The researchers have released the code for ReFT, and it is compatible with other fine-tuning libraries.

Problems with PEFT

Parameter-efficient finetuning (PEFT) methods reduce the costs of fine-tuning LLMs by updating only a small fraction of the model’s weights. PEFTs have been shown to achieve similar performance to full fine-tuning in many practical settings. Adapters such as LoRA can be used as an independent fine-tuned component that can be combined with the model at inference time. However, current state-of-the-art PEFTs modify weights rather than representations. 

“PEFTs lack the notion of time; they usually update representations across all layers and positions using the same weights,” Zhengxuan Wu and Aryaman Arora, PhD students at Stanford and lead authors of the paper, told TechTalks. “Earlier representation editing-methods like prefix-tuning, on the other hand, are not competitive with new PEFTs like LoRA. We wanted the best of both worlds.”

(Prefix-tuning keeps the weights of the main model frozen and trains an embedding model to add soft tokens to the input in order to fine-tune the behavior of the model toward the desired task.)

Previous work on interpreting deep learning systems has shown that models learn representations that encode rich semantic information. This means that there are subspaces within the model that represent specific concepts, suggesting that editing representations might be a more powerful alternative to weight updates.

“Our previous work shows that learned interventions on representation subspaces are powerful: they can find useful subspaces for controlling model behavior even when the model has not learned the task,” Wu and Arora said. “We thought this approach may be useful for controlling models in a top-down manner.”

ReFT and LoReFT

Representation Fine-Tuning (ReFT) is a family of methods based on the hypothesis that modifying models based on learned representations can be much more efficient than current PEFT methods.

ReFT methods operate on a frozen base model and learn task-specific interventions on hidden representations. Instead of adapting model weights, ReFT methods train interventions that manipulate a small fraction of model representations to steer model behaviors to solve downstream tasks at inference time. ReFT methods can be drop-in replacements for weight-based PEFTs. 

The researchers created a highly efficient ReFT method called Low-rank Linear Subspace ReFT (LoReFT). “This is a parametrization of ReFT that intervenes on hidden representations in the linear subspace spanned by a low-rank projection matrix,” the researchers write. Basically, LoReFT is the representation-based equivalent of LoRA.

The researchers tested LoReFT on both autoregressive models such as Llama and bidirectional encoders such as RoBERTa. The fine-tuned models ranged from 125 million to 13 billion parameters.

They tested the fine-tuned models on several benchmarks including common-sense reasoning, arithmetic reasoning, instruction-following, and natural language understanding. Compared to LoRA, LoReFT uses 10–50 times fewer parameters while achieving state-of-the-art performance for several datasets, including commonsense reasoning tasks. 

On arithmetic reasoning, LoReFT lags behind LoRA but continues to improve with scale. One interesting finding was in instruction-following tasks, where LoReFT outperformed PEFT methods and full fine-tuning, and achieved a win rate of within 1% of GPT-3.5 Turbo. Even after halving the parameter count, LoReFT still outperformed other fine-tuning methods. 

“These findings indicate that ReFT methods are worthy of further exploration, as they may emerge as more efficient and effective alternatives to weight-based PEFTs,” the researchers write.

The researchers released pyreft, an open-source Python library made for training and sharing ReFTs. Pyreft can fine-tune any pre-trained language model on HuggingFace.

The main limitation of ReFT is the large hyperparameter search space, which makes it difficult to find the best interventions. Wu and Arora told TechTalks that in the future, they will further explore the potential of ReFT in other modalities, including visual data. 

Leave a Reply

This site uses Akismet to reduce spam. Learn how your comment data is processed.