Enhancing Chain-of-Thought Reasoning with Critical Representation Fine-tuning

TL;DR

This paper introduces Critical Representation Fine-Tuning (CRFT), a novel approach that identifies and optimizes key representations in language models to significantly improve complex reasoning tasks efficiently.

Contribution

CRFT is a new method that dynamically fine-tunes critical representations in models, enhancing reasoning performance without altering the entire model.

Findings

CRFT improves reasoning accuracy across eight benchmarks.

CRFT boosts one-shot accuracy by 16.4%.

CRFT operates efficiently within a low-rank subspace.

Abstract

Representation Fine-tuning (ReFT), a recently proposed Parameter-Efficient Fine-Tuning (PEFT) method, has attracted widespread attention for significantly improving parameter efficiency by editing representation space alone. In this work, we investigate applying ReFT to complex reasoning tasks. However, directly using the native ReFT method, which modifies fixed representations at the beginning and end of each layer, yields suboptimal performance, as these fixed-position representations have uncertain impact on the outputs. We observe that, in complex reasoning tasks, there often exist certain critical representations. These representations either integrate significant information from preceding layers or regulate subsequent layer representations. Through layer-by-layer propagation, they exert a substantial influence on the final output. Naturally, fine-tuning these critical representations has the potential to greatly enhance reasoning performance. Building upon these insights, we propose Critical Representation Fine-Tuning (CRFT), a novel method that identifies and optimizes these critical representations through information flow analysis. CRFT operates within a supervised learning framework, dynamically optimizing critical representations in a low-rank linear subspace while freezing the base model. The effectiveness and efficiency of our method are validated across eight benchmarks for arithmetic and commonsense reasoning, using LLaMA and Mistral model families. Furthermore, our method also adapts effectively to few-shot settings, boosting one-shot accuracy by 16.4%. Our work highlights the untapped potential of representation-level optimization for CoT reasoning, offering a lightweight yet powerful alternative to traditional PEFT methods.