Offline DiffusionNFT as KL-Regularized Reward Tilting

Setup (From DiffusionNFT)

Introduce a binary optimality variable $o\in{0,1}$ and a prompt/context $c$. DiffusionNFT defines the reward as an optimality probability

Given a reference (old) model $\pi_{\mathrm{old}}(x_0\mid c)$, define the positive (optimality-conditioned) distribution

and the corresponding forward-noised marginals at diffusion time $t$ under a fixed kernel $\pi_{t\mid 0}(x_t\mid x_0)$:

DiffusionNFT’s training objective yields an optimal velocity field of the form

with guidance direction $\Delta(x_t,c,t)=\alpha(x_t,c)\bigl(v^{+}(x_t,c,t)-v^{\mathrm{old}}(x_t,c,t)\bigr)$.

And the mixture coefficient is defined as

Derivation of $\alpha(x_t,c)=p(o=1\mid x_t,c)$

Apply the forward diffusion to the positive distribution:

where we use the identity

we obtain

Note that $\alpha(x_t,c)=\mathbb{E}_{x_0\sim \pi _{\mathrm{old}}(\cdot \mid c)}[r(x_0,c)\mid x_t]$. This shows that the mixture coefficient equals the optimality posterior distribution at the noisy state $x_t$ under the old model.

Remark

This equality is a direct consequence of the definition of $\pi_{+}$ as an optimality-conditioned distribution and Bayes’ rule under the fixed forward noising kernel. It may be omitted or only implicit in the original DiffusionNFT paper.

DiffusionNFT optimal distribution at each step

In order to compute the optimal distribution, we need to simplify the residual term $\Delta(x_t,c,t)=\alpha(x_t,c)\bigl(v^{+}(x_t,c,t)-v^{\mathrm{old}}(x_t,c,t)\bigr)$.

Using the Bayes relation for the positive marginal

and i) utilizing the relation between velocity fields and score functions under fixed Gaussian noising, and ii) applying log and gradient, one has

Closed-form optimal distribution (density-level)

Since $\alpha(x_t,c)=p(o=1\mid x_t,c)$ and $\alpha\nabla\log\alpha=\nabla\alpha$, we have:

This implies the explicit density:

At $t=0$ this reduces to $p^{*}(x_0\mid c)\propto p^{\mathrm{old}}(x_0\mid c)\exp(\frac{2}{\beta}r(x_0,c))$.

Thus, DiffusionNFT with fixed $p^{\mathrm{old}}=p^{\mathrm{ref}}$ can be viewed as learning a KL-regularized exponential tilt of the reference model, where the “reward” is the optimality posterior at the noisy state.

Online DiffusionNFT Leads to Reward Hacking via Accumulated Tilting

In practice, DiffusionNFT can be run online: after fitting $v_\theta$ for one epoch, the old model is updated (e.g., by copying weights (hard) or by EMA (soft)) and the next epoch is trained against this updated old model. This induces a recursion over the initial reference distributions.

Idealized recursion with exact per-epoch optima

Let $p^{(k)}_t(x_t\mid c)$ denote the old marginal at epoch $k$ (the distribution induced by the current old velocity field). Treat the reward/optimality posterior $\alpha_t(x_t,c)=p(o=1\mid x_t,c)$ as fixed across epoches.

From the closed-form optimum, the population-optimal update at epoch $k$ is

Unrolling the recursion yields

As $K\to\infty$, the previous expression concentrates mass on the set of maximizers of the reward.

Remarks on EMA references

If the reference is updated via EMA in parameter space, the induced distribution recursion is not exactly the simple multiplicative update above. Nevertheless, EMA typically acts as a trust-region mechanism that interpolates between keeping $p^{(k)}_t$ fixed and fully replacing it by the latest student, effectively reducing the rate at which the tilt coefficient grows with $k$. The qualitative conclusion remains: in the absence of an anchoring term to the initial model, repeated online improvement tends to accumulate reward tilt and can become increasingly peaked around high-reward regions and thus prone to reward hacking.

Online DiffusionNFT with an Additional KL to the Initial Reference

To prevent unbounded drift from the original model, one can augment the per-epoch objective with an additional regularizer that penalizes deviation from the initial reference distribution $p^{(0)}_t(\cdot\mid c)=p^{\mathrm{ref}}_t(\cdot\mid c)$.

At the distribution level, consider the KL-regularized problem at epoch $k$:

Here $\eta_1$ controls the trust region to the current reference and $\eta_0$ controls anchoring to the initial reference.

The unique optimizer has the closed form

with weights and effective tilt coefficient

Thus, the per-epoch solution is an exponential tilt of a geometric mixture of the current and initial references.

It is straightforward to see that the iterates converge to a limiting distribution:

Notably, the limiting distribution depends only on the anchoring strength $\eta_0$; the trust-region parameter $\eta_1$ affects only the dynamics (convergence rate and stability) through $w$.

References