Papers
Topics
Authors
Recent
Search
2000 character limit reached

REAR: Test-time Preference Realignment through Reward Decomposition

Published 29 Jun 2026 in cs.CL and cs.LG | (2606.30339v1)

Abstract: Aligning LLMs with diverse user preferences is a critical yet challenging task. While post-training methods can adapt models to specific needs, they often require costly data curation and additional training. Test-time scaling (TTS) presents an efficient, training-free alternative, but its application has been largely limited to verifiable domains like mathematics and coding, where response correctness is easily judged. To extend TTS to preference alignment, we introduce a novel framework that models the task as a realignment problem, since the base model often fails to sufficiently align with the stated preference. Our key insight is to decompose the underlying reward function into two components: one related to the question and the other to preference information. This allows us to derive a REAlignment Reward (REAR) that selectively rescales the proportions of these two reward terms. We then show that REAR can be formulated as a linear combination of token-level policy log-probabilities, making it computationally efficient and easy to integrate with various TTS algorithms such as best-of-$N$ sampling and tree search. Experiments show that compared to other test-time baselines, REAR not only enables scalable test-time realignment for preference alignment tasks under diverse user requirements, but also generalizes to mathematical and visual tasks under appropriate preference settings.

Summary

  • The paper introduces REAR, a reward decomposition framework that realigns LLM outputs at test time to balance query fidelity with user preferences.
  • It employs closed-form scoring in Best-of-N and DVTS sampling, achieving 5-11% improvements over baseline preference alignment benchmarks.
  • The methodology enhances model efficiency and robustness across diverse domains, including long-context and multimodal tasks.

Test-Time Preference Realignment for LLMs via Reward Decomposition

Problem Formulation and Motivation

Aligning LLM generations with user-specific preferences at inference time remains a central but under-addressed challenge. Traditional post-training alignment methods, such as RLHF and instruction tuning, require significant overhead in data collection and fine-tuning, making them unsuitable for highly dynamic or diverse user preference settings. Existing test-time scaling (TTS) methods have shown strong results in verifiable domains (math, code) where correctness is easily assessed but have not generalized to open-ended, preference-sensitive dialog tasks. The core difficulty lies in defining a tractable and faithful reward function at inference that reliably balances question fidelity and user preference satisfaction.

The REAR framework identifies the fundamental issue as a realignment problem: pretrained models implicitly encode a fixed trade-off between answering the user's query and adhering to stated preferences, which may not suffice for downstream tasks. The proposed solution is to decompose the latent reward into question-related and preference-related components and introduce an inference-time mechanism to rescale their trade-off. Figure 1

