PaperlayerView exact revision on arXiv

Source-linked AI summary

ReFT: Representation Finetuning for Language Models

Zhengxuan Wu, Aryaman Arora, Zheng Wang, Atticus Geiger, Dan Jurafsky, Christopher D. Manning, Christopher Potts

arXiv:2404.03592v3cs.CLcs.AIcs.LG

TL;DR

Large-model adaptation is costly, and weight-based PEFT may underuse the semantic information in hidden representations. The paper introduces ReFT, including LoReFT and DiReFT, which learn interventions on a frozen model; across diverse evaluations, LoReFT combines strong performance with 15×–65× greater parameter efficiency than LoRA.

  • Problem

    Adapting large language models by finetuning weights is expensive, while prior PEFTs modify weights despite evidence that representations encode rich semantic information.

  • Method

    ReFT learns task-specific interventions on hidden representations of a frozen base model; LoReFT restricts edits to low-rank linear subspaces, while DiReFT is a more efficient ablation.

  • Results

    Across commonsense reasoning, arithmetic reasoning, instruction-following, and natural language understanding, ReFTs provide a strong efficiency–performance balance; LoReFT uses 15×–65× fewer parameters than LoRA and achieves state-of-the-art results on several domains.

  • Takeaways & Limitations

    Representation interventions can serve as effective, highly parameter-efficient alternatives to weight-based PEFTs, with LoReFT also demonstrating strong long-form instruction-following performance.

  • Takeaways & Limitations

    The evaluation mainly explores LLaMA-family models, and the precise causal mechanisms by which ReFT works remain unresolved.

Abstract

from arXiv · show

Parameter-efficient finetuning (PEFT) methods seek to adapt large neural models via updates to a small number of weights. However, much prior interpretability work has shown that representations encode rich semantic information, suggesting that editing representations might be a more powerful alternative. We pursue this hypothesis by developing a family of Representation Finetuning (ReFT) methods. ReFT methods operate on a frozen base model and learn task-specific interventions on hidden representations. We define a strong instance of the ReFT family, Low-rank Linear Subspace ReFT (LoReFT), and we identify an ablation of this method that trades some performance for increased efficiency. Both are drop-in replacements for existing PEFTs and learn interventions that are 15x--65x more parameter-efficient than LoRA. We showcase LoReFT on eight commonsense reasoning tasks, four arithmetic reasoning tasks, instruction-tuning, and GLUE. In all these evaluations, our ReFTs deliver the best balance of efficiency and performance, and almost always outperform state-of-the-art PEFTs. We release a generic ReFT training library publicly at https://github.com/stanfordnlp/pyreft.

1 Introduction

PEFT reduces the cost of adapting large language models by updating few weights, while ReFT instead learns interventions on hidden representations. LoReFT delivers strong benchmark performance with substantially fewer parameters than LoRA.

  • PEFT reduces memory usage and training time by updating a small number of weights instead of fully finetuning large language models.
  • ReFT learns task-specific interventions on hidden representations rather than adapting model weights, serving as a drop-in replacement for weight-based PEFTs.
  • Across four benchmarks and multiple model families, LoReFT achieves competitive or state-of-the-art performance despite training far fewer parameters than existing PEFTs.
  • LoReFT uses 15×–65× fewer parameters than LoRA while achieving state-of-the-art performance on commonsense reasoning, instruction-following, and natural language understanding.

2 Related work

Prior PEFTs modify weights, add modules, or train soft prompt tokens, while representation-editing and interpretability work motivates controlling language models through their internal representations.

  • Parameter-efficient finetuning methods: Adapter-based PEFTs add trainable modules to frozen models, but these components impose additional inference burden because they cannot be easily folded into model weights.
  • Parameter-efficient finetuning methods: LoRA and DoRA approximate additive weight updates with low-rank matrices that can be merged into the model without additional inference overhead.
  • Parameter-efficient finetuning methods: Prompt-based methods train randomly initialized soft input tokens while keeping language-model weights frozen, but often have inferior performance and significant inference overhead.
  • Representation editing: Representation-editing methods steer generations or erase concepts by modifying residual-stream activations without resource-intensive finetuning.
  • Interventional interpretability: Interventions on representation subspaces provide evidence that human-interpretable concepts are encoded linearly in model representations.

3 ReFT

