💬 NLP

REAR: Test-time Preference Realignment through Reward Decomposition

This paper introduces REAR, a training-free test-time framework that aligns large language models with diverse user preferences by decomposing reward functions to selectively rescale reward components, thereby extending efficient test-time scaling beyond verifiable domains to complex preference alignment tasks.

Original authors: Fuxiang Zhang, Pengcheng Wang, Chenran Li, Yi-Chen Li, Yuxin Chen, Lang Feng, Chenfeng Xu, Masayoshi Tomizuka, Bo An

Published 2026-06-30

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Fuxiang Zhang, Pengcheng Wang, Chenran Li, Yi-Chen Li, Yuxin Chen, Lang Feng, Chenfeng Xu, Masayoshi Tomizuka, Bo An

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a very smart, well-trained robot assistant (a Large Language Model). This robot has learned to be helpful, polite, and accurate based on the massive amount of data it was trained on. However, sometimes you have a very specific, personal request. Maybe you want the robot to explain math in a way that is "fun but not like a video game," or you want it to act like a "grumpy old detective."

The problem is that the robot's default setting might not perfectly match your specific mood or preference. Usually, to fix this, you have to send the robot back to "school" (retraining) with new data, which is expensive, slow, and requires a lot of resources.

This paper introduces a clever trick called REAR (Realignment through Reward Decomposition) that lets you fix the robot's behavior right now, while it's talking to you, without sending it back to school.

Here is how it works, broken down into simple analogies:

1. The Problem: The Robot's "Default" vs. Your "Wish"

Think of the robot's brain as having two different voices speaking at the same time:

Voice A (The Question): "How do I answer this math problem correctly?"
Voice B (The Preference): "How do I answer this while being grumpy and avoiding video game metaphors?"

When the robot is trained, it learns a mix of these two voices. But when you ask a specific question, that mix might be wrong. Maybe it's too focused on being "correct" and ignores your request to be "grumpy."

2. The Solution: The "Volume Knob" (Reward Decomposition)

The authors realized they could mathematically separate these two voices. They treat the robot's internal "reward" (the feeling of doing a good job) as two separate ingredients:

The Question Ingredient: How well does it answer the prompt?
The Preference Ingredient: How well does it follow your specific style?

REAR is like a volume knob that you can turn during the conversation.

If you turn the knob up, the robot listens more to your "Preference" voice.
If you turn it down, it listens more to the "Question" voice.

The magic is that they figured out how to calculate this volume knob using only the robot's own internal thoughts (its probability scores). You don't need a new teacher or a new dataset; you just use the robot's existing brain to re-balance itself.

3. The Strategy: "Try, Try, Try Again" (Test-Time Scaling)

Since the robot can't instantly know which answer is perfect, REAR uses a strategy called Test-Time Scaling. Think of it like this:

The "Best of N" Approach: Imagine you ask the robot to write a story. Instead of accepting the first draft, you ask it to write 16 different versions in the time it usually takes to write one.
The Scorecard: For each of those 16 versions, REAR acts as a judge. It uses that "volume knob" math to score them. It asks: "Which of these 16 stories best answers the question AND follows the grumpy detective style?"
The Winner: The robot picks the highest-scoring version and gives that to you.

They also use a more advanced version called Tree Search (like a detective exploring different paths in a maze). Instead of just writing 16 full stories, the robot explores different sentence choices step-by-step, constantly checking the score to make sure it stays on the right path.

4. Why This is a Big Deal

No Training Needed: You don't need to retrain the robot. It's like tuning a radio station while listening, rather than buying a new radio.
Works on Anything: They showed it works not just for "grumpy detectives," but also for complex math problems and even visual tasks (like looking at a picture and describing it accurately without making things up).
Efficient: Because it uses the robot's own internal math, it doesn't need to load a separate, heavy "judge" program to check the answers. It's faster and cheaper than previous methods.

Summary

REAR is a tool that lets you instantly customize a smart AI's personality or focus while it is working. It does this by mathematically separating "answering the question" from "following your style," and then using a "try many options and pick the best one" strategy to find the perfect response. It's like having a remote control for the AI's personality that works instantly, without needing to rebuild the AI from scratch.

Technical Summary: REAR: Test-time Preference Realignment through Reward Decomposition

Problem Statement

Aligning Large Language Models (LLMs) with diverse and personalized user preferences is a critical challenge. While post-training methods like Reinforcement Learning from Human Feedback (RLHF) are effective, they require costly data curation and additional training, making them inflexible for dynamic or specific user needs. Test-Time Scaling (TTS) offers a training-free alternative by leveraging additional computation during inference to improve output quality. However, existing TTS methods have largely been limited to verifiable domains (e.g., mathematics, coding) where correctness is easily judged. Applying TTS to preference alignment is difficult because response quality is holistic and not reducible to a simple verifiable answer. Current test-time alignment approaches often impose fixed computational overheads or require training value heads, limiting their scalability and efficiency.

Methodology

The paper introduces REAR (REAlignment Reward), a novel framework that frames test-time preference alignment as a realignment problem. The core insight is that a base model's training objective may not optimally satisfy a specific user's preference at inference time. REAR addresses this by decomposing the underlying reward function into two components and dynamically rescaling their importance.

