Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

InfoPO: On Mutual Information Maximization for Large Language Model Alignment

Teng Xiao, Zhen Ge, Sujay Sanghavi, Tian Wang, Julian Katz-Samuels, Marc Versage, Qingjun Cui, Trishul Chilimbi

TL;DR

This work reframes large-language-model alignment with human preferences as a conditional mutual-information maximization problem, avoiding the Bradley–Terry assumption that underpins many direct preference optimization methods. By adopting the NWJ estimator, InfoPO directly optimizes the mutual information between responses and human preferences given prompts, yielding a conservative, stable update that prevents the chosen response likelihood from monotonically decreasing. The authors show that DPO corresponds to an InfoNCE-based MI objective, and provide theoretical results demonstrating reverse-KL, mode-seeking behavior under InfoPO. Empirically, InfoPO consistently outperforms strong baselines on open benchmarks, particularly in reasoning-heavy tasks, across multiple model families and evaluation settings, including on-policy and off-policy regimes. The work offers a practical, principled path for aligning LLMs with human preferences and suggests avenues for exploring alternative MI estimators in future research.

Abstract

We study the post-training of large language models (LLMs) with human preference data. Recently, direct preference optimization and its variants have shown considerable promise in aligning language models, eliminating the need for reward models and online sampling. Despite these benefits, these methods rely on explicit assumptions about the Bradley-Terry (BT) model, which makes them prone to overfitting and results in suboptimal performance, particularly on reasoning-heavy tasks. To address these challenges, we propose a principled preference fine-tuning algorithm called InfoPO, which effectively and efficiently aligns large language models using preference data. InfoPO eliminates the reliance on the BT model and prevents the likelihood of the chosen response from decreasing. Extensive experiments confirm that InfoPO consistently outperforms established baselines on widely used open benchmarks, particularly in reasoning tasks.

InfoPO: On Mutual Information Maximization for Large Language Model Alignment

TL;DR

This work reframes large-language-model alignment with human preferences as a conditional mutual-information maximization problem, avoiding the Bradley–Terry assumption that underpins many direct preference optimization methods. By adopting the NWJ estimator, InfoPO directly optimizes the mutual information between responses and human preferences given prompts, yielding a conservative, stable update that prevents the chosen response likelihood from monotonically decreasing. The authors show that DPO corresponds to an InfoNCE-based MI objective, and provide theoretical results demonstrating reverse-KL, mode-seeking behavior under InfoPO. Empirically, InfoPO consistently outperforms strong baselines on open benchmarks, particularly in reasoning-heavy tasks, across multiple model families and evaluation settings, including on-policy and off-policy regimes. The work offers a practical, principled path for aligning LLMs with human preferences and suggests avenues for exploring alternative MI estimators in future research.

Abstract

We study the post-training of large language models (LLMs) with human preference data. Recently, direct preference optimization and its variants have shown considerable promise in aligning language models, eliminating the need for reward models and online sampling. Despite these benefits, these methods rely on explicit assumptions about the Bradley-Terry (BT) model, which makes them prone to overfitting and results in suboptimal performance, particularly on reasoning-heavy tasks. To address these challenges, we propose a principled preference fine-tuning algorithm called InfoPO, which effectively and efficiently aligns large language models using preference data. InfoPO eliminates the reliance on the BT model and prevents the likelihood of the chosen response from decreasing. Extensive experiments confirm that InfoPO consistently outperforms established baselines on widely used open benchmarks, particularly in reasoning tasks.
Paper Structure (28 sections, 1 theorem, 24 equations, 3 figures, 3 tables)

This paper contains 28 sections, 1 theorem, 24 equations, 3 figures, 3 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Reinforcement Learning from Human Feedback.
  4. Contrastive Preference Fine-tuning.
  5. Mutual Information Estimation.
  6. Notations and Preliminaries
  7. Methodology
  8. Mutual Information Maximization for Large Language Alignment
  9. Gradient Analysis of DPO
  10. The Proposed InfoPO
  11. Theoretical Analysis
  12. Experiments
  13. Experimental Setup
  14. Datasets.
  15. Models.
  16. ...and 13 more sections

Key Result

Theorem 4.1

Minimizing the InfoPO objective in Equation (Eq:InfoPO) with respect to ${\theta}$ will encourage mode-seeking behavior by minimizing the reverse KL divergence between $\pi_\theta (\mathbf{y}\mid\mathbf{x})$ and unknown distribution of chosen response $\pi_{\rm{chosen}}(\mathbf{y}\mid\mathbf{x})$.

Figures (3)

  • Figure 1: The training dynamics of average likelihood of InfoPO and DPO on the Mistral-7B. We observe that InfoPO exhibits the less decline in average chosen likelihoods, while still achieving the significant increase in margins of rejected and chosen likelihood, compared to DPO. Results on Llama3-8B are given in Section \ref{['exp:anly']}.
  • Figure 2: The training dynamics of average likelihood of InfoPO and DPO on the Llama3-8B. We observe that InfoPO exhibits the less decline in the average chosen likelihoods, while still achieving the significant increase in margins of rejected and chosen likelihood, compared to DPO.
  • Figure 3: The performance comparison on coding tasks.

Theorems & Definitions (2)

  • Theorem 4.1
  • proof