1. One structure, two sources of mismatch

Training–inference mismatch in LLM RL and CFG mismatch in diffusion RL look like unrelated engineering problems. They share one mathematical structure: the distribution that supplies the training samples is not the reference implicit in the training gradient. The causes differ; the resulting weighted gradients decompose in the same way.

For LLMs, a Megatron trainer and a vLLM sampler can produce different token probabilities because of numerical execution, sampling transforms, or checkpoint staleness [1]. For diffusion models, CFG changes the sampling procedure on purpose, while training may still use only the conditional branch [5].

We derive the diffusion case first, then map it to LLMs. Throughout, the rollout distribution, reward weights, and estimated statistics are held fixed during differentiation.

2. From IGO to exponentially weighted flow matching

Take the distribution-first view of IGO. Let $Q(X\mid c)$ be the current CFG rollout distribution for input $c$ and $A(X,c)$ an advantage. A local improvement trades reward against staying close to $Q$:

Setting the functional derivative to zero, $A-\beta(\log(P/Q)+1)+\lambda=0$, gives the familiar exponential tilt

with $Z(c)=\mathbb E_{Q}[W]$ assumed finite. This exponential weighting comes from this particular KL-regularized objective; it is the same tilt as in AWR [3], not a definition of all IGO algorithms [2].

To fit $Q^\star$ with flow matching, draw $X\sim Q$ and independent $\epsilon\sim\mathcal N(0,I)$, and set

so $t=1$ is noise, $t=0$ is data, and $u$ is the denoising velocity. The exponentially weighted objective is

For fixed $c$, weighting by $W$ turns the endpoint measure $Q$ into $Q^\star$ up to the constant $Z(c)$, so the ideal regression target is the flow of $Q^\star$. (FM is a field regression, not literally a log-likelihood or forward KL; across inputs, $Z(c)$ also sets their relative training weight [4].)

3. Two forces inside one gradient

Fix a noisy state $z$. Two conditional averages matter: the rollout mean field and the mean weight,

Writing $G=G_\theta(z)$, the output gradient is $\mathbb E_Q[W(G-u)\mid z]=m(z)G-\mathbb E_Q[Wu\mid z]$. Splitting $\mathbb E_Q[Wu\mid z]=m(z)F_Q(z)+\operatorname{Cov}_Q(W,u\mid z)$ gives

The orange term fits the guided rollouts

Gradient descent on the orange term moves $G$ toward $F_Q$. It survives when all weights equal one, so it expresses no preference between samples at all. At the pre-update field $G_0$ it vanishes only where $G_0=F_Q$, which is exactly what CFG breaks.

With CFG rollouts and conditional-only training, this term transfers guided sampling behaviour into the conditional branch. Sampling with CFG afterwards may amplify that behaviour again.

Note that the relevant reference is $F_Q$, not automatically the cached CFG velocity: re-noising endpoints and following the actual CFG trajectory define different joint distributions [5].

The green term is the reward-induced flow change

Under exponential reweighting, $F_{Q^\star}(z)=\mathbb E_Q[Wu\mid z]/m(z)$. Therefore

The covariance isolates exactly the change from the rollout flow to the reward-tilted flow. If the goal is to distil $Q^\star$ completely, the orange term is appropriate. If the goal is to apply the reward-induced change to an existing policy $G_0$, the orange term is an extra, unwanted update.

4. Two ways to remove the orange term

The fix has two parts: isolate the green covariance, then anchor the update to the frozen pre-update policy with a purple constraint. Centering plus this constraint and a fixed-reference STE give the same gradient when the conditional quantities are exact.

A. Center the weight, then anchor the update

Replace $W$ by $W-m(z)$, treating the baseline as a constant:

Because $\mathbb E_Q[W-m(z)\mid z]=0$, the coefficient of $G$ cancels and the orange term disappears:

Centering alone has no finite optimum. With fixed rollout statistics, this gradient does not depend on $G$: the objective is affine in $G$ and unbounded below. Repeated inner steps keep pushing the field even after the intended improvement has been reached. (Try the Centered only mode in the demo below.)

