PaperlayerView exact revision on arXiv

Source-linked AI summary

Test-Time Preference Optimization: On-the-Fly Alignment via Iterative Textual Feedback

Yafu Li, Xuyang Hu, Xiaoye Qu, Linjie Li, Yu Cheng

arXiv:2501.12895v1cs.CL

TL;DR

LLMs need faster adaptation to human preferences than retraining permits. TPO addresses this by converting reward signals into textual critiques and iterative output revisions during inference, improving benchmark alignment and sometimes surpassing aligned models after only a few steps.

  • Problem

    LLMs can produce unexpected or unsafe responses, while conventional preference alignment relies on iterative parameter retraining that limits swift adaptation.

  • Method

    TPO keeps model parameters fixed and iteratively translates reward-model signals into textual critiques and improvements that refine responses at test time.

  • Results

    Across instruction-following, preference, safety, and mathematics benchmarks, TPO improves alignment; after two steps, Llama-3.1-70B-SFT can match or exceed aligned models.

  • Takeaways & Limitations

    TPO provides a lightweight inference-time alternative that can bridge or exceed the performance gap between unaligned and aligned LLMs while scaling through search width and depth.

  • Takeaways & Limitations

    TPO depends on the policy model’s ability to interpret and execute textual instructions, and weaker models may fail to maintain alignment.

Abstract

from arXiv · show

Large language models (LLMs) demonstrate impressive performance but lack the flexibility to adapt to human preferences quickly without retraining. In this work, we introduce Test-time Preference Optimization (TPO), a framework that aligns LLM outputs with human preferences during inference, removing the need to update model parameters. Rather than relying on purely numerical rewards, TPO translates reward signals into textual critiques and uses them as textual rewards to iteratively refine its response. Evaluations on benchmarks covering instruction following, preference alignment, safety, and mathematics reveal that TPO progressively improves alignment with human preferences. Notably, after only a few TPO steps, the initially unaligned Llama-3.1-70B-SFT model can surpass the aligned counterpart, Llama-3.1-70B-Instruct. Furthermore, TPO scales efficiently with both the search width and depth during inference. Through case studies, we illustrate how TPO exploits the innate capacity of LLM to interpret and act upon reward signals. Our findings establish TPO as a practical, lightweight alternative for test-time preference optimization, achieving alignment on the fly. Our code is publicly available at https://github.com/yafuly/TPO.

1. Introduction

TPO aligns LLM outputs with human preferences during inference by converting reward feedback into textual critiques and iterative output revisions, without retraining. Across diverse benchmarks, it improves alignment and can allow an initially unaligned model to surpass an aligned counterpart after a few steps.

  • TPO addresses the need to adapt LLMs to human preferences without repeatedly retraining model parameters.
  • TPO interprets numerical reward signals as textual critiques and uses them to iteratively improve responses during inference.
  • TPO evaluates instruction following, preference alignment, safety, and mathematics on both unaligned and aligned models.
  • After only two TPO steps, both unaligned and aligned models achieve substantial benchmark gains.
  • TPO can let the initially unaligned Llama-3.1-70B-SFT surpass Llama-3.1-70B-Instruct on nearly all benchmarks.
  • TPO leverages LLMs’ ability to interpret and act on reward signals to align with human preferences on the fly.

2. Related Work

Prior preference-optimization methods generally update model parameters during training, while inference-time alignment methods modify context or decoding. TPO instead iteratively transforms reward feedback into textual guidance that refines the entire response.

  • Preference Optimization: Preference optimization methods include point-wise and pairwise approaches that typically use training-time gradient optimization.
  • Inference-Time Alignment: Inference-time alignment methods can optimize input context through in-context learning, retrieval augmentation, prompt rewriting, or in-context DPO.
  • Inference-Time Alignment: TPO differs by using reward-model feedback to iteratively interpret and refine the entire response rather than only finding an optimal context.
  • Inference-Time Alignment: TPO converts numerical feedback into textual feedback that continuously shapes model outputs.

3. Preliminary

Preference optimization aligns a policy with human preferences by increasing preferred outputs and decreasing dispreferred ones. Traditional methods optimize model parameters numerically using scoring functions, reward models, or gradient descent.

  • Preference optimization seeks to increase the likelihood of preferred outputs while decreasing the likelihood of misaligned ones.
  • A general scoring function quantifies policy alignment using prompts, preferred responses, and dispreferred responses.
  • RLHF incorporates human feedback into a reward model that guides policy optimization.
  • DPO replaces the reward model with a reparameterized reward expressed through the optimal policy.
  • Training-time preference methods update neural parameters with gradient descent to alter the output distribution toward preferred generations.

4. Method

