- 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
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: 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 (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(s⊕xp,a)=r0(s,a)+αrp(s⊕xp,a)
where r0 denotes the question-related reward, rp the preference-related reward (as determined by preference context xp), and α the (latent, model-dependent) trade-off coefficient. At test time, REAR substitutes α with a user- or task-controlled coefficient α^, yielding a realignment reward:
rREAR(s⊕xp,a)=r0(s,a)+α^rp(s⊕xp,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 τ is:
SREAR(τ)=t=0∑Tγt[(1−λ)logπ(at∣st)+λlogπ(at∣st⊕xp)]
where r00 controls the preference intensification.
Integration with Test-time Scaling Algorithms
REAR is compatible with standard TTS paradigms:
- Best-of-N (BoN) Sampling: Multiple samples are generated; the candidate with maximal r01 is selected.
- 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 mathematical reasoning, REAR extends the reward decomposition to task-specific instructions, consistently outperforming majority voting in answer selection accuracy, even when r02 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: Scaling curves for answer selection accuracy with REAR versus majority vote on math problems, showing consistent improvements as r03 increases.
REAR exhibits strong cross-domain generalizability:
Ablations and Analysis
The r04 parameter interpolates between response helpfulness (question fidelity) and strict preference adherence. Empirical sweeps demonstrate a non-monotonic relationship; moderate r05 values (e.g., 20) provide optimal trade-offs for most settings, and results are robust to hyperparameter selection.
Figure 5: Best-of-N performance varies with r06, 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: Scaling performance on PrefEval as r07 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 (r08), 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 r09 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).