ReFT is a framework for modifying selected hidden representations during a frozen Transformer’s forward pass. LoReFT performs low-rank subspace edits, while DiReFT removes constraints for greater efficiency.

  • Motivation: Distributed interchange intervention tests whether concepts are encoded in linear subspaces by replacing a base representation’s subspace with a counterfactual source representation.
  • Two low-rank ReFT instantiations: LoReFT adapts distributed interchange intervention by learning representation edits that steer model predictions toward task labels.
  • Two low-rank ReFT instantiations: LoReFT edits hidden representations only within a low-rank subspace, using an orthonormal projection R and a learned projected source Wh + b while freezing the language model.
  • Two low-rank ReFT instantiations: DiReFT removes LoReFT’s orthogonality constraint and difference operation, reducing training time while retaining a representation-level low-rank formulation.
  • Training objectives: The framework supports generation and classification objectives, including language modeling with teacher forcing and classification heads for final-layer representations.
  • The ReFT family of methods: ReFT interventions apply learned functions to selected positions and layers, overwrite representations, and thereby affect computations in subsequent layers.
  • The ReFT family of methods: A ReFT method is a set of non-overlapping interventions whose functions have independent parameters.

4 Experiments

The experiments compare LoReFT and DiReFT with existing PEFTs across commonsense reasoning, arithmetic reasoning, instruction tuning, and GLUE. LoReFT achieves leading commonsense performance, while arithmetic results vary and GLUE performance is comparable to PEFTs.

  • Hyperparameter configuration: LoReFT and DiReFT edit fixed prefix and suffix positions across selected layers, with tied or untied intervention parameters.The four tuned choices are prefix count, suffix count, intervened layers, and whether parameters are tied across positions within a layer.
  • Hyperparameter configuration: The fixed number of edited positions gives LoReFT and DiReFT an additional inference cost that does not scale with prompt length.This position choice is the main additional consideration relative to LoRA in the stated hyperparameter search.
  • Commonsense reasoning: LoReFT sets state-of-the-art performance on eight commonsense reasoning tasks, outperforming other methods by a considerable margin.DiReFT is more compute-efficient but achieves only slightly worse performance consistently.
  • Arithmetic reasoning: Arithmetic results vary substantially: LoReFT and DiReFT trail LoRA and adapters but outperform prefix-tuning, with no clear winner across all tasks.The authors suggest longer generations and greater task difficulty may make chain-of-thought reasoning harder for ReFTs; performance is better with the 13B than the 7B model.
  • Natural language understanding: On GLUE, LoReFT has comparable performance with PEFT methods at matched parameter count, whereas DiReFT performs worse than most PEFTs.The comparison covers RoBERTa-base and RoBERTa-large.

5 Limitations

The paper identifies limitations involving evaluation practice, model-family coverage, and incomplete understanding of ReFT’s causal mechanisms and interpretability implications.

  • The experiments mainly explore LLaMA-family models, leaving ReFT’s effectiveness on other model families and vision–language models for future work.
  • The large hyperparameter search space means ReFT’s capabilities have not yet been fully explored.
  • The precise mechanisms by which representation interventions affect model behavior remain insufficiently understood.
  • The paper leaves tracking upstream causal effects and modifying complex causal pathways with more structured ReFTs to future research.
  • Test-set hill-climbing in PEFT literature can overfit specific tasks and impede fair comparison.
  • The authors argue for compute- or time-matched hyperparameter tuning that excludes test-set-based tuning and model selection.

6 Conclusion

The conclusion presents LoReFT and DiReFT as efficient ReFT alternatives to PEFTs and describes a public library for training and sharing ReFTs.

  • LoReFT is proposed as a strong PEFT alternative, while DiReFT trades some performance for increased efficiency.
  • 15×–65× more efficiency than LoRA is reported for LoReFT across benchmarks from four domains.
  • LoReFT establishes new state-of-the-art performance on commonsense reasoning, instruction-following, and natural language understanding against the strongest PEFTs.
  • The authors release a Python library for training and sharing ReFTs.
  • The library supports pretrained HuggingFace language models through activation-intervention training infrastructure.

B Describing existing methods under the ReFT framework

