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 may produce unexpected or unsafe outputs without alignment, and whether they can be aligned at inference using interpretable feedback while approaching training-time methods remains an open question. TPO uses textual critiques to iteratively refine outputs without retraining, and after two steps an unaligned Llama-3.1-70B-SFT surpasses its aligned counterpart on nearly all benchmarks.

  • Problem

    Whether LLMs can be aligned during inference with performance on par with training-time methods using interpretable textual feedback remains an open question, given that unaligned models may produce unexpected or unsafe responses.

    Source: §1 · p. 1
  • Method

    TPO translates reward signals into textual critiques and uses them to iteratively refine outputs at test time without updating model parameters.

    Source: §1 · p. 2, §4 · p. 4
  • Results

    Across benchmarks, two TPO steps yield substantial gains, with unaligned Llama-3.1-70B-SFT surpassing its aligned counterpart on nearly all benchmarks.

    Source: §1 · p. 2
  • Takeaways & Limitations

    TPO offers a practical, lightweight alternative for aligning LLM outputs with human preferences during inference without retraining.

    Source: §8 · p. 9

    TPO depends on instruction-following proficiency; Llama-3.1-8B-Instruct loses alignment over successive iterations, and evaluated models were not explicitly trained for TPO.

    Source: §7.4 · p. 9

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 model outputs during inference without retraining by turning reward-model feedback into textual critiques and revision suggestions. The introduction reports benchmark gains and flexible test-time scaling.

  • RLHF (Ouyang et al., 2022) and DPO (Rafailov et al., 2024) update model parameters, while iterative retraining can hinder swift adaptation to changing data and requirements.Recent approaches also use numerical reward-model feedback to guide alignment during inference (Khanov et al., 2024; Qiu et al., 2024).Source: §1 · p. 1
  • At each step, the reward model scores responses, and the policy uses the highest- and lowest-scoring responses to form critiques and textual-gradient suggestions for revising outputs without changing weights.The critiques interpret numerical feedback as textual rewards.Source: §1 · p. 2
  • Evaluations cover instruction following, preference alignment, safety, and mathematics; reward-model scores rise over optimization steps for both unaligned and aligned models.The reward model serves as a proxy for human preferences.Source: §1 · p. 2
  • After two TPO steps, both unaligned and aligned models gain substantially, and Llama-3.1-70B-SFT surpasses Llama-3.1-70B-Instruct on nearly all benchmarks.The comparison is reported across the benchmark evaluations.Source: §1 · p. 2
  • With 22B parameters, an aligned model using TPO achieves an LC score of 53.4% on AlpacaEval 2 and a WR score of 72.2% on Arena-Hard, outperforming established leaderboard entries.Source: §1 · p. 2
  • TPO scales test-time computation through search width and depth, with depth-wise revision compensating for the lower efficiency of purely width-based sampling.Source: §1 · p. 2

2. Related Work

Prior preference optimization typically updates model parameters through training, while inference-time alignment methods intervene during decoding. TPO instead uses reward-model feedback to iteratively refine whole responses in textual form while keeping parameters fixed.

Related Work

  • Training-time preference optimization uses gradient descent, with point-wise methods optimizing individual examples and pairwise methods modeling relative preferences.Examples include PPO (Schulman et al., 2017), ReMax, and KTO (Ethayarajh et al., 2024) as point-wise methods, and DPO (Rafailov et al., 2024), SimPO (Meng et al., 2024), sDPO (Kim et al., 2024), and TR-DPO (Gorbatovski et al., 2024) as pairwise approaches.Source: §2 · p. 2
  • Inference-time alignment includes context optimization, in-context DPO, and reward-guided candidate selection through Best-of-N sampling.Context methods include in-context learning (Lin et al., 2023), retrieval augmentation (Xu et al., 2023), prompt rewriting (Cheng et al., 2023), and in-context DPO (Song et al., 2024); speculative variants include Speculative Best-of-N (Zhang et al.) and TreeBoN (Qiu et al., 2024).Source: §2 · p. 2
  • Other reward-guided decoding methods score generation at token or segment granularity, whereas TPO interprets feedback to optimize the entire response iteratively.Khanov et al. (2024) use token-level reward scores, Li et al. (2024a) evaluate segments, and Liu et al. (2024) use implicit and explicit value functions at token and chunk levels.Source: §2 · pp. 2–3
  • TPO iteratively interprets numerical reward feedback as textual critiques to refine entire responses, aligning preferences at test time without updating model parameters.This distinguishes it from token- or segment-level scoring approaches and frames textual feedback as an inference-time optimization process.Source: §2 · pp. 2–3

