Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Distribution-Aligned Decoding for Efficient LLM Task Adaptation

Senkang Hu, Xudong Han, Jinqi Jiang, Yihang Tao, Zihan Fang, Yong Dai, Sam Tak Wu Kwong, Yuguang Fang

TL;DR

The paper tackles the cost of adapting large language models (LLMs) to downstream tasks by reframing adaptation as output distribution alignment rather than weight updates. It introduces Steering Vector Decoding (SVDecode), which builds a task-specific steering vector from the KL divergence gradient between a briefly warm-started model and the pre-trained model, projects it to logit space, and applies it during decoding with an optimally chosen strength $\bar{\mu}$. The authors prove that a first-order SVDecode step is equivalent to a gradient step in full fine-tuning and derive a globally optimal steering strength using a Gauss-Newton approximation, enabling a decoding-time, training-free adaptation. Empirically, SVDecode consistently improves PEFT baselines across multiple tasks and models (up to 5 percentage points in MC accuracy and ~2 points in truthfulness) with minimal computational overhead, highlighting its potential for rapid, scalable deployment. Overall, the approach offers a principled, lightweight path to stronger task adaptation by shifting the emphasis from weight updates to distributional control during generation.

Abstract

Adapting billion-parameter language models to a downstream task is still costly, even with parameter-efficient fine-tuning (PEFT). We re-cast task adaptation as output-distribution alignment: the objective is to steer the output distribution toward the task distribution directly during decoding rather than indirectly through weight updates. Building on this view, we introduce Steering Vector Decoding (SVDecode), a lightweight, PEFT-compatible, and theoretically grounded method. We start with a short warm-start fine-tune and extract a task-aware steering vector from the Kullback-Leibler (KL) divergence gradient between the output distribution of the warm-started and pre-trained models. This steering vector is then used to guide the decoding process to steer the model's output distribution towards the task distribution. We theoretically prove that SVDecode is first-order equivalent to the gradient step of full fine-tuning and derive a globally optimal solution for the strength of the steering vector. Across three tasks and nine benchmarks, SVDecode paired with four standard PEFT methods improves multiple-choice accuracy by up to 5 percentage points and open-ended truthfulness by 2 percentage points, with similar gains (1-2 percentage points) on commonsense datasets without adding trainable parameters beyond the PEFT adapter. SVDecode thus offers a lightweight, theoretically grounded path to stronger task adaptation for large language models.

Distribution-Aligned Decoding for Efficient LLM Task Adaptation

TL;DR

The paper tackles the cost of adapting large language models (LLMs) to downstream tasks by reframing adaptation as output distribution alignment rather than weight updates. It introduces Steering Vector Decoding (SVDecode), which builds a task-specific steering vector from the KL divergence gradient between a briefly warm-started model and the pre-trained model, projects it to logit space, and applies it during decoding with an optimally chosen strength . The authors prove that a first-order SVDecode step is equivalent to a gradient step in full fine-tuning and derive a globally optimal steering strength using a Gauss-Newton approximation, enabling a decoding-time, training-free adaptation. Empirically, SVDecode consistently improves PEFT baselines across multiple tasks and models (up to 5 percentage points in MC accuracy and ~2 points in truthfulness) with minimal computational overhead, highlighting its potential for rapid, scalable deployment. Overall, the approach offers a principled, lightweight path to stronger task adaptation by shifting the emphasis from weight updates to distributional control during generation.

Abstract

Adapting billion-parameter language models to a downstream task is still costly, even with parameter-efficient fine-tuning (PEFT). We re-cast task adaptation as output-distribution alignment: the objective is to steer the output distribution toward the task distribution directly during decoding rather than indirectly through weight updates. Building on this view, we introduce Steering Vector Decoding (SVDecode), a lightweight, PEFT-compatible, and theoretically grounded method. We start with a short warm-start fine-tune and extract a task-aware steering vector from the Kullback-Leibler (KL) divergence gradient between the output distribution of the warm-started and pre-trained models. This steering vector is then used to guide the decoding process to steer the model's output distribution towards the task distribution. We theoretically prove that SVDecode is first-order equivalent to the gradient step of full fine-tuning and derive a globally optimal solution for the strength of the steering vector. Across three tasks and nine benchmarks, SVDecode paired with four standard PEFT methods improves multiple-choice accuracy by up to 5 percentage points and open-ended truthfulness by 2 percentage points, with similar gains (1-2 percentage points) on commonsense datasets without adding trainable parameters beyond the PEFT adapter. SVDecode thus offers a lightweight, theoretically grounded path to stronger task adaptation for large language models.

Paper Structure

This paper contains 50 sections, 2 theorems, 47 equations, 4 figures, 12 tables, 2 algorithms.

Table of Contents

  1. Introdcution
  2. Rethinking LLM Task Adaptation from the Perspective of Output Distribution Alignment
  3. Method
  4. Steering Vector Construction
  5. Confidence-Aware Steering Vector Constraint.
  6. Task-Aware Steering Vector Decoding
  7. Experiments
  8. Experimental Setup
  9. Main Results
  10. Ablation Study
  11. Conclusion
  12. Acknowledgement
  13. Related Work
  14. Parameter-Efficient Fine-Tuning
  15. LLM Task Adaptation
  16. ...and 35 more sections

Key Result

Theorem 1

The NLL objective in Eq. eq:nll is equivalent to minimizing the expected Kullback-Leibler (KL) divergence between the empirical label distribution $\hat{P}_{\mathrm{task}}(y \mid x)$ and the model's output distribution: where $\hat{P}_{\mathrm{task}}(y \mid x)$ is typically a delta function centered on the ground-truth label.

Figures (4)

  • Figure 1: Illustration of the framework of our proposed SVDecode. It includes two steps: (a) steering vector construction and (b) task-aware steering vector decoding. After the decoding with the steering vector, we can see the warm-started model's output distribution is steered towards the task-specific target distribution, thereby enhancing the performance of the model on the downstream task.
  • Figure 2: Ablation study on logit-space projection. 'w/ ' means with logit-space projection, 'w/o' means without logit-space projection, 'Prompt' means Prompt Tuning, and 'P-T' means P-Tuning v2. We conduct the ablation study on multiple-choice tasks.
  • Figure 3: Analysis of warm-start steps. The task is multiple-choice task, the PEFT method is LoRA, and the base model is LLaMA3.1-8B.
  • Figure 4: Ablation study on the $\alpha$ parameter in the confidence-aware constraint.

Theorems & Definitions (4)

  • Theorem 1
  • proof
  • Theorem 2
  • proof