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

1. Problem Statement

Modern reasoning agents are increasingly evaluated on their ability to generate multiple valid solution paths (e.g., diverse chain-of-thoughts, plans, or tool-use traces) rather than a single correct answer. Standard Reinforcement Learning (RL) objectives, such as Proximal Policy Optimization (PPO) or Group Relative Policy Optimization (GRPO), focus on reward maximization. While effective for single-answer accuracy, these methods tend to suffer from mode collapse, where the policy concentrates probability mass on a single high-reward trajectory, eliminating the diversity of valid solutions required for robust evaluation (e.g., pass@k metrics).

Distribution-matching RL offers an alternative by training the policy to fit a target distribution where trajectory probability is proportional to the exponentiated reward:
$$\tilde{\pi}(o | q) = \frac{1}{Z(q)} \pi_{\text{ref}}(o | q) \exp(\beta \tilde{r}(q, o))$$
where $Z(q)$ is the partition function (a normalization constant dependent on the prompt $q$).

The core bottleneck in existing distribution-matching methods (e.g., FlowRL, GFlowNets) is the online coupling of the policy ($\pi_\theta$) and the partition function estimator ($Z_\phi$). These methods co-train both components on shared online rollouts using a joint loss (e.g., Trajectory Balance). This coupling creates two critical issues:

1. Error Propagation: Calibration errors in the partition function estimator $Z_\phi$ flow directly into the policy gradient, distorting updates.
2. Lack of Diagnostics: There is no independent quality signal for the partition function; its performance is entangled with the policy's loss, making it impossible to diagnose whether failures stem from the policy or the partition function estimation.

2. Methodology: Decoupled Importance-Sampled Anchoring (DISA)

The authors propose DISA, a three-stage pipeline that decouples the estimation of the partition function from the policy optimization loop. The key insight is that the partition function $Z(q)$ depends only on the reference policy $\pi_{\text{ref}}$, the reward function, and the inverse temperature $\beta$, but not on the trainable policy $\pi_\theta$. Therefore, it can be estimated offline before RL training begins.

Stage 1: Offline Importance Sampling (IS) Estimation

For a set of offline prompts $\mathcal{D}_Z$, the method draws $N$ proposal trajectories from a strong proposal model $p_T(\cdot|q)$ (e.g., a larger teacher model).

• Importance Weights: For each trajectory $o_i$, an importance weight is calculated: $w(q, o_i) = \frac{\pi_{\text{ref}}(o_i|q)}{p_T(o_i|q)} \exp(\beta \tilde{r}(q, o_i))$.
• Log-Space Aggregation: To ensure numerical stability and reduce bias, the partition function estimate for prompt $q$ is computed in log-space using a logsumexp aggregator:
  $$\log \hat{Z}_{\text{IS}}(q) = \text{logsumexp}_i \left( \log \pi_{\text{ref}}(o_i|q) - \log p_T(o_i|q) + \beta \tilde{r}(q, o_i) \right) - \log N$$
  This provides an unbiased estimate of $Z(q)$ on the linear scale, with a controlled $O(1/N)$ bias on the log scale.

Stage 2: Amortization via Regression

The per-prompt estimates $\log \hat{Z}_{\text{IS}}(q)$ serve as labels to train a lightweight prompt-conditioned regressor $g_\psi(q)$ (e.g., a small MLP).

• The regressor is trained via least-squares regression to minimize the squared error between its prediction and the offline IS labels.
• Once trained, the regressor parameters $\psi^*$ are frozen. This creates a reusable, offline estimate of the partition function that generalizes to prompts outside the initial estimation set.

Stage 3: Anchored Online Policy Learning

During the online RL phase, the frozen regressor $g_{\psi^*}(q)$ replaces the learnable partition function $Z_\phi(q)$ in the Trajectory Balance (TB) loss.

• Objective: The policy $\pi_\theta$ is optimized to minimize the residual:
  $$\mathcal{L}_{\text{DISA}}(\theta) = \mathbb{E} \left[ \sum_i \omega_i \left( g_{\psi^*}(q) + \log \pi_\theta(o_i|q) - \log \pi_{\text{ref}}(o_i|q) - \beta \tilde{r}(q, o_i) \right)^2 \right]$$