3. Preliminary

Preference optimization seeks to increase the likelihood of preferred outputs and decrease that of misaligned outputs. Existing approaches score preference data and generally optimize model parameters during training.

3. Preliminary

  • Preference optimization increases the likelihood of preferred outputs while decreasing the likelihood of misaligned ones.Its scoring function evaluates alignment using prompts paired with winning and losing responses.Source: §3 · p. 3
  • RLHF (Ouyang et al., 2022) uses human feedback in a reward model, while DPO (Rafailov et al., 2024) reparameterizes rewards through the optimal policy.The RLHF objective also penalizes large deviations from a reference policy.Source: §3 · p. 3
  • These training-time methods use gradient-based updates to optimize numerical model parameters, such as neural weights.Each step updates parameters using the learning rate and loss gradient, changing the output distribution to favor preferred generations.Source: §3 · p. 3

4. Method

TPO refines responses at inference time by using textual critiques and update instructions to alter the output distribution while keeping model parameters fixed. In reward-model alignment, it iteratively generates and revises candidates, selecting the highest-scoring response at termination.

4. Method

  • TPO optimizes a contextual parameter through textual updates to responses, reallocating output probability mass while keeping model parameters fixed.Its textual loss, gradient, and update steps are carried out in natural language rather than through numerical parameter updates.Source: §4 · p. 4, §4.1 · p. 4

4.1. Components of TPO

  • TPO adapts gradient descent through four prompt-based components: variable definition, loss calculation, gradient computation, and variable optimization.The response is the variable; prompts elicit critiques, textual update instructions, and a refined response.Source: §4.1 · p. 4

4.2. Test-Time Alignment

  • For test-time alignment, a reward model scores candidate responses, whose best and worst examples ground a textual critique and instructions for generating revised candidates.The reward model acts as a proxy for human preferences; the critique describes the stronger response’s advantages and the weaker response’s shortcomings.Source: §4.2 · p. 4
  • Across up to D iterations, TPO evaluates new candidates, caches their scores, selects the best and worst for the next revision, and returns the cache’s highest-scoring response.The cache’s current highest- and lowest-scoring responses become the chosen and rejected examples for the next iteration.Source: §4.2 · pp. 4–5
  • TPO revises multiple candidates using prior feedback, combining parallel sampling and sequential revision in contrast to static-feedback and tree-search approaches.It aims to optimize context ϕ to alter the model’s output distribution for preference alignment.Source: §4.2 · p. 5

5. Experimental Setup

The experimental setup compares unaligned and aligned policy models, applies TPO at inference, and evaluates preference, instruction-following, safety, and mathematical performance using specified benchmarks and metrics.

Experimental Setup

  • Experiments compare unaligned Llama-3.1-70B-SFT with Llama-3.1-70B-Instruct, Mistral-Small-Instruct-2409, and an on-policy Llama-3.1-70B-DPO baseline.The DPO baseline is trained on UltraFeedback pairs, selecting the highest- and lowest-scoring of five SFT responses per query using a reward model (Meng et al., 2024).Source: §5 · p. 5
  • DPO incorporates reward-model feedback into parameters, whereas TPO integrates it into test-time context; FsfairX-LLaMA3-RM-v0.1 scores all policies, with an additional 8B reward model for SFT (Dong et al., 2023).These baselines enable comparison of parameter updating with contextual updating.Source: §5 · p. 5
  • TPO uses TextGrad’s textual-gradient and update prompts with a customized preference-loss prompt; default inference uses N=5 samples and D=2 benchmark iterations.Training-curve analyses run for five iterations; an ultra setting uses D=5 and N=20.Source: §5 · p. 5, Table 1 · p. 6
  • Benchmarks cover instruction following, preference alignment, safety, and math, with 500 HH-RLHF examples and full test sets elsewhere; evaluation uses reward averages or benchmark-specific metrics (Li et al., 2023a; Li et al., 2024b; Bai et al., 2022; Ji et al., 2023; Röttger et al., 2024; Lightman et al., 2024).Metrics include AlpacaEval 2 raw and length-controlled win rates, Arena-Hard win rate against GPT-4-0314, XSTest WildGuard classification accuracy, MATH-500 zero-shot chain-of-thought pass@1, and average reward for HH-RLHF and BeaverTails.Source: §5 · p. 5

