LoFiT: Localized Fine-tuning on LLM Representations

TL;DR

LoFiT introduces a localized fine-tuning framework that identifies key attention heads in LLMs and trains offset vectors for task adaptation, achieving comparable results to other fine-tuning methods with significantly fewer parameter modifications.

Contribution

This work presents LoFiT, a novel approach that localizes fine-tuning to a sparse set of attention heads, improving efficiency and effectiveness over existing representation intervention techniques.

Findings

LoFiT localizes to 3%-10% of attention heads.

LoFiT outperforms representation intervention methods in truthfulness and reasoning tasks.

LoFiT matches the performance of LoRA while modifying 20x-200x fewer parameters.

Abstract

Recent work in interpretability shows that large language models (LLMs) can be adapted for new tasks in a learning-free way: it is possible to intervene on LLM representations to elicit desired behaviors for alignment. For instance, adding certain bias vectors to the outputs of certain attention heads is reported to boost the truthfulness of models. In this work, we show that localized fine-tuning serves as an effective alternative to such representation intervention methods. We introduce a framework called Localized Fine-Tuning on LLM Representations (LoFiT), which identifies a subset of attention heads that are most important for learning a specific task, then trains offset vectors to add to the model's hidden representations at those selected heads. LoFiT localizes to a sparse set of heads (3%-10%) and learns the offset vectors from limited training data, comparable to the settings used for representation intervention. For truthfulness and reasoning tasks, we find that LoFiT's intervention vectors are more effective for LLM adaptation than vectors from representation intervention methods such as Inference-time Intervention. We also find that the localization step is important: selecting a task-specific set of attention heads can lead to higher performance than intervening on heads selected for a different task. Finally, across 7 tasks we study, LoFiT achieves comparable performance to other parameter-efficient fine-tuning methods such as LoRA, despite modifying 20x-200x fewer parameters than these methods.