Joint Localization and Activation Editing for Low-Resource Fine-Tuning

TL;DR

JoLA is a novel method that jointly learns which model components to edit and how to edit them, significantly improving low-resource fine-tuning of large language models across various tasks.

Contribution

The paper introduces JoLA, a joint learning approach for localization and activation editing, enhancing stability and performance in low-resource scenarios.

Findings

JoLA outperforms existing methods on three benchmarks.

It effectively identifies relevant model modules for editing.

JoLA improves task performance in low-data settings.

Abstract

Parameter-efficient fine-tuning (PEFT) methods, such as LoRA, are commonly used to adapt LLMs. However, the effectiveness of standard PEFT methods is limited in low-resource scenarios with only a few hundred examples. Recent advances in interpretability research have inspired the emergence of activation editing (or steering) techniques, which modify the activations of specific model components. Due to their extremely small parameter counts, these methods show promise for small datasets. However, their performance is highly dependent on identifying the correct modules to edit and often lacks stability across different datasets. In this paper, we propose Joint Localization and Activation Editing (JoLA), a method that jointly learns (1) which heads in the Transformer to edit (2) whether the intervention should be additive, multiplicative, or both and (3) the intervention parameters themselves - the vectors applied as additive offsets or multiplicative scalings to the head output. Through evaluations on three benchmarks spanning commonsense reasoning, natural language understanding, and natural language generation, we demonstrate that JoLA consistently outperforms existing methods. The code for the method is released at https://github.com/wenlai-lavine/jola.