Inference-time Alignment in Continuous Space
This addresses the challenge of limited effectiveness in inference-time alignment for AI models, particularly when base policies are weak or candidate sets are small, offering a novel approach to improve performance in tasks like safety and mathematical reasoning.
The paper tackles the problem of aligning large language models with human feedback at inference time by proposing Simple Energy Adaptation (SEA), which uses gradient-based sampling in continuous latent space to adapt responses, resulting in relative improvements of up to 77.51% on AdvBench and 16.36% on MATH compared to baselines.
Aligning large language models with human feedback at inference time has received increasing attention due to its flexibility. Existing methods rely on generating multiple responses from the base policy for search using a reward model, which can be considered as searching in a discrete response space. However, these methods struggle to explore informative candidates when the base policy is weak or the candidate set is small, resulting in limited effectiveness. In this paper, to address this problem, we propose Simple Energy Adaptation ($\textbf{SEA}$), a simple yet effective algorithm for inference-time alignment. In contrast to expensive search over the discrete space, SEA directly adapts original responses from the base policy toward the optimal one via gradient-based sampling in continuous latent space. Specifically, SEA formulates inference as an iterative optimization procedure on an energy function over actions in the continuous space defined by the optimal policy, enabling simple and effective alignment. For instance, despite its simplicity, SEA outperforms the second-best baseline with a relative improvement of up to $ \textbf{77.51%}$ on AdvBench and $\textbf{16.36%}$ on MATH. Our code is publicly available at https://github.com/yuanyige/sea