TPO reframes preference optimization as inference-time textual optimization: it critiques candidate responses, derives textual update instructions, and iteratively refines outputs while keeping model parameters fixed. The procedure uses reward-model feedback to select responses, generate critiques, produce new candidates, and retain the highest-scoring result.

  • Core idea: TPO searches for an optimal context ϕ rather than updating model parameters θ, thereby changing the output distribution through textual variables.Its textual updates correspond to gradient-based optimization while remaining entirely in natural language.
  • Optimization loop: Each iteration performs loss calculation, gradient computation, and variable optimization through prompt-based interactions, with all calculations encoded in natural language.The model calls itself to critique responses, derive updates, and apply them to the next candidates.
  • Core idea: At each step, the reward model scores candidates, and the highest- and lowest-scoring responses become the chosen and rejected examples.The system stores response-score pairs in a cache and selects its best and worst entries for comparison.
  • Optimization loop: TPO converts the chosen-versus-rejected comparison into textual loss, textual gradients, and update instructions for generating refined responses.The textual loss identifies strengths and weaknesses, while the textual gradient supplies refinement instructions rather than numeric gradients.
  • Optimization loop: The procedure repeats for at most D iterations and returns the highest-scoring response in the cache as the final output.TPO combines parallel candidate sampling with sequential revision and differs from static feedback methods by revising responses interactively.

5. Experimental Setup

The experiments compare TPO across unaligned and aligned policy models, using reward-model feedback and benchmarks spanning instruction following, preference alignment, safety, and mathematics. Evaluation also examines test-time training curves and contrasts inference-time context updates with training-time preference optimization.

  • Models: The study compares unaligned Llama-3.1-70B-SFT with aligned Llama-3.1-70B-Instruct, Mistral-Small-Instruct-2409, and an on-policy Llama-3.1-70B-DPO baseline.The models differ in whether they underwent training-time preference optimization such as RLHF or DPO.
  • Models: TPO updates contextual prompts ϕ at test time, whereas DPO incorporates reward-model feedback into model parameters θ.This comparison isolates context updating from parameter updating as two modes of preference alignment.
  • Benchmarks: Evaluations cover instruction following, general preference alignment, safety, and mathematical ability using AlpacaEval 2, Arena-Hard, HH-RLHF, BeaverTails-Evaluation, XSTest, and MATH-500.The study samples 500 HH-RLHF test instances and uses the full test sets for the other listed benchmarks.
  • Evaluation procedure: TPO runs for up to five iterations, measuring average reward-model scores at each iteration against no-test-time-training baselines and a revision baseline.Figure 3 represents test-time training performance as reward-model score versus TPO iterations.
  • Evaluation procedure: Table 1 uses D=2 iterations and N=5 samples by default, with an ultra setting of D=5 and N=20 to show TPO’s potential.The table compares TPO with training-time aligned baselines and marks the reward models used for each TPO result.

6. Experimental Results

TPO progressively aligns both unaligned and aligned models during inference, with the largest improvement typically occurring in the first two steps. Across benchmarks, TPO improves performance, inference stability, and computational efficiency without parameter updates.

  • Test-time Training: After two steps, TPO raises unaligned Llama-3.1-70B-SFT to match or exceed aligned Llama-DPO and Llama-Instruct under two reward models.TPO achieves this while avoiding the extensive training-phase computation used to align the comparison models.
  • Test-time Training: The first optimization step yields the largest improvement, so benchmark evaluations use two steps to balance efficiency and performance.Further steps remain beneficial but are comparatively less impactful.
  • Test-time Training: TPO improves already aligned models, whereas the revision baseline provides limited additional benefit.The result holds for aligned and unaligned models, with TPO producing stronger iterative alignment gains.
  • Benchmark Performance: TPO produces consistent gains for unaligned Llama-3.1-70B-SFT, surpassing Llama-3.1-70B-DPO on all evaluation sets and Llama-3.1-70B-Instruct on nearly all metrics.With five iterations and 20 samples, the ultra-TPO setting surpasses Llama-3.1-70B-Instruct across all metrics.
  • Benchmark Performance: A 22B-parameter Mistral-Small-Instruct-2409 reaches an LC score of 53.4 on AlpacaEval 2 after TPO, comparable to GPT-4-Turbo.Applying TPO yields consistent gains across diverse tasks for already aligned models.
  • Benchmark Performance: TPO improves inference stability by lowering reward-score standard deviations, enabling the unaligned model to surpass aligned models’ stability.The analysis interprets this as redistributing probability mass toward higher-quality responses receiving higher rewards.
  • Benchmark Performance: TPO tailors outputs to reward-model feedback without changing parameters and at a negligible fraction of training-time preference-optimization cost.The method nevertheless confers notable improvements in performance metrics.

7. Analysis