The ReFT framework expresses existing representation-editing methods as learned interventions over hidden representations, while distinguishing them from weight-based PEFTs.

  • Existing representation-editing methods can be cast as ReFTs, demonstrating the framework’s expressivity.
  • Unlike PEFT frameworks, ReFT represents time or sequence and can intervene on only selected positions.
  • RED: RED applies learned scaling and bias transformations to hidden representations at every layer and position.
  • Activation addition: Activation addition computes a scaled difference between activations from contrastive prompts and adds it to representations at all positions in a layer.
  • RepE: Representation engineering derives reading or contrast vectors from prompt activations and uses operators to intervene on selected positions.
  • LoRRA: LoRRA is classified as a PEFT because it tunes model weights using a LoRA variant rather than intervening directly on representations.

C Datasets

The evaluation covers commonsense reasoning, arithmetic reasoning, instruction following, and GLUE, using separate data splits and task-specific evaluation procedures.

  • Commonsense reasoning: The commonsense evaluation uses eight open-ended question-answering datasets spanning different domains.
  • Arithmetic reasoning: The arithmetic evaluation uses seven datasets covering arithmetic and mathematical word problems.
  • Natural language understanding: GLUE evaluation splits validation data into in-training evaluation and testing subsets, using task-appropriate metrics.
  • Hyperparameter tuning and decoding: Commonsense hyperparameters are selected on a standalone GSM8K development set, with settings chosen on LLaMA-7B and reused on LLaMA-13B.
  • Instruction following: Instruction-following hyperparameters are selected on Alpaca-52K using Alpaca-Eval v1.0 and GPT-4 Turbo annotation.
  • Natural language understanding: LoReFT and DiReFT receive separate hyperparameter tuning for a fair comparison.
  • Hyperparameter tuning and decoding: ReFT can be sensitive to hyperparameter settings, with performance depending on intervention positions and the number of layers or tokens edited.

D.3 Additional hyperparameter-tuning results of LoReFT

Matched-hyperparameter experiments test whether LoReFT’s gains depend on training longer than competing PEFTs, while the appendix tables compare it across multiple benchmark families. The results support similar performance under matched training budgets, with comparisons covering commonsense, arithmetic, and GLUE settings.

  • D.3 Additional hyperparameter-tuning results of LoReFT: With matched hyperparameters, LoReFT significantly outperforms previous methods on eight commonsense reasoning datasets.The reruns use the same number of epochs and effective batch size as LoRA or DoRA.
  • D.3 Additional hyperparameter-tuning results of LoReFT: Table 14 compares LLaMA-7B and LLaMA-13B accuracy across eight commonsense reasoning datasets and reports parameter percentages for LoReFT.Baseline results are taken from prior work, while LoReFT results average three runs with distinct random seeds.
  • D.3 Additional hyperparameter-tuning results of LoReFT: Table 15 compares LLaMA-7B and LLaMA-13B accuracy across four arithmetic reasoning datasets under matched epoch conditions.The comparison reports averaged results from three distinct random seeds for the authors’ method.
  • D.3 Additional hyperparameter-tuning results of LoReFT: Table 16 compares RoBERTa-base and RoBERTa-large against PEFT methods on GLUE using median performance over five distinct random seeds.The baseline results are taken from Kopiczko et al. (2024).
  • D.3 Additional hyperparameter-tuning results of LoReFT: Table 17 reports arithmetic-reasoning ablation results for LLaMA-7B and LLaMA-13B with standard deviations across three runs.All variants use LoReFT hyperparameters except DiReFT.

E Ablating the parametrisation of LoReFT

The section ablates LoReFT’s intervention parametrisation on arithmetic reasoning, comparing alternatives that remove orthogonality constraints, reduce computation, or reduce parameter count. Variants with similar parameter counts achieve similar performance to LoReFT across both evaluated models.

  • E Ablating the parametrisation of LoReFT: The ablations reevaluate LLaMA-1 7B and 13B on arithmetic reasoning using the same hyperparameters across variants of the LoReFT intervention function.Arithmetic reasoning is selected because it is LoReFT’s most difficult benchmark and trains relatively quickly.
  • E Ablating the parametrisation of LoReFT: One variant removes the orthonormal constraint while preserving LoReFT’s trainable parameter count and reducing memory overhead.It uses two low-rank nonorthogonal projection matrices.
  • E Ablating the parametrisation of LoReFT: Another variant directly edits a learned linear subspace with LoReFT’s parameter count but reduced intervention computation.A further variant uses only half of LoReFT’s trainable parameters with less intervention computation.
  • E Ablating the parametrisation of LoReFT: Variants with a similar number of trainable parameters achieve similar performance to LoReFT across two models.The comparison is reported in table 17.

