Research map / Probabilistic and scientific modeling / Variational and generative sampling
BRAIn: Bayesian Reward-conditioned Amortized Inference for natural language generation from feedback
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes BRAIn, a distribution-matching RLHF method that targets a Bayes-rule reward-conditioned posterior and uses a self-normalized baseline to cut gradient variance, showing DPO-sft is a special case.
Problem
Distribution matching methods (GDC, DPG) for aligning LMs to rewards have high-variance gradients and have been overlooked relative to contrastive DPO-style RLHF methods.
Approach
Define the target as a reward-conditioned posterior p(y|x,G=1) via Bayes' rule with the SFT model as prior, and subtract a self-normalized importance-sampling baseline from the gradient; the resulting gradient is unbiased for the negative self-normalized KL, and with n=2 and argmax weights it reduces to DPO-sft.
Key results
- On Anthropic HH, BRAIn reaches 95.40% win rate against gold under the training reward model versus 54.60% for DPO and 87.37% for DPO-sft.
- With Mixtral-8x7B as judge, BRAIn achieves 74.36% win rate on Anthropic HH and 64.74% on Reddit TL;DR, above DPO, DPO-sft and RSO.
- Removing self-normalization from the baseline drops Anthropic HH win rate from 95.4 to 59.1, showing the baseline's key role.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use BRAIn instead of the strongest baseline, RSO | QLoRA-tuned Llama-7B alignment · AnthropicHH test | Train RM win-rate: 84.59$\pm$0.99 -> 95.40$\pm$0.57 | Helped |
| Use BRAIn instead of the strongest baseline, RSO | QLoRA-tuned Llama-7B alignment · AnthropicHH test | LLM eval win-rate: 67.22$\pm$1.29 -> 74.36$\pm$1.20 | Helped |
| Use BRAIn instead of the strongest baseline, RSO | CarperAI summarization LM alignment · Reddit TL;DR test | Train RM win-rate: 91.24$\pm$0.68 -> 95.21$\pm$0.52 | Helped |
| Use BRAIn instead of the strongest baseline, DPO-sft | CarperAI summarization LM alignment · Reddit TL;DR test | LLM eval win-rate: 60.55$\pm$1.18 -> 64.74$\pm$1.16 | Helped |
| Add self-normalization to the baseline (ablation) | BRAIn · Reddit TL;DR | Train RM win-rate: 61.4 -> 95.2 | Helped |
| Add the baseline (ablation) | BRAIn · Reddit TL;DR | Train RM win-rate: 61.1 -> 95.2 | Helped |
| Add self-normalization to the baseline (ablation) | BRAIn · AnthropicHH | Train RM win-rate: 59.1 -> 95.4 | Helped |
| Add the baseline (ablation) | BRAIn · AnthropicHH | Train RM win-rate: 58.3 -> 95.4 | Helped |
| Replace argmax importance weighting with softmax (ablation) | DPO-sft · AnthropicHH | Train RM win-rate: 87.37$\pm$0.91 -> 89.21$\pm$0.85 | Helped |
| Replace argmax importance weighting with softmax (ablation) | DPO-sft · AnthropicHH | LLM eval win-rate: 66.78$\pm$1.29 -> 69.28$\pm$1.27 | Helped |
| Use softmax over 32 outputs instead of two outputs (ablation) | DPO-sft+IW · AnthropicHH | Train RM win-rate: 89.21$\pm$0.85 -> 93.30$\pm$0.69 | Helped |
| Use softmax over 32 outputs instead of two outputs (ablation) | DPO-sft+IW · AnthropicHH | LLM eval win-rate: 69.28$\pm$1.27 -> 73.89$\pm$1.21 | Helped |
Surprising findings
- DPO-sft has much higher reward-model win rate than DPO but matches it under the LLM judge, suggesting reward hacking.
- Without self-normalization in the baseline, performance drops to roughly the level of DPO, nearly eliminating the gain.
Limitations
- Self-normalized KL divergence properties are not fully explored beyond the proofs given.
- GPT-4 evaluation limited to 500 examples per dataset due to budget.
Experimental setup
Datasets: Reddit TL;DR Anthropic HH
Benchmarks: Reddit TL;DR Anthropic HH
Metrics: win rate against gold head-to-head win/tie/loss rate
Baselines: DPO DPO-sft RSO GDC GDC++
Builds on: DPO GDC GDC++ DPG RSO
Closest papers by meaning
- Amortizing intractable inference in large language models (ICLR 2024 · similarity 0.89)
- Sequence to Sequence Reward Modeling: Improving RLHF by Language Feedback (AAAI 2025 · similarity 0.89)
- Approximated Variational Bayesian Inverse Reinforcement Learning for Large Language Model Alignment (AAAI 2025 · similarity 0.88)
- Active Reward Modeling: Adaptive Preference Labeling for Large Language Model Alignment (ICML 2025 · similarity 0.88)
- Reducing the Probability of Undesirable Outputs in Language Models Using Probabilistic Inference (NeurIPS 2025 · similarity 0.88)
- BOND: Aligning LLMs with Best-of-N Distillation (ICLR 2025 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.