PaperlayerView exact revision on arXiv

Source-linked AI summary

From Data to Rewards: a Bilevel Optimization Perspective on Maximum Likelihood Estimation

Abdelhakim Benechehab, Gabriel Singer, Corentin Léger, Youssef Attia El Hili, Giuseppe Paolo, Albert Thomas, Maurizio Filippone, Balázs Kégl

arXiv:2510.07624v3stat.MLcs.LG

TL;DR

The paper asks how to use policy-gradient training when its reward signal must be learned from data, rather than supplied directly. It formulates reward learning as a bilevel optimization problem and concludes that an optimal reward lets policy gradients use their advantages more effectively.

  • Problem

    Policy-gradient methods can improve preference alignment over supervised fine-tuning but depend on reward signals, motivating reward learning from data.

    Source: §1 · pp. 1–2

  • Method

    A bilevel framework optimizes the reward at the outer level using MLE, while policy gradients train model parameters at the inner level.

    Source: §1 · p. 2

  • Results

    The authors conclude that an optimal reward enables a more effective use of policy-gradient methods' advantages.

    Source: §7 · p. 10

  • Takeaways & Limitations

    The framework provides a way to choose rewards so policy-gradient training produces models optimal in maximum likelihood with respect to observed data.

    Source: §7 · pp. 9–10

    The reward parametrization is restrictive, and further large-scale experiments are needed to assess generalizability to complex applications.

    Source: §7 · p. 10

Abstract

from arXiv · show

Generative models form the backbone of modern machine learning, underpinning state-of-the-art systems in text, vision, and multimodal applications. While Maximum Likelihood Estimation has traditionally served as the dominant training paradigm, recent work have highlighted its limitations, particularly in generalization and susceptibility to catastrophic forgetting compared to Reinforcement Learning techniques, such as Policy Gradient methods. However, these approaches depend on explicit reward signals, which are often unavailable in practice, leaving open the fundamental problem of how to align generative models when only high-quality datasets are accessible. In this work, we address this challenge via a Bilevel Optimization framework, where the reward function is treated as the optimization variable of an outer-level problem, while a policy gradient objective defines the inner-level. We then conduct a theoretical analysis of this optimization problem in a tractable setting and extract insights that, as we demonstrate, generalize to applications such as tabular classification and model-based reinforcement learning. We release the code at https://github.com/abenechehab/nll_to_po .

1 INTRODUCTION

MLE remains a central training approach, while policy-gradient methods have shown advantages for aligning models with preferences but rely on reward signals. The paper proposes learning a reward from data through bilevel optimization, with MLE as the outer objective and policy gradient as the inner objective.

  • MLE maximizes the probability of observed data and remains a backbone of generative-model training, including next-token prediction in LLMs.Source: §1 · p. 1
  • Policy-gradient methods have shown stronger preference alignment and generation quality than supervised fine-tuning, using explicit or implicit reward signals.Source: §1 · p. 1
  • The paper asks whether an implicit reward can be learned from unlabeled data so policy optimization can train models more effectively than MLE.Source: §1 · p. 2
  • The framework optimizes a reward function at the outer level using MLE, while an inner policy-gradient objective trains the model parameters.Source: §1 · p. 2

2 RELATED WORK

The related work spans generative-model training, reward learning, and bilevel optimization. The paper positions its approach as a general framework for using data-derived rewards with policy gradients across MLE tasks.

  • Prior generative-model work uses MLE or its approximations, while policy-gradient methods offer an alternative when reward functions are available.Source: §2 · p. 2
  • Reward-learning paradigms include inverse reinforcement learning, goal-based rewards, and preference-based rewards, with preference learning supporting generative-model post-training.Source: §2 · p. 3
  • Bilevel optimization has been applied to hyperparameter tuning, reinforcement learning, meta-learning, and generative-model training, using approaches including implicit differentiation.Source: §2 · p. 3
  • Unlike the closest cited bilevel approach, which combines MLE with maximum-entropy inverse RL for episodic control, this work targets a general framework for data modalities and MLE tasks.Source: §2 · p. 3

3 PRELIMINARIES

The preliminaries motivate policy gradients where likelihood objectives do not match evaluation or where post-training must preserve prior knowledge, then formulate reward learning as a bilevel problem. The setup contrasts maximum-likelihood fitting with entropy-regularized policy-gradient training.

