Soft Best-of-n Sampling for Model Alignment

TL;DR

This paper introduces Soft Best-of-n sampling, a method that smoothly interpolates between original and reward-optimized distributions, with theoretical guarantees on convergence and insights into sampling limitations.

Contribution

It proposes Soft Best-of-n sampling, extending traditional BoN sampling with a temperature parameter, and provides theoretical analysis of its convergence and fundamental limitations.

Findings

Soft Best-of-n converges to the optimal tilted distribution at O(1/n) rate.

The method allows smooth control over reward and distribution distortion.

Analysis reveals fundamental limitations of blockwise sampling for discrete sequences.

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 $\lambda$. 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.