The analysis examines TPO’s textual-feedback interaction, test-time scaling, computational efficiency, and dependence on instruction-following ability. TPO benefits from increased search width and depth, can outperform larger-sample Best-of-N sampling, and costs far less than training-time optimization, but weaker instruction-following can undermine alignment.

  • Policy-Reward Interaction with Textual Feedback: TPO uses iterative textual feedback from reward-model comparisons to refine responses during inference.The policy interprets chosen and rejected responses, producing textual critiques that guide subsequent revisions.
  • Test-Time Scaling in Width and Depth: Increasing search width from 5 to 20 consistently boosts performance before plateauing, while additional iterations let narrower searches catch up.TPO-D2-N15 reaches a reward-model score comparable to TPO-D4-N5 and TPO-D3-N10.
  • Test-Time Scaling in Width and Depth: TPO-D2-N5 samples 15 responses yet surpasses Best-of-N with 30 and 60 samples, achieving average win-rates of 65.2% and 57.5%, respectively.The comparison uses GPT-4 win-rate evaluation on 100 instances from AlpacaEval 2, Arena-Hard, and HH-RLHF.
  • Scaling Computing from Training-time to Test-time: Training Llama-3.1-70B-DPO requires approximately 72,840 PFLOPs, whereas TPO costs around 9.3 PFLOPs per query, less than 0.01% of training overhead.TPO uses textual updates at test time instead of finetuning the model on a large dataset.
  • Instruction Following as a Prerequisite: TPO relies on the policy model’s ability to interpret and execute textual instructions, and Llama-3.1-8B-Instruct loses alignment across iterations.Its decreasing reward-model scores indicate that instruction-following proficiency is important for TPO success.

8. Conclusion

The conclusion presents TPO as an inference-time alternative to retraining that converts reward signals into textual critiques and improvements. Across instruction following, safety, and mathematical reasoning benchmarks, a few optimization steps can bridge or exceed the gap between unaligned and aligned models.

  • Conclusion: TPO aligns LLMs with human preferences during inference without retraining by translating reward signals into textual losses and gradients.The method iteratively refines outputs and shifts probability mass toward preferred responses.
  • Conclusion: Experiments across instruction following, safety, and mathematical reasoning show that a few TPO steps can bridge or exceed the performance gap between unaligned and aligned LLMs.The conclusion also reports flexible scaling through search width and depth.
  • Conclusion: TPO is characterized as a lightweight, interpretable, and efficient alternative to training-time preference optimization.Its success depends on models being able to interpret and act on textual feedback.

A. Prompt Design

The prompt design defines textual loss prompts for comparing chosen and rejected responses or evaluating a single response. These prompts instruct the model to analyze strengths and weaknesses without generating a new answer to the query.

  • TPO Loss Prompt: The TPO loss prompt asks the model to compare chosen and rejected responses and explain why one is preferred.It is designed for textual loss calculation from pairwise reward-model outcomes.
  • TPO Loss Prompt: The loss prompt explicitly instructs the model not to generate a response to the user query.The instruction emphasizes concise analysis rather than answer generation.
  • Revision Baseline: The revision baseline prompt evaluates one model response by analyzing its strengths and weaknesses step by step.Unlike the TPO loss prompt, it does not compare chosen and rejected responses.

B. Data Statistics

The appendix identifies Table 3 as presenting statistical data for all benchmark datasets used in the study.

  • Data Statistics: Table 3 presents the statistical data for all benchmark datasets used in the work.The accompanying table is labeled as benchmark-dataset data statistics.

C. Test-time Training

Figures 8–10 present test-time training curves across the Llama-3.1-70B-SFT and aligned models using different reward models.

  • C. Test-time Training: Test-time training curves are reported for different datasets and models in Figures 8, 9, and 10.Figures 8 and 9 cover Llama-3.1-70B-SFT with two reward models, while Figure 10 covers aligned models.

D. Case Study

The case studies illustrate TPO’s inference-time optimization process through initial responses, textual feedback, and optimized outputs.

  • D. Case Study: Case studies show the query, chosen and rejected initial responses, textual loss, textual gradient, and optimized response.The textual gradient functions as an optimal contextual parameter that reallocates probability mass during optimization.

(The following content may contain unsafe or sensitive information.)

The examples show how textual feedback can revise responses by adding structure, detail, safety-oriented explanation, or broader coverage. The resulting outputs include expanded answers for factual, educational, practical, and coding queries.

  • Response refinement: Textual feedback expands terse or incomplete answers into more detailed and structured responses.Examples include organizing animation tools, elaborating natural-selection explanations, and adding practical methods for obtaining water.
  • Response refinement: TPO can revise responses by adding contextual explanations and examples to improve accessibility.The optimized natural-selection response adds examples such as finch beak variation and peppered moth coloration.
  • Coding example: The coding example contrasts a multiprocessing pool with a less structured job-based approach.The preferred response includes worker setup and pool.map for distributing items in parallel.
  • Caveat: One generated response demonstrates that optimization can produce unsafe or irrelevant content despite the surrounding preference-feedback process.The example presents code framed around sneezing and coughing etiquette based on user preference or bias.
Loading 2501.12895v1…