1. Reward Decomposition

The authors formalize the reward structure of a pretrained LLM within the Maximum Entropy RL framework. They posit that the reward $r(s \oplus x_p, a)$ for a state $s$ (question) and preference $x_p$ can be decomposed into:

Question-related reward ( $r_0$ ): The implicit reward the base model maximizes to answer the question correctly.
Preference-related reward ( $r_p$ ): The reward component capturing how well the response satisfies the preference $x_p$ .

The composite reward is modeled as $r(s \oplus x_p, a) = r_0(s, a) + \alpha r_p(s \oplus x_p, a)$ , where $\alpha$ is an inherent balancing coefficient.

2. Realignment Reward (REAR)

To realign the model for a specific user, REAR introduces a controllable coefficient $\hat{\alpha}$ (or $\lambda = \hat{\alpha}/\alpha$ ) to rescale the preference component during inference:
$r_{REAR}(s \oplus x_p, a) = r_0(s, a) + \hat{\alpha} r_p(s \oplus x_p, a)$

Crucially, the authors derive a computationally efficient formulation for this reward. By leveraging the relationship between optimal policies and value functions in MaxEnt RL, they show that $r_{REAR}$ can be expressed as a linear combination of token-level policy log-probabilities:
$S_{REAR}(\tau) = \sum_{t=0}^{T} \gamma^t \left( (1-\lambda) \log \pi(a_t|s_t) + \lambda \log \pi(a_t|s_t \oplus x_p) \right)$
This score function depends only on the log-probabilities of the base model under two contexts: the question-only prompt and the full prompt including preferences. It does not require training a separate reward model or generating additional tokens for scoring.

3. Integration with Test-Time Scaling

The REAR score is integrated into two TTS algorithms:

Best-of-N (BoN) Sampling: Generate $N$ responses and select the one with the highest $S_{REAR}$ score.
Diverse Verifier Tree Search (DVTS): A tree search algorithm that uses REAR scores to guide branch expansion and selection, promoting diversity while optimizing for preference alignment.

Key Contributions

Formalization of Realignment: The paper formalizes test-time preference alignment as a realignment problem, proposing REAR as a computationally efficient reward derived from a decomposed objective.
Training-Free Scoring: REAR is derived solely from the base model's internal probabilities, requiring no additional training, reward model training, or speculative generation steps.
Scalable Algorithms: The authors develop two scalable test-time algorithms (BoN and DVTS) guided by REAR that can enhance generation performance with varying sample budgets.
Generalizability: The method is shown to generalize beyond explicit preference alignment to mathematical reasoning and visual tasks by injecting task-specific preference text.

Experimental Results

The authors evaluated REAR on several benchmarks using Qwen and Llama models:

Preference Alignment Benchmarks: On PrefEval (covering explicit/implicit preferences and choices), Multifaceted Bench, and Ping-Pong (role-playing), REAR-guided methods (both BoN and DVTS) consistently outperformed baselines including greedy decoding, Amulet, Linear Alignment (LA), and Best-of-N with Generative Reward Models (GenRM). For instance, DVTS with REAR improved over greedy decoding by 11.6% on average across metrics.
Long-Context Robustness: REAR demonstrated superior robustness in long-context scenarios (up to 16k tokens) compared to GenRM, which degraded significantly as context length increased.
Mathematical Reasoning: When applied to mathematical tasks (MATH500, AIME, AMC) with appropriate preference text (e.g., "Act as a rigorous mathematician"), REAR-guided selection outperformed majority voting, showing consistent scaling gains.
Visual Tasks: On MMHal-Bench for visual hallucination detection, REAR improved overall scores and reduced hallucination rates compared to greedy decoding and GenRM.
Cross-Model Transfer: The method generalized to Llama-3.1-8B-Instruct without per-model tuning, confirming that the reward decomposition is not specific to the Qwen family.
Efficiency: REAR-guided methods were significantly faster than baselines requiring external reward models or iterative optimization (like Amulet), as they avoid loading separate models and rely on a single forward pass for scoring.

Significance and Claims

The paper claims that REAR provides a scalable and general solution for aligning LLMs with user preferences at test time without the overhead of training or external reward models. By decomposing the reward into question and preference components, REAR enables dynamic control over the trade-off between answering the question and adhering to preferences.

The authors emphasize that REAR allows test-time scaling to extend to subjective, open-ended domains (like preference alignment and role-playing) where traditional verifiers fail. They note that while a default $\lambda = 20$ works robustly across many tasks, the method allows for tuning the strength of realignment. The work positions itself as a step toward making test-time alignment more accessible and efficient, bridging the gap between the success of TTS in verifiable domains and the complexities of subjective preference alignment.

Limitations Acknowledged: The authors note that REAR exposes a tunable hyperparameter ( $\lambda$ ) to the user, which may require small validation sweeps for tasks with significantly different preference distributions. Additionally, while TTS is lighter than fine-tuning, identifying the optimal computational budget remains a direction for future work.