Elamur Research⌕
☰
home / papers / 2024 / BOND: Aligning LLMs with Best-of-N Distillation

BOND: Aligning LLMs with Best-of-N Distillation

Paper · 2024 · Pier Giuseppe Sessa, Robert Dadashi, Leonard Hussenot et al.

Summary

Best-of-N sampling draws N responses from a reference model and keeps the one a reward model scores highest. It gives excellent reward-versus-KL trade-offs, but inference costs N times more. This Google DeepMind paper introduces BOND (Best-of-N Distillation), an RLHF method that trains a policy to produce Best-of-N-quality outputs from a single sample.

The authors first derive a closed-form expression for the Best-of-N output distribution. They show it is the optimum of KL-regularized RLHF under two conditions: a rank-based reward (the log quantile of a response's reward) and a regularization strength tied to N. Alignment then becomes a distribution-matching task.

Their practical variant, J-BOND, has three components:

  • It minimizes the Jeffreys divergence, a blend of mode-covering forward KL and mode-seeking backward KL.
  • It estimates quantiles by Monte Carlo.
  • It repeatedly distils Best-of-N of a slowly moving EMA anchor rather than a fixed reference.

On XSum summarization and Gemma models, J-BOND reaches a better KL-reward Pareto front than REINFORCE baselines and improves benchmark results. The work matters because it moves compute from inference to training.

Key points

  • Best-of-N matches the solution of standard KL-regularized RLHF with β = 1/(N−1). Its implied reward depends only on ranks, so it is unaffected by monotone rescaling of the reward model. The authors conjecture this makes it more resistant to reward hacking.
  • Baselines were REINFORCE with 2 samples per prompt and a leave-one-out baseline, tested at β_RL ∈ {0.001, 0.01, 0.1, 1}. J-BOND used anchor mixing η = 0.02.
  • REINFORCE results depended strongly on the choice of β_RL. J-BOND needs no fixed regularization level: its reward rose steadily while KL grew roughly linearly.
  • J-BOND was used to fine-tune Gemma 1.1 2B/7B, RecurrentGemma 2B/9B and CodeGemma 1.1. Gemma 1.1 IT 7B beat Mistral 7B v0.2 Instruct on safety and instruction following.
  • The closest prior work, Amini et al. (2024), uses only backward KL. BOND's Jeffreys objective and dynamic-anchor iteration are what distinguish it, and the authors report the iterative anchor is critical for the best results.

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

 Fields, topics and concepts

📚 Shelves 4The field of science this work belongs to (OpenAlex fields).Computer Science · 19,503 worksArtificial Intelligence · 18,805 worksArts and Humanities · 17,959 worksLanguage and Linguistics · 17,789 works

🧭 Topics 2Research subjects and languages the work is about.Language models & pretraining · 4,145 worksEfficient models & compression · 328 works

Cite

Cite as Pier Giuseppe Sessa, Robert Dadashi, Leonard Hussenot, Johan Ferret, Nino Vieillard, Alexandre Rame, Bobak Shariari, Sarah Perrin, Abe Friesen, Geoffrey Cideron, Sertan Girgin, Piotr Stanczyk, Andrea Michi, Danila Sinopalnikov, Sabela Ramos, Amelie Heliou, Aliaksei Severyn, Matt Hoffman, Nikola Momchev and Olivier Bachem. 2024. BOND: Aligning LLMs with Best-of-N Distillation. arXiv:2407.14622.

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.48550/arxiv.2407.14622 · arXiv 2407.14622

@misc{sessa2024bond,
  title = {{BOND: Aligning LLMs with Best-of-N Distillation}},
  author = {Pier Giuseppe Sessa, Robert Dadashi, Leonard Hussenot, Johan Ferret, Nino Vieillard, Alexandre Rame, Bobak Shariari, Sarah Perrin, Abe Friesen, Geoffrey Cideron, Sertan Girgin, Piotr Stanczyk, Andrea Michi, Danila Sinopalnikov, Sabela Ramos, Amelie Heliou, Aliaksei Severyn, Matt Hoffman, Nikola Momchev and Olivier Bachem},
  year = {2024},
  doi = {10.48550/arxiv.2407.14622},
  eprint = {2407.14622},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2407.14622}
}

References

