Title: PITA: Preference-Guided Inference-Time Alignment for LLM Post-Training
Primary area: foundation or frontier models, including LLMs
Keywords: large language models, post-training, preference learning, Soft-Q learning

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, bypassing the pre-trained reward model dependency. The problem is framed as identifying an underlying preference distribution, solved through 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.