6. Experimental Results

TPO uses iterative reward-model feedback to improve unaligned and already aligned models during inference. The experiments report gains in benchmark performance and inference stability without changing model parameters.

6.1. Test-time Training

  • After two TPO steps, Llama-3.1-70B-SFT reaches or exceeds Llama-DPO and Llama-Instruct under two reward models, with larger gains than revision from step two onward.TPO also further improves already aligned models, while revision adds limited benefit; the first optimization step yields the largest gain, so benchmark evaluations use two steps. The curves average reward-model scores over up to five iterations and compare TPO with revision, which refines a cached response without rejected-response preference cues.Source: §6.1 · pp. 6–7

6.2. Benchmark Performance

  • TPO makes Llama-3.1-70B-SFT outperform Llama-3.1-70B-DPO on every evaluation set and Llama-3.1-70B-Instruct on all metrics except AlpacaEval 2 LC.With the Tulu-3-8B reward model, it reaches 70.5 WR on Arena-Hard, surpassing Llama-3.1-405B-Instruct on the official leaderboard. In an ultra setting of five iterations and 20 samples, SFT exceeds Llama-3.1-70B-Instruct across all metrics.Source: §6.2 · p. 7
  • LC 53.4 brings 22B-parameter Mistral-Small-Instruct-2409 to a level comparable to GPT-4-Turbo on the official leaderboard.TPO also consistently improves Llama-3.1-70B-Instruct and Mistral-Small-Instruct-2409 across diverse tasks; the smaller model gains substantially after two iterations.Source: §6.2 · p. 7
  • TPO improves inference stability: SFT with TPO exceeds aligned models, and TPO further stabilizes aligned models, as measured by lower reward-score standard deviations.The authors interpret this as redistributing probability mass toward high-quality responses that receive higher reward-model scores.Source: §6.2 · p. 7, Figure 4 · p. 7
  • TPO changes no model parameters and uses a negligible fraction of training-time preference-optimization computation while improving performance metrics.The paper presents test-time optimization as tailoring outputs to reward-model feedback.Source: §6.2 · p. 7

7. Analysis

The analysis describes textual feedback as the mechanism for iterative refinement, examines width and depth scaling and computational cost, and identifies instruction-following as a prerequisite.

7.1. Policy-Reward Interaction with Textual Feedback

  • TPO turns numerical reward feedback into textual improvement suggestions, which the policy model follows to refine responses over successive iterations.In a desert-water example, textual feedback recommended adding examples and limitations, which appeared in the updated response.Source: §7.1 · pp. 7–8

7.2. Test-time Scaling in Width and Depth

  • Increasing search width from 5 to 20 consistently boosts performance before plateauing, while smaller widths can catch up through additional revision rounds.A larger width accelerates convergence; TPO-D2-N15 after two iterations has a reward-model score comparable to TPO-D4-N5 and TPO-D3-N10.Source: Figure 5 · p. 8, §7.2 · p. 8
  • 65.2% and 57.5% average win rates against BoN with 30 and 60 samples, respectively, were achieved by TPO-D2-N5 using 15 total responses.The comparison used 100 randomly selected instances from AlpacaEval 2, Arena-Hard, and HH-RLHF, with GPT-4 win-rate evaluation.Source: Figure 6 · p. 8, §7.2 · pp. 8–9

7.3. Scaling Computing from Training-time to Test-time

  • 72,840 PFLOPs are required to train Llama-3.1-70B-DPO, compared with about 9.3 PFLOPs per TPO query, less than 0.01% of that training cost.The comparison assumes 64k training instances of maximum length 2,048 and a TPO maximum context length of 4,096; the paper notes that TPO’s relative cost advantage over Instruct training is greater still.Source: §7.3 · p. 9

7.4. Instruction Following as a Prerequisite

  • Llama-3.1-8B-Instruct’s reward-model scores decrease across TPO iterations, indicating that instruction-following proficiency is necessary for TPO to succeed.The evaluated models were not explicitly trained for TPO-specific tasks such as textual loss calculation or gradient computation; more meticulous fine-tuning for those instruction-following tasks is left for future work.Source: Figure 7 · p. 9, §7.4 · p. 9

8. Conclusion

Test-time Preference Optimization (TPO) refines LLM outputs during inference by turning reward signals into textual feedback, without retraining. Across benchmarks, it can bridge or exceed the performance gap between unaligned and aligned LLMs, and scales through search width and depth.

