Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Beyond Imitation: Recovering Dense Rewards from Demonstrations

Jiangnan Li, Thuy-Trang Vu, Ehsan Abbasnejad, Gholamreza Haffari

TL;DR

This work reframes supervised fine-tuning (SFT) of large language models as inverse reinforcement learning (IRL) on a token-level MDP with discount $\gamma=1$, showing that SFT implicitly learns a dense token-level reward via the IQ-Learn reduction. It then derives a practical approach, Dense-Path REINFORCE (DPR), that extracts a baseline-relative dense reward from SFT logits and uses undiscounted REINFORCE to improve the policy, achieving gains over SFT on multiple instruction-following benchmarks. The authors provide theoretical results, including a dual contraction bound linking reward error to policy error and potential-based shaping to connect the SFT logit to a true reward, supported by extensive experiments across four backbones and four benchmarks. The findings demonstrate that density of credit assignment can be leveraged in LfD, offering a scalable and stable path to enhance instruction-following performance by exploiting the hidden reward structure embedded in demonstration data. Overall, the paper reframes SFT as a reward-learning process and offers a practical, empirically effective method for exploiting dense reward signals to surpass imitation-based baselines.

Abstract

Conventionally, supervised fine-tuning (SFT) is treated as a simple imitation learning process that only trains a policy to imitate expert behavior on demonstration datasets. In this work, we challenge this view by establishing a fundamental equivalence between SFT and Inverse Reinforcement Learning. We prove that the SFT objective is a special case of Inverse Q-Learning, which implies that the SFT process does not just learn a policy, but also an implicit, dense, token-level reward model that explains the expert demonstrations. We then show how to recover this dense reward signal directly from the SFT model by formulating a baseline-relative reward function. The availability of such a dense reward model offers numerous benefits, providing granular credit assignment for each token generated. We demonstrate one key application by using these recovered rewards to further improve the policy with reinforcement learning. Our method, Dense-Path REINFORCE, consistently outperforms the original SFT models on instruction-following benchmarks. This work reframes SFT not merely as policy imitation but as a powerful reward learning mechanism, opening new possibilities for leveraging expert demonstrations.

Beyond Imitation: Recovering Dense Rewards from Demonstrations

TL;DR

This work reframes supervised fine-tuning (SFT) of large language models as inverse reinforcement learning (IRL) on a token-level MDP with discount , showing that SFT implicitly learns a dense token-level reward via the IQ-Learn reduction. It then derives a practical approach, Dense-Path REINFORCE (DPR), that extracts a baseline-relative dense reward from SFT logits and uses undiscounted REINFORCE to improve the policy, achieving gains over SFT on multiple instruction-following benchmarks. The authors provide theoretical results, including a dual contraction bound linking reward error to policy error and potential-based shaping to connect the SFT logit to a true reward, supported by extensive experiments across four backbones and four benchmarks. The findings demonstrate that density of credit assignment can be leveraged in LfD, offering a scalable and stable path to enhance instruction-following performance by exploiting the hidden reward structure embedded in demonstration data. Overall, the paper reframes SFT as a reward-learning process and offers a practical, empirically effective method for exploiting dense reward signals to surpass imitation-based baselines.

Abstract

Conventionally, supervised fine-tuning (SFT) is treated as a simple imitation learning process that only trains a policy to imitate expert behavior on demonstration datasets. In this work, we challenge this view by establishing a fundamental equivalence between SFT and Inverse Reinforcement Learning. We prove that the SFT objective is a special case of Inverse Q-Learning, which implies that the SFT process does not just learn a policy, but also an implicit, dense, token-level reward model that explains the expert demonstrations. We then show how to recover this dense reward signal directly from the SFT model by formulating a baseline-relative reward function. The availability of such a dense reward model offers numerous benefits, providing granular credit assignment for each token generated. We demonstrate one key application by using these recovered rewards to further improve the policy with reinforcement learning. Our method, Dense-Path REINFORCE, consistently outperforms the original SFT models on instruction-following benchmarks. This work reframes SFT not merely as policy imitation but as a powerful reward learning mechanism, opening new possibilities for leveraging expert demonstrations.

Paper Structure

This paper contains 44 sections, 2 theorems, 49 equations, 4 figures, 3 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Preliminaries
  4. Problem setup and notation.
  5. Soft-optimality equations.
  6. MaxEnt IRL in occupancy space.
  7. IQ-Learn: a $Q$-space reduction.
  8. Methodology
  9. Step 1: SFT is equivalent to a special case of IQ-Learn
  10. Step 2: Reward error is controlled by policy error (IRL view)
  11. Step 3: From an SFT-derived dense reward to policy improvement
  12. Experiments
  13. Experimental setup
  14. Data.
  15. Backbones (pretrained only).
  16. ...and 29 more sections

Key Result

Proposition 1

On the token MDP with discount rate $\gamma=1$, maximizing $J^*(Q)$ is equivalent to minimizing the token-level SFT loss $\mathcal{L}_{\mathrm{SFT}}(\theta)=\mathbb{E}_{(s,a)\sim\rho_E}\!\left[-\log\pi_\theta(a\mid s)\right]$, where $\pi_\theta(a\mid s)\propto \exp Q_\theta(s,a)$.

Figures (4)

  • Figure 1: Credit assignment in Dense-Path REINFORCE (Best viewed in color). We provide two answers to a math question. The left is the correct response, and on the right is our modified response. Each token is colored according to the baseline-relative dense reward as expressed in Eq. \ref{['eq:baseline-reward']} (darker red means higher reward), using the trained SFT model and SFT checkpoint. We see that the model correctly identifies the erroneous number, without much change to the reward value of the other tokens, which indicates the ability to do credit assignment.
  • Figure 2: The effect of reward discount-rate ($\gamma\in\{0.93,0.94,...,1.00\}$) across four backbones. Performance (win rate vs. SFT, higher is better) peaks at the undiscounted setting $\gamma{=}1.0$. This is consistent with our analysis: (i) the SFT$\leftrightarrow$IQ‑Learn equivalence is derived for $\gamma{=}1$; (ii) with discounting, early tokens are over‑rewarded relative to later ones, weakening token‑level credit assignment.
  • Figure 3: Baseline checkpoint selection. We vary the baseline $\pi_{\mathrm{ref}}$ along the SFT training trajectory (x‑axis: SFT progress), keeping all else fixed. A baseline trained with around half of the total training samples yields the best results. Intuitively, an early baseline is too weak, over‑inflating rewards and increasing variance; a late baseline is too close to the full SFT, shrinking $\log\pi_{\mathrm{SFT}}{-}\log\pi_{\mathrm{ref}}$ and reducing signal‑to‑noise. The midpoint balances magnitude and discriminativeness, consistent with our bound in Appendix \ref{['app:baseline-tightness']}.
  • Figure :

Theorems & Definitions (2)

  • Proposition 1: SFT $\equiv$ IQ-Learn with a linear conjugate
  • Theorem 2: Dual contraction: reward error $\le$ policy error