Regularized Best-of-N Sampling with Minimum B ayes Risk Objective for Language Model Alignment (2025)

Abstract: Best-of-N (BoN) sampling with a reward model has been shown to be an effective strategy for aligning Large Language Models (LLMs) to human preferences at the time of decoding. BoN sampling is susceptible to a problem known as reward hacking when the accuracy of the reward model is not high enough. Because the reward model is an imperfect proxy for the true objective, over-optimizing its value can compromise its performance on the true objective. A common solution to prevent reward hacking in preference learning techniques is to optimize a reward using proximity regularization, which ensures that the language model remains close to the reference model. In this research, we propose MBR-BoN, a variant of BoN that aims to mitigate reward hacking at inference time by incorporating the Minimum Bayes Risk (MBR) objective as a proximity regularization term. We show empirically and analytically that the MBR objective quantifies the proximity of the response to the reference policy, serving as a proximity regularizer for BoN sampling. We evaluate MBR-BoN on the AlpacaFarm and Anthropic’s hh-rlhf datasets and show that it outperforms both BoN sampling and MBR decoding. As an application of MBR-BoN, we use it to generate a pairwise preference learning dataset. Experimental results show that DPO models trained on a dataset generated with MBR-BoN outperform a DPO model generated with vanilla BoN.

## introduction language model alignment is a widely used technique for optimizing the behavior of large language models (llms) to human preferences, steering the models to generate informative, harmless, and helpful responses @xcite @xcite . best-of-n (bon) sampling is widely used to align the llm at decoding time @xcite . bon samples n responses from the language model and selects the best response ac-cording to the proxy reward model as the output of the system. however, bon sampling is known to suffer from the reward hacking problem @xcite @xcite @xcite . the reward hacking is a phenomena where the learning agent overfits to the misspecified reward model, failing to optimise for the true intended objective @xcite . the problem occurs because of reward misspecification; the proxy reward trained from a human preference dataset of a limited quality or quantity does not perfectly reflect true human preferences. as a result, optimizing for the reward model does not always optimize for the preference of the true intended objective. for example, @xcite shows that with 25% label noise, which is the amount of disagreement observed in real-world preference annotations @xcite , bon sampling degrades performance with n greater than 16 @xcite . @xcite shows that even when the proxy reward model performs reasonably well relative to the reference model, it still exhibits overoptimization behavior. we also observe the degradation of performance with n greater than 32 when the amount of train data for the proxy reward model is limited (appendix a). given that human preferences depend on the domain, language, culture, and various other factors of the users @xcite @xcite @xcite @xcite , it is desirable to develop a method that is robust to the situation where the reward model is misspecified due to limited quality and/or quality of the preference dataset. a common approach to mitigate reward hacking in preference learning is to add a proximity regularization term to the loss function to keep the trained model close to the reference model @xcite @xcite . previous work in bon has shown that reducing the number of samples n mitigates the reward hacking @xcite @xcite . this approach successfully increases the proximity to the reference policy @xcite but at the expense of diminished improvement obtained by the method. to this end, we propose mbr-bon, a method that introduces the minimum bayes risk (mbr) objective @xcite @xcite as a proximity regularization term into the bon to mitigate the reward hacking problem. 1 the mbr objective serves as a proximity regularizer by its nature which we show in section 3. instead of optimizing the raw reward score, we optimize a sum of the reward score and a regularization term. mbr-bon can tune the regularization strength by the hyperparameter β, similar to the proximity regularization in rlhf and dpo. we evaluate the performance of mbr-bon on the alpacafarm @xcite and anthropic's hh-rlhf datasets @xcite and show that it outperforms the performance of vanilla bon in a wide range of settings. we also use mbr-bon to generate a pairwise preference learning dataset and show that a model trained by dpo on a dataset generated with mbr-bon outperforms a model trained on a dataset generated with vanilla bon. ## background first, we give an overview of preference learning algorithms including rlhf and dpo. then we introduce the decoding-time alignment algorithm, bon sampling. ## minimum bayes risk objective is a proximity regularizer although bon sampling is shown to be effective with a decent reward model, it is prone to the reward hacking problem under less accurate reward models @xcite . a naive approach to prevent reward hacking is to introduce a proximity regularizer to the bon sampling in the form of a kl-divergence term, as is common in preference learning methods @xcite @xcite . however, we observe that this strategy does not improve over bon in most cases (appendix d). to this end, we propose to use the mbr objective as a proximity regularizer. first, in section 3.1, we visually show that the mbr objective is correlated with the semantic proximity of the reference policy. then, we show an analytical result in section 3.2 that the mbr objective corresponds to the wasserstein distance @xcite , indicating that the mbr objective by its nature quantifies the proximity of the text to the reference policy. ## mbr-best-of-n (mbr-bon) sampling we propose mbr-best-of-n (mbr-bon) sampling, a variant of bon sampling with an mbr objective as the proximity regularizer, to mitigate the reward hacking problem of bon sampling. mbr-bon uses the mbr objective as the proximity regularizer: the hyperparameter β controls the tradeoff between the reward and proximity to the reference model. using a small β makes the output more aligned with the proxy reward, with β = 0 recovering vanilla bon sampling. a larger β makes the output closer to the behavior of the reference model π ref , with β = +∞ recovering mbr decoding. advantage of wd over kl-divergence. wd is a more suited regularizer than kl-divergence for inference-time algorithms where the number of samples is very small. while kl-divergence is useful for training-time alignment algorithms, it poses several challenges for inference-time algorithms with limited samples. theoretically, any high confidence lower bound on kl-divergence requires a sample size exponential in the value of @xcite . this suggests that estimating kldivergence is unreliable in finite-sample settings. for example, for the first instance of the alpacaeval instruction (what are the names of some famous actors that started their careers on broadway?), the kl-divergence of the randomly sampled 128 responses from mistral has a minimum of 627, a maximum of 5870, a mean of 1854, and a standard deviation of 1039. moreover, kl-divergence is sensitive to small differences in the sequences. specifically, kldivergence can be large even if the underlying sequences differ very little. for example, the two sentences: "yes i will do it." and "yes i'll do it." are considered completely different data instances when computing kl-divergence. conversely, wd considers them to be quite similar data instances. this is because the wd uses the utility function to quantify the divergence and represents the difference between two distributions in terms of the semantic distance between the sequences. this makes the wd a more robust measure against the minor variances that naturally occur in natural language texts. see appendix d for experimental evaluation of using kl-divergence as a regularization term. in addition to being a good proximal regularizer, the mbr objective is a useful text generation objective in its own right. the objective is shown to be effective, outperforming map decoding in a variety of text generation tasks, including instructionfollowing task @xcite @xcite . ## experiments we evaluate the performance of mbr-bon for two use cases. first, we evaluate the performance of mbr-bon for decoding time alignment (section 5.1). then, we evaluate mbr-bon as a sampling strategy to generate a preference learning dataset to be used for dpo (section 5.2). for each subset, we use the first 1000 entries of the train split as the development set and the first 1000 entries of the test split as a test dataset. we use mistral-7b-sft-beta (mistral) and dolly-v2-3b (dolly) as the language models @xcite @xcite . to evaluate mbr-bon under various conditions, we use shp-large, shp-xl (ethayarajh et al., 2022), and oasst @xcite as proxy reward models. we use eurus as a gold reference reward model as it is one of the most accurate reward models @xcite and is open-source which makes the experiments reproducible. the results using other reward models as a gold reference are reported in appendix e and g. the average spearman's rank correlation coefficient ρ @xcite to the gold reference reward (eurus) is reported in we compare the performance of bon, mbr, and mbr-bon. we sample up to @xmath0. we set the top-p to be @xmath1.9 and the temperature to be @xmath2.0 for the nucleus sampling @xcite . for a fair comparison, we use the same set of n responses for all algorithms. we use the sentence bert model @xcite based on mpnet @xcite to compute the senmbr-bon use the development set to select the optimal β. for each pair of a proxy reward and a gold reference reward, we run mbr-bon with β ∈ {10 -6 , 2•10 -6 , 5•10 -6 , 10 -5 , ..., 2•10 1 } and pick the best performing β for @xmath3. we use the same β for all n in evaluation. see appendix g and f for the ablation study on the regularization strength β. ## conclusions we propose mbr-bon, a variant of bon sampling with mbr objective as a proximity regularizer to mitigate the reward hacking problem. we show that the mbr objective is a proximity regularizer by its nature and show it in the experiments. we evaluate the performance of mbr-bon using the alpaca-farm and anthropic's hh-rlhf datasets. the result shows that mbr-bon outperforms bon when the proxy reward is weakly correlated with the reference objective. as an application of the method, we also show that mbr-bon is an effective strategy for generating a preference dataset for dpo. we believe that mbr-bon will be a practical choice for future decoding-time alignment methods because of its applicability and performance improvements.