8. Conclusion

  • Across instruction-following, safety, and mathematical-reasoning benchmarks, a few TPO steps can bridge or exceed the performance gap between unaligned and aligned LLMs.The conclusion reports this capability across a variety of benchmarks.Source: §8 · p. 9
  • TPO aligns LLM outputs during inference without retraining by turning reward signals into textual critiques and improvements that iteratively refine responses.The paper describes these textual signals as “textual loss” and “textual gradients,” shifting probability mass toward human-preferred responses.Source: §8 · p. 9
  • TPO scales flexibly through search width and depth, supporting more efficient exploration and refinement in resource-constrained settings.Source: §8 · p. 9
  • Case studies suggest that TPO’s success depends on a model’s ability to interpret and act on textual feedback, making strong instruction-following capabilities essential.Source: §8 · p. 9
  • Future work may refine textual interaction protocols for specialized tasks, explore advanced reward models, and adapt or pre-train weaker models for better TPO performance.Source: §8 · p. 9

A. Prompt Design

TPO retains TextGrad’s vanilla prompts for Pgrad and Pupdate while adding a customized textual loss prompt, Ploss, for preference optimization. Ploss compares chosen and rejected responses; the revision baseline evaluates a single response.

A. Prompt Design

  • TPO uses TextGrad’s vanilla prompts for Pgrad and Pupdate and adds a customized textual loss function, Ploss, for test-time preference optimization.Source: A. Prompt Design · p. 13
  • Ploss asks the model to compare chosen and rejected responses, analyze each response’s strengths and weaknesses, and explain the choice without generating a query response.The prompt also instructs the model to be concise.Source: A. Prompt Design · p. 13
  • Unlike TPO’s Ploss, the revision baseline evaluates one model response to a user query by analyzing its strengths and weaknesses.Its prompt also says not to generate a response and to be concise.Source: A. Prompt Design · p. 13

B. Data Statistics

The section presents statistical data for all benchmark datasets used in the work.

B. Data Statistics

C. Test-time Training

The section presents test-time training curves across datasets and models, with separate figures for two reward-model configurations and aligned models.

  • Figure 8 shows test-time training curves for Llama-3.1-70B-SFT using FsfairX-LLaMA3-RM-v0.1.Source: Figure 8 · p. 14
  • Figure 9 shows test-time training curves for Llama-3.1-70B-SFT using Llama-3.1-Tulu-3-8B-RM.Source: Figure 9 · p. 15
  • Figure 10 shows test-time training curves for Llama-3.1-70B-Instruct and Mistral-Small-Instruct-2409.Source: Figure 10 · p. 16

D. Case Study

The case studies illustrate TPO’s practical application across representative benchmark datasets and show how textual gradients refine outputs without changing model parameters.

  • The case studies were selected from several representative benchmark datasets to demonstrate TPO’s practical application and effectiveness.Source: D. Case Study · p. 17
  • Each case presents the original query, chosen and rejected initialization responses, textual loss and gradient, and the optimized response.Source: D. Case Study · p. 17
  • With model parameters θ fixed, the textual gradient acts as a contextual parameter ϕ that reallocates probability mass toward outputs better aligned with human preferences.Source: D. Case Study · p. 17

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

These examples illustrate iterative textual feedback refining responses across factual, practical, coding, and safety-sensitive prompts. The revisions add detail, structure, context, and constructive guidance tailored to each query.

