English

DISA: Offline Importance Sampling for Distribution-Matching LLM-RL

Machine Learning 2026-05-19 v1 Computation and Language

Abstract

Modern reasoning agents are increasingly evaluated on their ability to generate multiple valid solution paths, plans, or tool-use traces for a given input. Standard reward-maximizing RL tends to collapse onto the most easily reinforced high-reward mode, whereas distribution-matching RL aims to allocate probability mass across the entire reward-shaped solution set. Achieving this objective requires computing a prompt-dependent partition function over the trajectory space. Because existing distribution-matching methods learn this partition function online alongside the policy, calibration errors in the partition function directly distort policy updates and remain impossible to diagnose independently. We introduce DISA, short for Decoupled Importance-Sampled Anchoring, which moves this calibration problem outside the RL loop. DISA draws proposal trajectories offline, estimates the partition function via importance sampling, and freezes the resulting partition-function estimate before policy optimization begins. This decoupling preserves the distribution-matching objective while strictly separating partition-function estimation from policy learning in data, gradients, loss, and diagnostics. Empirically, on two open-weight backbones across six math and three code benchmarks, DISA matches or exceeds the online-coupled distribution-matching baseline FlowRL, outperforms rewardmaximization baselines GRPO and GSPO on math averages, and exceeds LoRASFT distillation by up to 13.8 Mean@8 points on the same offline trajectories. An LLM-as-judge evaluation further shows that DISA retains substantially more strategy-level diversity than reward-maximization baselines, and sensitivity studies on the proposal strength and inverse temperature follow the bias-variance pattern predicted by the analysis.

Keywords

Cite

@article{arxiv.2605.17295,
  title  = {DISA: Offline Importance Sampling for Distribution-Matching LLM-RL},
  author = {Shaobo Wang and Yujie Chen and Yafeng Sun and Wenjie Qiu and Zhihui Xie and Sihang Li and Yucheng Li and Huiqiang Jiang and Xingzhang Ren and Xuming Hu and Dayiheng Liu and Linfeng Zhang},
  journal= {arXiv preprint arXiv:2605.17295},
  year   = {2026}
}

Comments

21 pages, 7 figures, 7 tables. Abstract shortened to respect the arXiv limit of 1920 characters. Please see the PDF for the full abstract

R2 v1 2026-07-22T07:17:08.543Z