PITA: Preference-Guided Inference-Time Alignment for LLM Post-Training

TL;DR

PITA introduces a novel inference-time alignment method for LLMs that directly incorporates user preferences into token generation, avoiding the need for pre-trained reward models and reducing computational costs.

Contribution

It presents a new framework that learns preference-based guidance policies during inference, bypassing reward model dependence and enabling efficient alignment of LLM outputs.

Findings

Effective in aligning outputs with user preferences across tasks

Reduces reliance on pre-trained reward models

Demonstrates improved alignment in mathematical reasoning and sentiment classification

Abstract

Inference-time alignment enables large language models (LLMs) to generate outputs aligned with end-user preferences without further training. Recent post-training methods achieve this by using small guidance models to modify token generation during inference. These methods typically optimize a reward function KL-regularized by the original LLM taken as the reference policy. A critical limitation, however, is their dependence on a pre-trained reward model, which requires fitting to human preference feedback--a potentially unstable process. In contrast, we introduce PITA, a novel framework that integrates preference feedback directly into the LLM's token generation, eliminating the need for a reward model. PITA learns a small preference-based guidance policy to modify token probabilities at inference time without LLM fine-tuning, reducing computational cost and bypassing the pre-trained reward model dependency. The problem is framed as identifying an underlying preference distribution, solved through stochastic search and iterative refinement of the preference-based guidance model. We evaluate PITA across diverse tasks, including mathematical reasoning and sentiment classification, demonstrating its effectiveness in aligning LLM outputs with user preferences.