• Decoupling: Crucially, gradients are stopped from flowing through $g_{\psi^*}$. The partition function acts as a fixed, per-prompt constant (an "anchor") during policy updates.

3. Key Contributions

1. Identification of Coupling as a Core Obstacle: The paper identifies the entanglement of partition function estimation and policy learning as the primary source of calibration errors and diagnostic blindness in distribution-matching RL.
2. The DISA Pipeline: A novel three-stage approach that moves partition function estimation offline via importance sampling, amortizes it via regression, and freezes it for online RL.
3. Theoretical Guarantees:
   • Exact Preservation: If the offline partition function is exact, the global minimizer of the DISA objective matches the theoretical reward-tilted target distribution.
   • Stationarity with Bias: Even with a prompt-only bias (where the error depends only on $q$ and not the trajectory), the target distribution remains a stationary point of the on-policy gradient flow.
   • Bounded Perturbation: Residual estimation errors enter the policy objective only through a bounded perturbation envelope (derived via Cauchy-Schwarz), ensuring stability.
4. Empirical Validation: Comprehensive evaluation across math and code benchmarks demonstrating that DISA preserves the benefits of distribution matching while avoiding the pitfalls of online coupling.

4. Experimental Results

The authors evaluated DISA on two open-weight backbones (Qwen2.5-7B and Qwen3-4B-Base) across six math benchmarks (including AIME 2024/2025, AMC 2023, MATH-500) and three code benchmarks (LiveCodeBench, CodeForces, HumanEval+).

• Mathematical Reasoning:
  • DISA matched or slightly exceeded the online-coupled baseline FlowRL (e.g., 65.5 vs. 64.8 Mean@8 on Qwen3-4B-Base).
  • It significantly outperformed reward-maximization baselines GRPO and GSPO (e.g., +16.1 points over GRPO on Qwen3-4B-Base).
• Code Generation:
  • On pass@16 (a metric sensitive to solution diversity), DISA achieved the highest scores among all methods (e.g., 51.5 on Qwen3-4B-Base), outperforming FlowRL, GSPO, and GRPO.
  • On pass@1, DISA remained competitive, showing no significant regression compared to reward-maximization methods.
• Diversity Analysis:
  • An LLM-as-judge evaluation on solution strategy diversity (using a 1-5 Likert scale) showed that DISA retained the largest fraction of the backbone's original diversity (3.72/5.0), compared to FlowRL (3.40) and significantly better than GRPO/GSPO (~3.0).
• Ablation Studies:
  • Proposal Strength: Using a weaker proposal model (Qwen3-4B vs. 235B) increased the variance of the offline labels, leading to a performance drop, but the regressor still converged, confirming the bias-variance trade-off predicted by theory.
  • Inverse Temperature ($\beta$): Sensitivity studies confirmed that under-tilting ($\beta=10$) leads to mild signal loss, while over-tilting ($\beta=20$) causes a "variance cliff" and mode collapse, validating the theoretical analysis of the importance sampling weights.
• Distillation vs. RL: When compared to fine-tuning (SFT) on the same high-reward trajectories used for DISA's offline estimation, DISA provided an additional gain of up to 13.8 Mean@8 points, demonstrating that fitting the soft, reward-tilted distribution recovers landscape information discarded by hard SFT targets.

5. Significance and Claims

The paper claims that DISA represents a significant architectural shift in post-training LLMs for reasoning. By decoupling the partition function estimation from the policy loop, DISA:

• Eliminates Calibration Errors: It prevents partition function estimation errors from distorting policy gradients.
• Enables Independent Diagnostics: The quality of the partition function can be monitored separately from the policy training via the regressor's validation loss.
• Preserves Diversity: It successfully maintains the multi-modal structure of the solution space, which is critical for complex reasoning tasks where multiple valid paths exist.
• Theoretical Soundness: The method provides formal guarantees that the offline estimation process does not alter the theoretical fixed point of the distribution-matching objective, provided the estimation error is bounded.

The authors conclude that DISA offers a robust, theoretically grounded alternative to online-coupled distribution matching, achieving state-of-the-art performance in both accuracy and diversity without the instability associated with co-training partition functions and policies.