# BRAIn: Bayesian Reward-conditioned Amortized Inference for natural language generation from feedback

arXiv: https://arxiv.org/abs/2402.02479 · HTML: https://mlsanity.com/paper/2402.02479/

## Key facts

- Cited by 7 papers in the map (data through 14 Oct 2026).
- Placed in Reward modeling and RLHF, one of 10 directions in LLM alignment / RLHF / preference optimisation.
- Builds on 25 of its 37 references that are in the map.
- First posted on arXiv 4 Feb 2024.

## Placed in

- LLM alignment / RLHF / preference optimisation / [Reward modeling and RLHF](https://mlsanity.com/alignment/reward-modeling-and-rlhf/)

## Builds on (most cited first, 25 of 25)

- [Training language models to follow instructions with human feedback](https://mlsanity.com/paper/2203.02155/) (2022)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023)
- [Judging LLM-as-a-Judge with MT-Bench and Chatbot Arena](https://mlsanity.com/paper/2306.05685/) (2023)
- [Direct Preference Optimization: Your Language Model is Secretly a Reward Model](https://mlsanity.com/paper/2305.18290/) (2023)
- [Training a Helpful and Harmless Assistant with Reinforcement Learning from Human Feedback](https://mlsanity.com/paper/2204.05862/) (2022)
- [QLoRA: Efficient Finetuning of Quantized LLMs](https://mlsanity.com/paper/2305.14314/) (2023)
- [Learning to summarize from human feedback](https://mlsanity.com/paper/2009.01325/) (2020)
- [Fine-Tuning Language Models from Human Preferences](https://mlsanity.com/paper/1909.08593/) (2019)
- [Mixtral of Experts](https://mlsanity.com/paper/2401.04088/) (2024)
- [Decision Transformer: Reinforcement Learning via Sequence Modeling](https://mlsanity.com/paper/2106.01345/) (2021)
- [A General Theoretical Paradigm to Understand Learning from Human Preferences](https://mlsanity.com/paper/2310.12036/) (2023)
- [Zephyr: Direct Distillation of LM Alignment](https://mlsanity.com/paper/2310.16944/) (2023)
- [Reinforcement Learning and Control as Probabilistic Inference: Tutorial and Review](https://mlsanity.com/paper/1805.00909/) (2018)
- [RRHF: Rank Responses to Align Language Models with Human Feedback without tears](https://mlsanity.com/paper/2304.05302/) (2023)
- [SLiC-HF: Sequence Likelihood Calibration with Human Feedback](https://mlsanity.com/paper/2305.10425/) (2023)
- [Statistical Rejection Sampling Improves Preference Optimization](https://mlsanity.com/paper/2309.06657/) (2023)
- [Pretraining Language Models with Human Preferences](https://mlsanity.com/paper/2302.08582/) (2023)
- [Quark: Controllable Text Generation with Reinforced Unlearning](https://mlsanity.com/paper/2205.13636/) (2022)
- [Camels in a Changing Climate: Enhancing LM Adaptation with Tulu 2](https://mlsanity.com/paper/2311.10702/) (2023)
- [Defining and Characterizing Reward Hacking](https://mlsanity.com/paper/2209.13085/) (2022)
- [Controlling Linguistic Style Aspects in Neural Language Generation](https://mlsanity.com/paper/1707.02633/) (2017)
- [A Distributional Approach to Controlled Text Generation](https://mlsanity.com/paper/2012.11635/) (2020)
- [On Reinforcement Learning and Distribution Matching for Fine-Tuning Language Models with no Catastrophic Forgetting](https://mlsanity.com/paper/2206.00761/) (2022)
- [Distributional Reinforcement Learning for Energy-Based Sequential Models](https://mlsanity.com/paper/1912.08517/) (2019)

## Built on by (most cited first, 7 of 7)

- [HelpSteer2: Open-source dataset for training top-performing reward models](https://mlsanity.com/paper/2406.08673/) (2024)
- [Nemotron-4 340B Technical Report](https://mlsanity.com/paper/2406.11704/) (2024)
- [Reward-aware Preference Optimization: A Unified Mathematical Framework for Model Alignment](https://mlsanity.com/paper/2502.00203/) (2025)
- [Optimal Policy Minimum Bayesian Risk](https://mlsanity.com/paper/2505.17242/) (2025)
- [PIPA: Preference Alignment as Prior-Informed Statistical Estimation](https://mlsanity.com/paper/2502.05773/) (2025)
- [DIVERGE: Diversity-Enhanced RAG for Open-Ended Information Seeking](https://mlsanity.com/paper/2602.00238/) (2026)
- [DKL: Decoupled Knowledge Learning for Instruction-Tuned Language Models](https://mlsanity.com/paper/2609.02685/) (2026)

## Abstract (arXiv)

Distribution matching methods for language model alignment such as Generation with Distributional Control (GDC) and Distributional Policy Gradient (DPG) have not received the same level of attention in reinforcement learning from human feedback (RLHF) as contrastive methods such as Sequence Likelihood Calibration (SLiC), Direct Preference Optimization (DPO) and its variants. We identify high variance of the gradient estimate as the primary reason for the lack of success of these methods and propose a self-normalized baseline to reduce the variance. We further generalize the target distribution in DPG, GDC and DPO by using Bayes' rule to define the reward-conditioned posterior. The resulting approach, referred to as BRAIn - Bayesian Reward-conditioned Amortized Inference acts as a bridge between distribution matching methods and DPO and significantly outperforms prior art in summarization and Antropic HH tasks.