Figure 1: REAR combines question- and preference-related rewards, enabling realignment from a non-collaborative suggestion (conflicting with the user's stated aversion to gamification) to a collaborative response with a higher REAR score.

REAR: Theoretical Foundation and Formulation

REAR (REAlignment Reward) is grounded in a maximum entropy RL framework for token-level MDPs in text generation. The core insight is the following reward decomposition:

r(sxp,a)=r0(s,a)+αrp(sxp,a)r(s \oplus x_p, a) = r_0(s, a) + \alpha r_p(s \oplus x_p, a)

where r0r_0 denotes the question-related reward, rpr_p the preference-related reward (as determined by preference context xpx_p), and α\alpha the (latent, model-dependent) trade-off coefficient. At test time, REAR substitutes α\alpha with a user- or task-controlled coefficient α^\hat{\alpha}, yielding a realignment reward:

rREAR(sxp,a)=r0(s,a)+α^rp(sxp,a)r_{\mathrm{REAR}}(s \oplus x_p, a) = r_0(s, a) + \hat{\alpha} r_p(s \oplus x_p, a)

Crucially, this reward can be computed via a closed-form, potential-based shaping of LLM log-probabilities under two contexts (question only, question + preference), avoiding the need for explicit reward model training. The derived REAR score function for a response τ\tau is:

SREAR(τ)=t=0Tγt[(1λ)logπ(atst)+λlogπ(atstxp)]S_{\mathrm{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]

where r0r_00 controls the preference intensification.

Integration with Test-time Scaling Algorithms

REAR is compatible with standard TTS paradigms:

  1. Best-of-N (BoN) Sampling: Multiple samples are generated; the candidate with maximal r0r_01 is selected.
  2. Diverse Verifier Tree Search (DVTS): A tree search expands candidate sequences, scoring partial or full trajectories with REAR to promote exploration and preference alignment.

This architectural separation (reward scoring versus sequence proposal) enables direct comparison with reward-based sampling and search used in verifiable TTS.

Empirical Results and Numerical Analysis

Empirical evaluation utilizes a suite of preference alignment and reasoning benchmarks, with strong coverage of both explicit and implicit preference-following, conversation style, and role-play settings. Models include Qwen and Llama3 instruction-tuned variants; baselines comprise test-time alignment (Amulet, LA), BoN with generative/external RMs, and greedy decoding.

Key performance highlights:

  • On the PrefEval explicit and implicit benchmarks, BoN with REAR and DVTS with REAR provide absolute improvements over baselines ranging from 5% to 11% in LLM-as-a-judge scores.
  • In long-context scenarios (up to 16k tokens), REAR-guided methods outperform all baselines, demonstrating robustness to preference retention under increased context length.
  • Across summary metrics, DVTS with REAR achieves the highest aggregate preference alignment across all tested benchmarks. Figure 2

    Figure 2: REAR-guided methods exceed external reward model baselines in both BoN and DVTS regimes, demonstrating competitive or superior reward estimation without additional model overhead.

On mathematical reasoning, REAR extends the reward decomposition to task-specific instructions, consistently outperforming majority voting in answer selection accuracy, even when r0r_02 increases to 64 samples. Across MATH500, AIME24/25, and AMC23, gains are robust with respect to the exact preference phrasing and are not observed for irrelevant preference text. Figure 3

Figure 3: Scaling curves for answer selection accuracy with REAR versus majority vote on math problems, showing consistent improvements as r0r_03 increases.

REAR exhibits strong cross-domain generalizability:

  • On MMHal-Bench (visual hallucination detection), preference-injected REAR evaluation reduces hallucination rates while increasing factuality scores.
  • Cross-model transfer to Llama-3.1-8B-Instruct validates that performance gains are architecture-agnostic, provided the base model is sufficiently instruction-following-aligned. Figure 4

    Figure 4: Example multimodal hallucination task with preference-conditioned input, enabling REAR to guide hallucination suppression during generation.

Ablations and Analysis

The r0r_04 parameter interpolates between response helpfulness (question fidelity) and strict preference adherence. Empirical sweeps demonstrate a non-monotonic relationship; moderate r0r_05 values (e.g., 20) provide optimal trade-offs for most settings, and results are robust to hyperparameter selection. Figure 5

Figure 5: Best-of-N performance varies with r0r_06, with peaks at moderate preference amplification and degradation for overly aggressive scaling.

Efficiency evaluations highlight BoN with REAR as the fastest TTS preference-alignment method among those considered, due to its lack of reliance on standalone reward models or iterative backward passes. The overhead versus greedy decoding is minor compared to the gains in alignment, and the approach enables simple deployment for user-facing customization. Figure 6

Figure 6: Scaling performance on PrefEval as r0r_07 increases for BoN, DVTS, and baseline methods by task, validating the superiority and scalability of REAR-guided inference.

Practical Implications and Future Directions

REAR presents a highly practical solution to test-time preference realignment: it augments the capabilities of pretrained LLMs to accommodate dynamic and multifaceted user specifications without retraining or auxiliary reward models. The modular reward decomposition exposes a high-level realignment knob (r0r_08), offering end users or downstream applications control over preference strength—a critical requirement for customized deployment at scale.

From a theoretical perspective, REAR crystallizes the implicit nature of reward modeling within aligned LLMs and demonstrates that RL-derived models can, with suitable test-time transformations, double as process and preference verifiers for their own outputs.

Open questions remain around scaling REAR to highly complex or adversarial preference structures, formal guarantees under distribution shifts or adversarial contexts, and extending the approach to reinforcement agents and multimodal or agentic LLM settings. Further research could also refine methods for learning or inferring the optimal scaling coefficient r0r_09 based on user-level interaction data.

Conclusion

REAR introduces a principled, efficient, and general approach for LLM preference alignment at inference via reward decomposition and reweighting. By leveraging the internal log-probabilities of base models under different conditioning, it obviates the need for post-training, auxiliary reward models, or extensive annotations, while providing strong numerical gains and robustness across domains and architectures. This framework is an important advance for scalable, controllable, and user-aligned LLM deployment (2606.30339).

Paper to Video (Beta)

No one has generated a video about this paper yet.

Whiteboard

No one has generated a whiteboard explanation for this paper yet.

Open Problems

We haven't generated a list of open problems mentioned in this paper yet.

Collections

Sign up for free to add this paper to one or more collections.

Tweets

Sign up for free to view the 2 tweets with 13 likes about this paper.