F.1 A single vector is worth a thousand tokens

Memorisation experiments test how much information a rank-1 representation intervention can encode in one vector or across many input-output mappings. The results show strong recovery for familiar text, weaker recovery for scrambled or random sequences, and reliable storage of up to 256 pairs.

  • F.1 A single vector is worth a thousand tokens: A single rank-1 LoReFT intervention at one residual-stream layer is trained to recover a specific output sequence from the last prompt token.The simplified intervention removes Wh, making it input-independent.
  • F.1 A single vector is worth a thousand tokens: 100% recovery holds for up to 2,048 tokens across most layers in both LLaMA-1 7B and 13B.The 13B model shows memorisation up to 2,560 tokens, while the 7B model fails catastrophically beyond 2,048 tokens.
  • F.1 A single vector is worth a thousand tokens: Recovery rates are significantly worse for scrambled Alice’s Adventures in Wonderland and random token sequences than for the original book.The authors suggest pretraining-data memorisation may contribute because the original book was likely in the pretraining corpus.
  • F.2 A single vector can memorise a codebook with 256 entries: A single rank-1 intervention can reliably memorise up to 256 random input-output pairs, with near-perfect recall at layer 20 of the 13B model.The experiments vary the intervention layer for LLaMA-1 7B and 13B.
  • F.2 A single vector can memorise a codebook with 256 entries: The codebook result suggests that a learned projection can partially disentangle token identities superpositioned in the original representation basis.The simplified intervention learns only a single input-dependent scalar b.

G.1 Multi-task learning: Learned ReFTs are like puzzle pieces

The section demonstrates that learned ReFT subspaces can be combined as task-specific components, and illustrates this with German completion, instruction following, and safety-oriented behavior. Separate subspaces are trained for different abilities and then used compositionally.

  • G.1 Multi-task learning: Learned ReFTs are like puzzle pieces: ReFT partitions a low-rank projection into orthogonal subspaces, assigns training examples to partitions, and allows the partitions to be combined at inference time.This frames learned abilities as compositional components rather than only as multi-task behavior.
  • G.1 Multi-task learning: Learned ReFTs are like puzzle pieces: A rank-8 LoReFT divides its subspace into two equal-rank groups trained separately for German completion and instruction following.The experiments use 1K German-completion examples and 1K Ultrafeedback instruction-following examples without hyperparameter tuning.
  • G.1 Multi-task learning: Learned ReFTs are like puzzle pieces: For air-pollution and Linux-process prompts, the German-completion subspace produces German responses while preserving task-specific content.The reported responses address each prompt in German.
  • 5-shot GOODY-2 Imitator with ReFTs: The GOODY-2 imitation experiment uses five sampled training examples and a rank-4 LoReFT with 32,772 learned parameters.Training takes 100 epochs and finishes within 30 seconds on one Nvidia A100 40G GPU.
  • G.1 Multi-task learning: Learned ReFTs are like puzzle pieces: LoReFT responses to the dog-breed prompt frame the issue as subjective preference rather than selecting one universally cuter breed.The response contrasts with Llama-2-chat’s factual correction and GOODY-2’s safety-oriented refusal rationale.

H Inference overhead analysis of ReFT with our ReFT library

ReFT introduces inference-time compute overhead because interventions remain hooked into the language model’s computation graph. The overhead increases with intervention rank and the number of intervening layers, while prompt-token interventions can reduce the additional cost.

  • Overhead source: ReFT interventions remain in the computation graph at inference, unlike mergeable LoRA weights, so they introduce runtime overhead.The ReFT library uses pyvene for model interventions.
  • Efficiency trade-off: Intervening only on prompt tokens can significantly reduce overhead because extra computation is incurred while populating the initial key-value cache.This optimization limits additional inference work relative to intervening throughout generation.
  • Experimental setup: The runtime benchmark measures greedy generation of up to 256 new tokens on LLaMA-1 7B, comparing LoReFT with an un-intervened language model.LoReFT settings vary rank, intervening layers, and intervening prompt positions.
  • Overhead trends: Higher intervention rank and more intervening layers positively correlate with larger generation-time overhead.The figure measures generation time for a fixed 256-token output across different prompt lengths.
  • Representative result: 0.05 second is the reported overhead for intervening across 10 layers at rank 8 on the last prompt token.This setting illustrates the overhead measured for a specific LoReFT configuration.
Loading 2404.03592v3…