Let $G_0(z)$ be the policy output before the update window, held fixed throughout its inner steps. Add a constraint weighted by the same conditional mean $m(z)$:

Its gradient and curvature are

Since $W=e^{A/\beta}>0$, $m(z)>0$ and the problem is strongly convex. The purple term is zero at $G=G_0$ and restores the policy as it moves. It anchors to $G_0$, not to the CFG field $F_Q$, so it does not bring the orange term back. The pointwise optimum is

The corrected objective applies the reward-induced flow change to the pre-update policy, and nothing else. Two details matter in practice: $G_0$ is the pre-update policy, not necessarily the pretrained model, and must stay fixed through the inner window; and the coefficient is the state-conditional $m(z)$, not a batch average of $W$.

Estimating $m(z)$ with a tempered posterior

For each input $c$, let $X_1,\dots,X_K$ be the $K$ rollout endpoints generated for that input. A uniform empirical prior and the Gaussian bridge $x_t\mid X\sim\mathcal N((1-t)X,\,t^2I)$ give, for $0<t<1$,

Use the detached $\widehat m_T(z)$ both in the centered weight and as the constraint coefficient. At $T=1$ these are the exact posterior responsibilities of the empirical bank; $T>1$ inflates the assumed noise variance and smooths a concentrated posterior, at the cost of bias relative to the population conditional mean. Compute the full summed latent squared norm and use a stable log-softmax.

What an inexact baseline costs

Replacing $m$ by a fixed estimate $\widehat m$ in both places gives

A residue of the orange term survives in proportion to the baseline error. Positive regularization does not remove it; only an accurate conditional mean does. At low noise a concentrated posterior may correctly identify the source endpoint, so temperature smoothing is a regularizer, not a guarantee of a better estimate.

B. Correct the forward field with a fixed-reference STE

Instead of estimating the scalar $m(z)$, estimate the vector field $F_Q$ directly by ordinary FM on the CFG endpoints, as GuidanceNFT does with online flow matching on guided samples [5]:

With the same frozen $G_0$, define the straight-through surrogate

Its forward value is $G_\theta-G_0+F_\phi$, its Jacobian is that of $G_\theta$. With an exact reference, ordinary weighted FM through this surrogate has gradient

identical to the centered-and-anchored objective above: the purple constraint is implicit in the fixed reference. Our contribution here is this equivalence; the learned-reference correction itself follows [5].

If $G_0$ is replaced by the current $G_\theta$ inside the stop-gradient at every step, $\widetilde G_\theta=G_\theta+\operatorname{sg}(F_Q-G_\theta)$, the purple term disappears and only the green covariance remains: the refreshed STE is centering without the constraint.

Which one to use

A learned reference means maintaining an auxiliary model next to the policy; at trillion-parameter scale its weights, training, and communication are expensive, while the kernel estimator needs no extra parameters. Both still need a fixed $G_0$: cached predictions suffice at fixed training queries, fresh queries need extra reference evaluations or an explicit approximation, and the reference and student must use the same branch or CFG composition. STE itself does not require an auxiliary model; the difference is how the reference field is obtained. Learning both $F_Q$ and $F_{Q^\star}$ explicitly is a third way to build the difference target. None of these output-space identities guarantees that a shared network can realize the pointwise target exactly.

5. The same picture for LLMs

The decomposition, and both fixes, carry over to logit space. At a fixed prefix $h$, let $q$ be the sampler's probabilities, $p_\theta=\operatorname{softmax}(\ell_\theta)$ the trainer's, $e_y$ the one-hot sampled token, and $m(h)=\mathbb E_q[W\mid h]$. The logit score is $s_\theta(y)=e_y-p_\theta$ with $\mathbb E_q[s_\theta\mid h]=q-p_\theta$, so negative weighted log-likelihood has gradient