65 references, 21 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

  1. R. Agarwal, N. Vieillard, Y. Zhou, P. Stanczyk, S. R. Garea, M. Geist, and O. Bachem. On-policy distillation of language models: Learning from self-generated mistakes. In ICLR, 2024. (p. 6)
  2. A. Ahmadian, C. Cremer, M. Gallé, M. Fadaee, J. Kreutzer, A. Üstün, and S. Hooker. Back to basics: Revisiting REINFORCE style optimization for learning from human feedback in LLMs. arXiv preprint, 2024. (pp. 3 and 12) find
  3. A. Amini, T. Vieira, and R. Cotterell. Variational Best-of-N alignment. arXiv preprint, 2024. (p. 12) find
  4. A. Askell, Y. Bai, A. Chen, D. Drain, D. Ganguli, T. Henighan, A. Jones, N. Joseph, B. Mann, N. DasSarma, N. Elhage, Z. Hatfield-Dodds, D. Hernandez, J. Kernion, K. Ndousse, C. Olsson, D. Amodei, T. Brown, J. Clark, S. McCandlish, C. Olah, and J. Kaplan. A general language assistant as a laboratory… find
  5. A. Beirami, A. Agarwal, J. Berant, A. D’Amour, J. Eisenstein, C. Nagpal, and A. T. Suresh. Theoretical guarantees on the Best-of-N alignment policy. arXiv preprint, 2024. (p. 12) find
  6. A. Botev, S. De, S. L. Smith, A. Fernando, G.-C. Muraru, R. Haroun, L. Berrada, R. Pascanu, P. G. Sessa, R. Dadashi, L. Hussenot, J. Ferret, S. Girgin, O. Bachem, A. Andreev, K. Kenealy, T. Mesnard, C. Hardin, S. Bhupatiraju, S. Pathak, L. Sifre, M. Rivière, M. S. Kale, J. Love, P. Tafti, A. Joulin,… find
  7. D. Calandriello, D. Guo, R. Munos, M. Rowland, Y. Tang, B. A. Pires, P. H. Richemond, C. L. Lan, M. Valko, T. Liu, et al. Human alignment of large language models through online preference optimisation. In ICML, 2024. (p. 12) find
  8. S. Casper, X. Davies, C. Shi, T. K. Gilbert, J. Scheurer, J. Rando, R. Freedman, T. Korbak, D. Lindner, P. Freire, et al. Open problems and fundamental limitations of reinforcement learning from human feedback. TMLR, 2023. (p. 1) find
  9. P. Christiano, B. Shlegeris, and D. Amodei. Supervising strong learners by amplifying weak experts. arXiv preprint, 2018. (p. 13)
  10. P. F. Christiano, J. Leike, T. Brown, M. Martic, S. Legg, and D. Amodei. Deep reinforcement learning from human preferences. In NeurIPS, 2017. (p. 1) find
  11. J. Clark and D. Amodei. Faulty Reward Functions in the Wild. https://openai.com/research/ faulty-reward-functions , 2016. (p. 1) ↗
  12. CodeGemma Team. Codegemma: Open code models based on gemma. arXiv preprint, 2024. (p. 12)
  13. A. Cotra. Iterated distillation and amplification. https://ai-alignment.com/iterated-dis tillation-and-amplification-157debfd1616 , 2018. (p. 13) ↗
  14. T. M. Cover. Elements of information theory. 1999. (p. 21) find
  15. W. Dabney, M. Rowland, M. G. Bellemare, and R. Munos. Distributional reinforcement learning with quantile regression. In AAAI Conference on Artificial Intelligence, 2017. (p. 21) find
  16. H. Dong, W. Xiong, D. Goyal, Y. Zhang, W. Chow, R. Pan, S. Diao, J. Zhang, K. SHUM, and T. Zhang. RAFT: Reward ranked finetuning for generative foundation model alignment. TMLR, 2023. (p. 12)
  17. Y. Flet-Berliac, N. Grinsztajn, F. Strub, E. Choi, C. Cremer, A. Ahmadian, Y. Chandak, M. G. Azar, O. Pietquin, and M. Geist. Contrastive policy gradient: Aligning llms on sequence-level scores in a supervised-friendly fashion. arXiv preprint, 2024. (p. 12) find
  18. R. M. French. Semi-distributed representations and catastrophic forgetting in connectionist networks. Connection Science, 1992. (pp. 1 and 3) find
  19. L. Gao, J. Schulman, and J. Hilton. Scaling laws for reward model overoptimization. In ICML, 2023.
  20. (pp. 1, 7, and 12) M. Geist, B. Scherrer, and O. Pietquin. A theory of regularized markov decision processes. In ICML, 2019. (p. 3) find
  21. Gemini Team. Gemini: A family of highly capable multimodal models. 2023. (p. 1)
  22. Gemma Team. Gemma: Open models based on gemini research and technology. arXiv preprint, 2024. find
  23. (pp. 2, 10, and 12) A. Gorbatovski, B. Shaposhnikov, A. Malakhov, N. Surnachev, Y. Aksenov, I. Maksimov, N. Balagansky, and D. Gavrilov. Learn your reference model for real good alignment. arXiv preprint, 2024. (p. 12) find
  24. L. Gui, C. Gârbacea, and V. Veitch. Bonbon alignment for large language models and the sweetness of Best-of-N sampling. arXiv preprint, 2024. (p. 12) find
  25. S. Guo, B. Zhang, T. Liu, T. Liu, M. Khalman, F. Llinares, A. Rame, T. Mesnard, Y. Zhao, B. Piot, J. Ferret, and M. Blondel. Direct language model alignment from online ai feedback. arXiv preprint, 2024. (p. 12) find
  26. J. Hilton. KL divergence of max-of-n, 2023. (p. 12) find
  27. J. Hoffmann, S. Borgeaud, A. Mensch, E. Buchatskaya, T. Cai, E. Rutherford, D. d. L. Casas, L. A. Hendricks, J. Welbl, A. Clark, et al. Training compute-optimal large language models. In NeurIPS, 2022. (p. 13)
  28. H. Jeffreys. An invariant form for the prior probability in estimation problems. Proceedings of the Royal Society of London. Series A, Mathematical and Physical Sciences, 186(1007):453–461, 1946. find
  29. (p. 5) A. Q. Jiang, A. Sablayrolles, A. Roux, A. Mensch, B. Savary, C. Bamford, D. S. Chaplot, D. d. l. Casas, E. B. Hanna, F. Bressand, et al. Mixtral of experts. arXiv preprint arXiv:2401.04088, 2024. (p. 2) find
  30. D. P. Kingma and J. Ba. Adam: A method for stochastic optimization. In ICLR, 2015. (p. 10)
  31. W. Kool, H. van Hoof, and M. Welling. Buy 4 REINFORCE samples, get a baseline for free! 2019. find
  32. S. Kullback. Information Theory and Statistics. New York, 1959. (p. 4) find
  33. A. Lazaridou, A. Potapenko, and O. Tieleman. Multi-agent communication meets natural language: Synergies between functional and structural language learning. In ACL, 2020. (p. 3) find
  34. Z. Li, T. Xu, Y. Zhang, Y. Yu, R. Sun, and Z.-Q. Luo. Remax: A simple, effective, and efficient reinforcement learning method for aligning large language models. arXiv preprint, 2023. (p. 3) find
  35. R. Munos, M. Valko, D. Calandriello, M. G. Azar, M. Rowland, Z. D. Guo, Y. Tang, M. Geist, T. Mesnard, A. Michi, et al. Nash learning from human feedback. arXiv preprint, 2023. (p. 12)
  36. R. Nakano, J. Hilton, S. Balaji, J. Wu, L. Ouyang, C. Kim, C. Hesse, S. Jain, V. Kosaraju, W. Saunders, et al. Webgpt: Browser-assisted question-answering with human feedback. arXiv preprint, 2021.
  37. (pp. 1 and 12) S. Narayan, S. B. Cohen, and M. Lapata. Don’t give me the details, just the summary! topic-aware convolutional neural networks for extreme summarization. In E. Riloff, D. Chiang, J. Hockenmaier, and J. Tsujii, editors, Conference on Empirical Methods in Natural Language Processing, 20… — Don't Give Me the Details, Just the Summary! Topic-Aware Convolutional Neural Networks for Extreme Summarization
  38. R. Ngo, L. Chan, and S. Mindermann. The alignment problem from a deep learning perspective. arXiv preprint, 2022. (p. 1)
  39. OpenAI. Gpt-4 technical report. 2023. (p. 1)
  40. L. Ouyang, J. Wu, X. Jiang, D. Almeida, C. Wainwright, P. Mishkin, C. Zhang, S. Agarwal, K. Slama, A. Ray, et al. Training language models to follow instructions with human feedback. NeurIPS, 2022.
  41. (p. 1) A. Pace, J. Mallinson, E. Malmi, S. Krause, and A. Severyn. West-of-n: Synthetic preference generation for improved reward modeling. arXiv preprint, 2024. (p. 12) find
  42. A. Pan, K. Bhatia, and J. Steinhardt. The effects of reward misspecification: Mapping and mitigating misaligned models. In ICLR, 2022. (p. 1) find
  43. A. S. Pinto, A. Kolesnikov, Y. Shi, L. Beyer, and X. Zhai. Tuning computer vision models with task rewards. In ICML, 2023. (p. 12) find
  44. J. Qiping Yang, S. Salamatian, Z. Sun, A. Theertha Suresh, and A. Beirami. Asymptotics of language model alignment. arXiv, 2024. (p. 3) find
  45. A. Radford, K. Narasimhan, T. Salimans, and I. Sutskever. Improving language understanding by generative pre-training. 2018. (p. 1) find
  46. A. Radford, J. Wu, R. Child, D. Luan, D. Amodei, and I. Sutskever. Language models are unsupervised multitask learners. 2019. (p. 1) find
  47. R. Rafailov, A. Sharma, E. Mitchell, S. Ermon, C. D. Manning, and C. Finn. Direct preference optimization: Your language model is secretly a reward model. arXiv preprint, 2023. (p. 4)
  48. C. Raffel, N. Shazeer, A. Roberts, K. Lee, S. Narang, M. Matena, Y. Zhou, W. Li, and P. J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. JMLR, 2020. (p. 1)
  49. A. Ramé, N. Vieillard, L. Hussenot, R. Dadashi, G. Cideron, O. Bachem, and J. Ferret. WARM: On the benefits of weight averaged reward models. In ICML, 2024. (p. 13) find
  50. A. Ramé, J. Ferret, N. Vieillard, R. Dadashi, L. Hussenot, P.-L. Cedoz, P. G. Sessa, S. Girgin, A. Douillard, and O. Bachem. WARP: On the benefits of weight averaged rewarded policies. arXiv preprint, 2024. find
  51. (pp. 10, 12, and 13) M. Reid, N. Savinov, D. Teplyashin, D. Lepikhin, T. Lillicrap, J.-b. Alayrac, R. Soricut, A. Lazaridou, O. Firat, J. Schrittwieser, et al. Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context. arXiv preprint, 2024. (p. 1)
  52. P. Roit, J. Ferret, L. Shani, R. Aharoni, G. Cideron, R. Dadashi, M. Geist, S. Girgin, L. Hussenot, O. Keller, et al. Factually consistent summarization via reinforcement learning with textual entailment feedback. In ACL, 2023. (p. 6)
  53. J. Schulman, F. Wolski, P. Dhariwal, A. Radford, and O. Klimov. Proximal policy optimization algorithms. arXiv preprint, 2017. (p. 3) find
  54. J. M. V. Skalse, N. H. R. Howe, D. Krasheninnikov, and D. Krueger. Defining and characterizing reward gaming. In NeurIPS, 2022. (pp. 1 and 3) find
  55. N. Stiennon, L. Ouyang, J. Wu, D. Ziegler, R. Lowe, C. Voss, A. Radford, D. Amodei, and P. F. Christiano. Learning to summarize with human feedback. NeurIPS, 2020. (pp. 1 and 12) find
  56. R. S. Sutton and A. G. Barto. Reinforcement Learning: An Introduction. MIT Press, 1998. (p. 6) find
  57. Y. Tang, D. Z. Guo, Z. Zheng, D. Calandriello, Y. Cao, E. Tarassov, R. Munos, B. Ávila Pires, M. Valko, Y. Cheng, and W. Dabney. Understanding the performance gap between online and offline alignment algorithms. arXiv preprint, 2024. (p. 3) find
  58. H. Touvron, L. Martin, K. Stone, P. Albert, A. Almahairi, Y. Babaei, N. Bashlykov, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint, 2023. (pp. 1, 3, and 12)
  59. A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polosukhin. Attention is all you need. In NeurIPS, 2017. (p. 2)
  60. N. Vieillard, T. Kozuno, B. Scherrer, O. Pietquin, R. Munos, and M. Geist. Leverage the average: an analysis of kl regularization in reinforcement learning. NeurIPS, 2020. (p. 4) find
  61. J. Wei, M. Bosma, V. Zhao, K. Guu, A. W. Yu, B. Lester, N. Du, A. M. Dai, and Q. V. Le. Finetuned language models are zero-shot learners. In ICLR, 2022. (p. 1)
  62. R. J. Williams. Simple statistical gradient-following algorithms for connectionist reinforcement learning. Reinforcement learning, 1992. (pp. 1, 3, 6, 12, and 19) find
  63. Y. Wu, Z. Sun, H. Yuan, K. Ji, Y. Yang, and Q. Gu. Self-play preference optimization for language model alignment. arXiv preprint, 2024. (p. 12) find
  64. J. Q. Yang, S. Salamatian, Z. Sun, A. T. Suresh, and A. Beirami. Asymptotics of language model alignment. arXiv preprint, 2024. (p. 12) find
  65. D. M. Ziegler, N. Stiennon, J. Wu, T. B. Brown, A. Radford, D. Amodei, P. Christiano, and G. Irving. Fine-tuning language models from human preferences. arXiv preprint, 2019. (p. 1) find

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/sessa-2024-bond-aligning-llms-best-n/