Improving AI response alignment without retraining model parameters
Inference-Time Nash Alignment
Artificial Intelligence
Summary
Aligning AI models to make their answers better usually requires retraining the model or using a lot of data, which is costly. The authors explore a way to improve AI responses at the time of use, without changing the model itself. They treat the problem like a game between two decision-makers and introduce two methods to find a balanced solution. Tests show their methods perform close to fully retrained models and remain effective under different settings.
preference-based fine-tuninginference-time alignmentNash equilibriumtwo-player zero-sum gameBest-of-Nash algorithmNash Mirror Descentduality gapregularization parameterRLHFDPO
Authors
Hadi Hosseini, Debmalya Mandal, Duohan Zhang
Abstract
Preference-based fine-tuning methods such as RLHF and DPO require substantial compute and large preference datasets. They also need direct access to the model parameters which are not provided by many state-of-the art models. Inference-time alignment offers a cost-effective alternative without updating model parameters. However, existing inference-time methods rely on a scalar reward model derived under a Bradley-Terry assumption, which cannot represent general preferences. Following recent work on fine-tuning with generalized preferences, in this work, we initiate the study of inference-time alignment under general preferences. We formulate the problem as obtaining a Nash equilibrium of a two-player zero-sum game between policies. We propose two algorithms: Best-of-Nash (BoN) and Nash Mirror Descent (NMD). We prove that both algorithms achieve a duality gap that matches the problem lower bound. Empirically, we implement the two methods on three datasets, which shows that our methods substantially outperform the base policy, converging to the performance of the fine-tuned models. Moreover, our results show that NMD remains robust across the regularization parameter.