the categorical twin of $\fit{m(z)(G-F_Q)}-\rew{\operatorname{Cov}_Q(W,u\mid z)}$. The orange term pulls the trainer toward the sampler's distribution, whatever the reward says.

Baseline or score centering: two routes to the same expectation

Score centering [1] subtracts the conditional mean of the score, $s_\theta-\mathbb E_q[s_\theta\mid h]=e_y-q$. Subtracting the conditional weight baseline instead gives the same expected update:

The diffusion baseline uses the first form; LLM score centering uses the second, with sampler probabilities. They agree in expectation with exact conditional means, not necessarily as finite-batch estimators.

Straight-through, in logits

The logit-level STE is

For positive $q$ its forward softmax is $q$ and its logit Jacobian is the identity, so weighted cross-entropy yields the green term alone, like the refreshed diffusion STE. Freezing the correction at the pre-update logits $\ell_0$, $\widetilde\ell_\theta=\ell_\theta+\operatorname{sg}(\log q-\ell_0)$, additionally keeps a restoring term $\anc{m(h)(\widetilde p_\theta-q)}$ with $\widetilde p_\theta=\operatorname{softmax}(\widetilde\ell_\theta)$, the analogue of the purple constraint. Correcting only the sampled token's scalar log-probability leaves its derivative unchanged and is not enough.

The practical difference: LLM score centering can read the sampler's conditional token probabilities directly, while diffusion needs a conditional expectation over latent endpoints, either the scalar mean weight or the vector mean field. Kernel estimation and auxiliary-model training are two approximations to these. In both domains, centering removes the identified mismatch term; it does not in general recover an exact on-policy gradient under the trainer's own distribution.

6. What we see at CFG = 5

We compare NFT [6] with the corrected objective at guidance scale 5, using $K=8$ rollouts per input. The corrected run keeps rising over its logged range; NFT improves at first and then gives back part of the gain.

Figure 1. Reward versus training step at CFG = 5, $K=8$ rollouts per input. Blue: corrected objective; brown: NFT [6]. Values are approximate readings of a single run per method, shown up to the last logged step (32 and 41) without smoothing or extrapolation. Hover over a point to see its value.

This is one run per method, so it motivates a controlled ablation rather than establishing the mismatch correction as the cause. Removing the orange term also does not by itself guarantee visual fidelity, prevent reward-induced oversaturation, or make reward improvement monotonic.

Citation

If you find this note useful, please cite it as:

@misc{yang2026tim,
  title        = {Training--Inference Mismatch in Generative {RL}},
  author       = {Yang, Zeshi and Wang, Yujie},
  year         = {2026},
  month        = oct,
  howpublished = {\url{https://zeshiyang.github.io/blog/training-inference-mismatch.html}},
  note         = {Blog post}
}

References

The cross-domain correspondence and the equivalence identities are derived in this note; the learned-reference correction for diffusion follows [5]. The experiment is separate from the published results below.

1. Marek & Ryabinin. Score Centering Stabilizes Off-policy Reinforcement Learning. 2026.
   https://arxiv.org/html/2609.20807v1#S4

2. Ollivier, Arnold, Auger & Hansen. Information-Geometric Optimization Algorithms: A Unifying Picture via Invariance Principles. 2017.
   https://www.jmlr.org/papers/v18/14-467.html

3. Peng, Kumar, Zhang & Levine. Advantage-Weighted Regression: Simple and Scalable Off-Policy Reinforcement Learning. 2019.
   https://arxiv.org/abs/1910.00177

4. Lipman et al. Flow Matching for Generative Modeling. 2022.
   https://arxiv.org/abs/2210.02747

5. Liu, Li, Qian & Cheng. GuidanceNFT: NFT-Based Reward Fine-Tuning of Guided Diffusion Models. 2026.
   https://arxiv.org/html/2610.15856v1#S3

6. Zheng et al. DiffusionNFT: Online Diffusion Reinforcement with Forward Process. 2025.
   https://arxiv.org/abs/2509.16117