Water in the desert

  • The revised desert-water answer adds historical and regional examples and conservation practices, while noting that available methods depend on local conditions.Examples include wells, rainwater harvesting, oases, desalination, and fog collection.Source: (The following content may contain… · p. 18

2D animation tools

  • The animation-tools answer groups options by professional, online, and open-source use, adds descriptions and pricing, and advises considering system compatibility.The textual feedback had recommended organizing the list and adding pros, cons, pricing, and compatibility information.Source: (The following content may contain… · p. 20

Spaghetti-tree hoax

  • The spaghetti-tree answer explains the hoax’s setup and context, then adds public reactions and a reported viewer response.It notes that spaghetti was relatively unfamiliar in the UK and that the broadcast aired on April Fools’ Day.Source: (The following content may contain… · p. 21

Natural selection

  • The revised natural-selection explanation adds examples, outlines the process from variation to speciation, and clarifies that selection is not conscious or directed toward perfection.It also discusses genetic drift, epigenetics, and human-driven selective pressures.Source: (The following content may contain… · p. 23
  • The selected despill algorithm detects spill with an RGB-distance threshold and alpha-based spill map, then blurs and normalizes the map to weight a chroma-color blend.It recombines the despilled image with the original alpha channel, and the threshold and blur parameters should be adjusted to the image and desired removal level.Source: (The following content may contain… · pp. 24–25
  • The despill feedback recommends testing on known images with varied spill, examining edge cases, and illustrating intermediate steps such as the spill map and weighted blend.These additions are proposed to demonstrate effectiveness and guide parameter choices.Source: (The following content may contain… · p. 25
  • The multiprocessing answer’s pool-and-map skeleton is judged clearer and more comprehensive than the rejected version’s manually started jobs.The chosen response uses multiprocessing.cpu_count() to set the worker count; the rejected worker only prints items.Source: (The following content may contain… · p. 26
  • The optimized multiprocessing example adds exception handling and a customizable safe-processing function so failures can be logged while remaining items continue processing.The example returns None for a failed item and appends “_processed” in its sample processing function.Source: (The following content may contain… · pp. 27–28
  • The sneezing-etiquette response better meets the prompt by branching on health, environmental, and convenience biases and including a dry-run, unlike the generic rejected answer.The chosen function returns recommendations for each bias and an invalid-bias message.Source: (The following content may contain… · pp. 29–30

Examples

  • The etiquette example turns feedback into a bias-specific Python function with recommendations for health, environmental impact, convenience, and cultural practices.The described program uses if/else cases and includes explanations and examples for applying the recommendations.Source: (The following content may contain… · p. 31
  • The frozen-yogurt revision clarifies upfront that it differs from ice cream, contrasts their ingredients, and advises checking labels for dietary needs.It also notes frozen yogurt’s typically lower fat content and tangy flavor, while warning that added sweeteners and flavorings may affect sugar and nutritional content.Source: (The following content may contain… · p. 32
  • The motion-sickness revision retains the sensory-conflict explanation and adds contributing factors, common travel settings, and prevention or management techniques.It names individual sensitivity, prior experiences, and medical conditions, and suggests focusing on the horizon, deep breathing, or over-the-counter medication.Source: (The following content may contain… · p. 33
  • For a politically violent query, the optimized response acknowledges the quote’s historical context while rejecting violence and redirecting toward peaceful advocacy resources.It recommends respectful dialogue and legal and democratic processes, and names the ACLU and Nonviolent Communication as resources.Source: (The following content may contain… · p. 34
  • The weight-guidance revision presents BMI, waist-to-hip ratio, body-fat percentage, and professional consultation as estimates, emphasizing individual differences.It gives BMI and ratio thresholds and adds examples and advice for interpreting results, including possible steps when measurements indicate elevated risk.Source: (The following content may contain… · p. 36
  • The same weight response adds measurement resources, realistic goal-setting, progress tracking, patience, and alternatives for people facing access barriers.It points readers to CDC and American Council on Exercise calculators and mentions online resources, community health programs, and support groups.Source: (The following content may contain… · p. 36
  • For a request to craft offensive language, the optimized refusal explains the concern and offers a respectful alternative using “I” statements.The preceding feedback also suggests communication resources, examples, reassurance, and discussion of potential consequences, though the shown optimized response is truncated after beginning its resource suggestion.Source: (The following content may contain… · p. 37

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

  • For the relationship dilemma, feedback expanded the response with criteria for valid exceptions, long-term effects on trust, cultural differences, alternatives to lying, and communication and boundaries.The optimized response also identifies fear of judgment, conflict, or causing discomfort as possible motivations for lying.Source: (The following content may contain… · p. 39
  • The chosen refusal was preferred because it explains safety and privacy concerns and invites further questions, unlike the brief rejected refusal.Feedback then made the refusal more explicit about harm and illegality, and added peaceful political-engagement alternatives such as dialogue, voting, and contacting representatives.Source: (The following content may contain… · pp. 40–41
  • For money laundering, feedback added examples for each stage, detection and prevention measures, and broader economic and security impacts to the structured explanation.The optimized response illustrates placement with small cash deposits, layering with transfers across accounts and countries, and integration with real estate or business investment; it also mentions the Bank Secrecy Act and terrorism financing.Source: (The following content may contain… · pp. 42–43
Loading 2501.12895v1…