3.1 MOTIVATION

  • For sequence prediction, policy gradients optimize sequence-level rewards such as BLEU or ROUGE, whereas autoregressive MLE optimizes token-level likelihoods.Source: §3.1 · p. 4
  • Prior studies suggest on-policy RL fine-tuning preserves prior knowledge while specializing to new distributions more effectively than supervised fine-tuning.Source: §3.1 · p. 4
  • The proposed framework interprets data signals as reward functions, enabling policy-gradient optimization instead of direct likelihood maximization.Source: §3.1 · p. 4

3.2 PROBLEM SETUP

  • MLE fits a conditional model to observed input-output data by maximizing log-likelihood, equivalently minimizing KL divergence.Source: §3.2 · p. 4
  • The policy-gradient objective maximizes a reward evaluating generated outputs against observations, with entropy regularization controlled by λ.Source: §3.2 · p. 4
  • The bilevel formulation makes the reward function the outer-level optimization variable, with the MLE objective determining its optimum.Source: §3.2 · p. 4

4 SOLVING BI-O IN A TRACTABLE CASE

Under Gaussian data and reward assumptions, the paper derives closed-form bilevel solutions and interprets the optimal reward as a covariance-scaled distance. Synthetic experiments show that the optimal-reward policy-gradient variant matches the NLL baseline in NLL while converging faster on moment matching.

4.1 THEORETICAL ANALYSIS

  • With Gaussian conditional densities and a quadratic reward, the bilevel problem has a unique closed-form solution, including an isotropic-case solution.Source: §4.1 · p. 5
  • The optimal reward is a negative squared Mahalanobis distance, with its matrix inversely proportional to the data-generating covariance.This implies noisier data receive a weaker penalty for deviations from samples; λ balances reward maximization and entropy regularization.Source: §4.1 · p. 5
  • Substituting the optimal reward makes the inner policy-gradient objective equivalent to minimizing reverse KL between the model and data distributions.Source: §4.1 · p. 5

4.2 EMPIRICAL VALIDATION

  • The optimal-reward PG variant nearly matches the NLL baseline in validation NLL and converges faster on mean and covariance moment matching.Vanilla PG with U = I_n can have diverging NLL when variance shrinks to zero for some λ values, causing numerical instabilities.Source: §4.2 · p. 6, Figure 1 · p. 6
  • For one training point, the optimal-reward PG variant fits the target distribution naturally, while the NLL baseline initially increases variance sharply before matching it.The paper relates this behavior to reverse-KL mode seeking, which encourages concentration on the data distribution's mode.Source: §4.2 · p. 6, Figure 2 · p. 6

5 SOLVING BI-O IN GENERAL

The section develops implicit differentiation as a way to optimize the outer objective beyond tractable distributional settings. On synthetic data, the solver improves the outer loss and moves the reward parameter near its theoretical optimum.

5 SOLVING BI-O IN GENERAL

  • Direct outer optimization is more general than the tractable analysis, and implicit differentiation avoids unrolled-gradient memory costs and gradient-free dimensionality limits.The paper focuses on implicit differentiation because explicit methods can require storing long computational graphs, while gradient-free methods face the curse of dimensionality.Source: §5 · p. 7

5.1 IMPLICIT DIFFERENTIATION

  • Implicit differentiation treats the inner solution as a function of the reward parameter and computes its derivative through a linear system for the outer chain-rule gradient.The inner solution is characterized as a root of an operator; the derivative is obtained using the implicit function theorem.Source: §5.1 · p. 7
  • The operator is the inner-objective gradient scaled by the learning rate, so the method assumes inner optimization reaches a local minimum where that gradient vanishes.The paper describes convergence to such a local minimum as a plausible hypothesis for modern stochastic optimizers.Source: §5.1 · p. 7

5.2 EMPIRICAL VALIDATION

  • The synthetic-data solver improves outer NLL and brings the reward parameter near its theoretical optimum, with a roughly convex loss landscape whose minimum is nearby.The experiment uses 100 outer iterations and 50 inner iterations; the parameter starts at 1, corresponding to Euclidean distance.Source: §5.2 · pp. 7–8, Figure 3 · p. 8

6 APPLICATIONS

The applications use theoretical and implicit-differentiation reward parametrizations in policy-gradient training, evaluating them on tabular classification and model-based reinforcement learning. Results favor optimal-reward policy gradients on tabular tasks and show favorable NLL and MSE outcomes in reinforcement learning, while downstream policy returns remain untested.

6 APPLICATIONS

  • The proposed applications estimate the optimal reward from empirical covariance or use an implicit-differentiation bilevel solver, which avoids dependence on covariance-estimation accuracy and the theoretical assumptions.The methods are benchmarked against vanilla PG and NLL losses, with the identity reward serving as the squared-distance baseline.Source: §6 · p. 8

