Soft Best-of-n Sampling for Model Alignment

Abstract

Best-of-n (BoN) sampling is a practical approach for aligning language model outputs with human preferences without expensive fine-tuning. BoN sampling is performed by generating n responses to a prompt and then selecting the sample that maximizes a reward function. BoN yields high reward values in practice at a distortion cost, as measured by the KL-divergence between the sampled and original distribution. This distortion is coarsely controlled by varying the number of samples: larger n yields a higher reward at a higher distortion cost. We introduce Soft Best-of-n sampling, a generalization of BoN that allows for smooth interpolation between the original distribution and reward-maximizing distribution through a temperature parameter . We establish theoretical guarantees showing that Soft Best-of-n sampling converges sharply to the optimal tilted distribution at a rate of O(1/n) in KL and the expected (relative) reward. For sequences of discrete outputs, we analyze an additive reward model that reveals the fundamental limitations of blockwise sampling.

Tasks

Reproductions