Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

LLMdoctor: Token-Level Flow-Guided Preference Optimization for Efficient Test-Time Alignment of Large Language Models

Tiesunlong Shen, Rui Mao, Jin Wang, Heming Sun, Jian Zhang, Xuejie Zhang, Erik Cambria

TL;DR

The paper introduces LLMdoctor, a test-time alignment framework that learns token-level preferences from a frozen patient LLM via token-level reward acquisition and trains a smaller doctor model with token-level flow-guided preference optimization (TFPO). The doctor provides a flow-based reward model during online decoding, and decoding combines base model probabilities with token-level rewards to preserve diversity while aligning with human preferences. The TFPO objective enforces flow conservation over subtrajectories and uses a value head to discriminate among next-token options, enabling multi-dimensional preference balancing without retraining the large model. Empirical results across multiple datasets show that LLMdoctor outperforms existing test-time methods and even matches or exceeds full fine-tuning approaches like DPO, while maintaining generation diversity and enabling real-time preference control.

Abstract

Aligning Large Language Models (LLMs) with human preferences is critical, yet traditional fine-tuning methods are computationally expensive and inflexible. While test-time alignment offers a promising alternative, existing approaches often rely on distorted trajectory-level signals or inefficient sampling, fundamentally capping performance and failing to preserve the generative diversity of the base model. This paper introduces LLMdoctor, a novel framework for efficient test-time alignment that operates via a patient-doctor paradigm. It integrates token-level reward acquisition with token-level flow-guided preference optimization (TFPO) to steer a large, frozen patient LLM with a smaller, specialized doctor model. Unlike conventional methods that rely on trajectory-level rewards, LLMdoctor first extracts fine-grained, token-level preference signals from the patient model's behavioral variations. These signals then guide the training of the doctor model via TFPO, which establishes flow consistency across all subtrajectories, enabling precise token-by-token alignment while inherently preserving generation diversity. Extensive experiments demonstrate that LLMdoctor significantly outperforms existing test-time alignment methods and even surpasses the performance of full fine-tuning approaches like DPO.

LLMdoctor: Token-Level Flow-Guided Preference Optimization for Efficient Test-Time Alignment of Large Language Models

TL;DR

The paper introduces LLMdoctor, a test-time alignment framework that learns token-level preferences from a frozen patient LLM via token-level reward acquisition and trains a smaller doctor model with token-level flow-guided preference optimization (TFPO). The doctor provides a flow-based reward model during online decoding, and decoding combines base model probabilities with token-level rewards to preserve diversity while aligning with human preferences. The TFPO objective enforces flow conservation over subtrajectories and uses a value head to discriminate among next-token options, enabling multi-dimensional preference balancing without retraining the large model. Empirical results across multiple datasets show that LLMdoctor outperforms existing test-time methods and even matches or exceeds full fine-tuning approaches like DPO, while maintaining generation diversity and enabling real-time preference control.

Abstract

Aligning Large Language Models (LLMs) with human preferences is critical, yet traditional fine-tuning methods are computationally expensive and inflexible. While test-time alignment offers a promising alternative, existing approaches often rely on distorted trajectory-level signals or inefficient sampling, fundamentally capping performance and failing to preserve the generative diversity of the base model. This paper introduces LLMdoctor, a novel framework for efficient test-time alignment that operates via a patient-doctor paradigm. It integrates token-level reward acquisition with token-level flow-guided preference optimization (TFPO) to steer a large, frozen patient LLM with a smaller, specialized doctor model. Unlike conventional methods that rely on trajectory-level rewards, LLMdoctor first extracts fine-grained, token-level preference signals from the patient model's behavioral variations. These signals then guide the training of the doctor model via TFPO, which establishes flow consistency across all subtrajectories, enabling precise token-by-token alignment while inherently preserving generation diversity. Extensive experiments demonstrate that LLMdoctor significantly outperforms existing test-time alignment methods and even surpasses the performance of full fine-tuning approaches like DPO.
Paper Structure (56 sections, 5 theorems, 32 equations, 9 figures, 4 tables)

This paper contains 56 sections, 5 theorems, 32 equations, 9 figures, 4 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Preliminaries
  4. Methodology
  5. Token-Level Reward Acquisition
  6. TFPO-Based Fine-Grained Preference Tuning
  7. Flow-Guided Optimization for Token Sequences.
  8. Subtrajectory Balance Objective for TFPO.
  9. Value Discrimination Loss.
  10. Overall TFPO Training Objective.
  11. Training Procedure.
  12. Online Alignment
  13. Flow-Guided Reward Model Formulation.
  14. Reward-Guided Decoding Algorithm.
  15. Flexible Online Alignment.
  16. ...and 41 more sections

Key Result

Lemma A.1

For any fixed prefix $s_t$, the unique policy $\pi^\ast(\cdot \mid s_t)$ that maximizes the objective in Eq. eq:token-objective is given by: where $Z(s_t) = \sum_{y' \in \mathcal{V}} \pi_0(y' \mid s_t) \exp(r(s_t, y')/\tau)$ is the partition function.

Figures (9)

  • Figure 1: Comparison of test-time alignment approaches.
  • Figure 2: Overall framework of LLMdoctor
  • Figure 3: Pareto frontier comparison for helpfulness and harmlessness.
  • Figure 4: Weak-to-strong guidance performance. Comparison of length-controlled (LC) and raw AlpacaEval 2 win rates across different base model scales. All test-time methods employ a 7B guidance model, while DPO involves full fine-tuning at each respective scale.
  • Figure 5: Alignment signal dynamics.
  • ...and 4 more figures

Theorems & Definitions (10)

  • Lemma A.1: Optimal Policy Form
  • proof
  • Theorem A.2: Ceiling Effect
  • proof
  • Theorem B.1: Discriminative Importance as KL-Divergence Contribution
  • proof
  • Theorem C.1: Distribution Matching Property of TFPO
  • proof
  • Theorem C.2: Entropy Lower Bound
  • proof