6.1 TABULAR CLASSIFICATION

  • Tabular experiments use multiclass logistic regression with PG on Poker and imbalanced Credit default; AUC supplements accuracy for the imbalanced task.The reward uses distances between one-hot labels and generations sampled from the model's softmax distribution.Source: §6.1 · pp. 8–9
  • On imbalanced Credit default, accuracy is high across methods but AUC distinguishes the optimal-reward PG variant; on Poker, that variant performs best despite the task challenging all methods.The tables report accuracy and Credit default AUC; no specific metric values are supplied in these passages.Source: Table 1 · p. 8, §6.1 · p. 9, Table 2 · p. 9

6.2 MODEL-BASED REINFORCEMENT LEARNING

  • The reinforcement-learning evaluation fits Gaussian transition models on three D4RL HalfCheetah datasets collected by simple, medium, and expert policies.Models train for 400 epochs with Adam at learning rate 10^-3 and λ = 1.Source: §6.2 · p. 9
  • In HalfCheetah model learning, optimal-reward PG improves NLL over identity-reward PG, while the implicit-solver reward achieves second-best NLL and optimal MSE.NLL-optimized models have the strongest NLL; MSE is relevant when using deterministic planners.Source: §6.2 · p. 9

7 CONCLUSION

The paper frames reward learning as a bilevel route to MLE-optimal models and reports validation in practical applications. It also identifies restrictive reward parameterization and limited large-scale testing as boundaries on the current evidence.

7 CONCLUSION

  • On the MBRL experiment, optimal-reward PG ranks second to NLL on NLL and first on MSE.The comparison is reported for Table 3.Source: Table 3 · p. 9
  • The paper introduces a bilevel framework and closed-form solutions under specific reward-model and data-distribution assumptions, then validates the approach in practical applications.The authors describe the framework as learning rewards for policy gradient that produce models optimal in the MLE sense with respect to observed data.Source: §7 · p. 9
  • The authors conclude that an optimal reward choice lets the framework make more effective use of policy-gradient methods.Source: §7 · p. 10
  • The reward parameterization may restrict flexibility, and further large-scale experiments are needed to assess generalization to complex applications.The reported experiments focus predominantly on tabular data; future directions include LLM fine-tuning, structured prediction, and time-series forecasting.Source: §7 · p. 10

REPRODUCIBILITY STATEMENT

The appendix records code and experiment details, establishes a positive-semidefinite lemma, and proves that the inner-level optimization has a unique solution under its stated assumptions.

REPRODUCIBILITY STATEMENT

A.1.1 A USEFUL LEMMA

A.1.2 AN INTERMEDIATE PROPOSITION: SOLUTION OF THE INNER-LEVEL PROBLEM

  • Under Assumptions 4.1 and 4.2, the inner-level problem has exactly one solution, established by deriving its objective and gradient and showing strict concavity.Source: §A.1.2 · pp. 20, 23

A.1.3 CONCLUDING THE PROOF: SOLUTION OF THE OUTER-LEVEL PROBLEM

The proof establishes convexity and uniqueness for the outer-level solution, then treats the isotropic case through concavity and first-order optimality conditions.

A.1.3 CONCLUDING THE PROOF: SOLUTION OF THE OUTER-LEVEL PROBLEM

  • The outer-level problem's objective is convex in U, and the proof establishes that its solution is unique.Source: §A.1.3 · p. 23

A.2 PROOF OF COROLLARY 4.1.1

  • For the isotropic case B = σ^2In, the solution set is characterized by showing concavity in (A, σ^2) and solving the first-order conditions.Substitution into the DKL expression yields Fλ,Σ.Source: §A.2 · p. 25

A.3 PROOF OF COROLLARY 4.1.2

The appendix establishes that maximizing the objective is equivalent to minimizing the reverse KL divergence. It also shows that a bijection between parameter spaces preserves the corresponding optimization solutions.

A.3 PROOF OF COROLLARY 4.1.2

  • Maximizing J(θ) is equivalent to minimizing the reverse KL divergence.Source: §A.3 · p. 26

A.4 ON THE DEFINITION OF H IN ASSUMPTION 4.2

  • The appendix justifies reparameterizing the search space by S++ through a bijection.Source: §A.4 · p. 27
  • The two optimization problems are equivalent: solutions map in both directions through the bijection ϕ.Source: §A.4 · pp. 27–28

B ADDITIONAL EXPERIMENTS

Figure 4 compares distributions for different values of λ.

B.1 DISTRIBUTION COMPARISON

  • Figure 4 compares distributions across different values of λ.Source: Figure 4 · p. 29

Related papers

Loading 2510.07624v3…