Proximal WDCE in Diffusion Neural Samplers

• Proximal WDCE is a reverse-KL training objective that decomposes diffusion sampling into tractable proximal subproblems using weighted denoising losses.
• It formulates learning as a stochastic optimal control problem on path measures, tempering updates via importance weights to enhance mode coverage.
• The framework supports both continuous and discrete settings by integrating bridge-based and masked-token strategies to balance robustness and convergence.

Proximal Weighted Denoising Cross-Entropy (WDCE) is a denoising-style reverse-KL training objective for diffusion-based neural samplers when learning is organized as a proximal point method on path measures. In "Proximal Diffusion Neural Sampler" (Guo et al., 4 Oct 2025), each proximal subproblem is realized by a weighted denoising loss: in continuous state space it is a bridge-drift matching objective, and in discrete masked diffusion it is a weighted masked-token cross-entropy. The same acronym WDCE also appears in "Robust Diffusion Models via Divergence-Induced Weighted Denoising" (Li et al., 21 Jun 2026), where it denotes the KL, or cross-entropy, specialization of a local divergence-induced denoising framework; there the loss reduces to the standard weighted MSE term and the paper explicitly states that it does not introduce proximal operators or a proximal-point method. The term therefore spans two closely related but distinct constructions: a path-space proximal training scheme and a KL denoising objective inside a broader divergence family.

1. Definition and scope

In the PDNS framework, the task of learning a diffusion-based neural sampler for an unnormalized target distribution is formulated as a stochastic optimal control problem on path measures, and the learning process is decomposed into a sequence of simpler subproblems by a proximal point method (Guo et al., 4 Oct 2025). Proximal WDCE is the practical realization of each such subproblem. Its defining features are the same across continuous and discrete settings: an outer importance weight inherited from the proximal interpolation in path space, a denoising target defined by a bridge or masked noising mechanism, and a reverse-KL interpretation in which pathwise log-likelihoods are replaced by denoising terms.

The terminology is precise. “Weighted” refers to importance weights derived from Radon–Nikodym derivatives and tempered by the proximal step size. “Denoising” refers either to matching the conditional score of a reference bridge in continuous time or to predicting masked coordinates in discrete diffusion. “Cross-entropy” refers to reverse-KL projection onto a target path measure or its proximal surrogate.

A separate but related usage arises in divergence-induced weighted denoising for DDPMs. There, WDCE is the KL instance of a local $f$-divergence construction, with generator

$$G_{\mathrm{KL}}(\delta) = (1+\delta)\log(1+\delta)-\delta,$$

local divergence $h_{\mathrm{KL}}(m)=m$, derivative $h'_{\mathrm{KL}}(m)=1$, and global objective

$$L_{\mathrm{WDCE}}(\theta)=\sum_{t=1}^T \mathbb{E}\!\left[w_t\|\epsilon-\epsilon_\theta(x_t,t)\|^2\right].$$

In that framework, WDCE is simply the KL specialization and has constant residual-space influence (Li et al., 21 Jun 2026).

2. Path-space proximal construction

The path-space formulation begins with a target density or mass function

$\pi(x)\propto e^{-\beta V(x)},$

with inverse temperature $\beta>0$, energy $V$, and path measure $P$ on trajectories $X=(X_t)_{t\in[0,T]}$ (Guo et al., 4 Oct 2025). If the reference path measure is memoryless and has terminal marginal $$G_{\mathrm{KL}}(\delta) = (1+\delta)\log(1+\delta)-\delta,$$0, the terminal reward is

$$G_{\mathrm{KL}}(\delta) = (1+\delta)\log(1+\delta)-\delta,$$1

and the optimal path measure satisfies

$$G_{\mathrm{KL}}(\delta) = (1+\delta)\log(1+\delta)-\delta,$$2

Under memorylessness, $$G_{\mathrm{KL}}(\delta) = (1+\delta)\log(1+\delta)-\delta,$$3.

PDNS applies a proximal step with step size $$G_{\mathrm{KL}}(\delta) = (1+\delta)\log(1+\delta)-\delta,$$4:

$$G_{\mathrm{KL}}(\delta) = (1+\delta)\log(1+\delta)-\delta,$$5

The solution is a geometric interpolation in path space,

$$G_{\mathrm{KL}}(\delta) = (1+\delta)\log(1+\delta)-\delta,$$6

with density ratio

$$G_{\mathrm{KL}}(\delta) = (1+\delta)\log(1+\delta)-\delta,$$7

If one initializes with $$G_{\mathrm{KL}}(\delta) = (1+\delta)\log(1+\delta)-\delta,$$8 and solves each subproblem to optimality, then

$$G_{\mathrm{KL}}(\delta) = (1+\delta)\log(1+\delta)-\delta,$$9

and the terminal marginal obeys

$h_{\mathrm{KL}}(m)=m$0

This gives a staged path from the reference terminal law $h_{\mathrm{KL}}(m)=m$1 to the target $h_{\mathrm{KL}}(m)=m$2, while keeping successive updates local in path space.

The significance of this construction is operational rather than purely formal. Large barriers between modes can cause one-shot training to collapse onto incomplete mode sets. The proximal interpolation tempers each update, producing a progressively refined path to the desired distribution and promoting exploration across modes (Guo et al., 4 Oct 2025).

3. Continuous and discrete proximal WDCE objectives

In continuous state space, the controlled dynamics are

$h_{\mathrm{KL}}(m)=m$3

with base drift $h_{\mathrm{KL}}(m)=m$4, noise scale $h_{\mathrm{KL}}(m)=m$5, and control $h_{\mathrm{KL}}(m)=m$6 (Guo et al., 4 Oct 2025). The practical proximal WDCE objective replaces pathwise log-likelihood with a denoising or bridge-matching term that depends only on $h_{\mathrm{KL}}(m)=m$7:

$h_{\mathrm{KL}}(m)=m$8

Using importance sampling from the previous iterate,

$h_{\mathrm{KL}}(m)=m$9

Here the outer factor is the proximal importance weight, the denoising target is the conditional score of the reference bridge, and the reverse-KL interpretation is retained because the squared error replaces the pathwise $h'_{\mathrm{KL}}(m)=1$0 term up to additive constants.

For discrete masked diffusion, the reference process is a continuous-time Markov chain whose masked generator replaces masked entries uniformly, while the model generator uses a categorical distribution $h'_{\mathrm{KL}}(m)=1$1 over replacements (Guo et al., 4 Oct 2025). The proximal WDCE becomes a weighted denoising cross-entropy:

$h'_{\mathrm{KL}}(m)=1$2

with

$h'_{\mathrm{KL}}(m)=1$3

The weighting acts on each terminal sample and its masked views; the denoising term is the masked-prediction cross-entropy; and the cross-entropy meaning is again a reverse-KL projection onto the proximal target.

The Radon–Nikodym factors are explicit. In continuous time they are given by Girsanov,

$h'_{\mathrm{KL}}(m)=1$4

whereas in discrete masked diffusion they are computed from jumpwise log-ratios between the reference and learned generators (Guo et al., 4 Oct 2025).

4. WDCE as KL denoising and robustness comparisons

The non-proximal WDCE formulation in robust diffusion training begins from the DDPM reverse Gaussian kernel structure and a local divergence construction (Li et al., 21 Jun 2026). With per-step residual

$h'_{\mathrm{KL}}(m)=1$5

and mismatch

$h'_{\mathrm{KL}}(m)=1$6

the per-step likelihood ratio is lognormal, $h'_{\mathrm{KL}}(m)=1$7, so the conditional $h'_{\mathrm{KL}}(m)=1$8-divergence becomes a scalar function $h'_{\mathrm{KL}}(m)=1$9. Summing local divergences gives

$$L_{\mathrm{WDCE}}(\theta)=\sum_{t=1}^T \mathbb{E}\!\left[w_t\|\epsilon-\epsilon_\theta(x_t,t)\|^2\right].$$0

and the gradient takes the form

$$L_{\mathrm{WDCE}}(\theta)=\sum_{t=1}^T \mathbb{E}\!\left[w_t\|\epsilon-\epsilon_\theta(x_t,t)\|^2\right].$$1

This yields divergence-induced weighted denoising in residual space, with influence weight

$$L_{\mathrm{WDCE}}(\theta)=\sum_{t=1}^T \mathbb{E}\!\left[w_t\|\epsilon-\epsilon_\theta(x_t,t)\|^2\right].$$2

For the KL generator, WDCE is exactly the standard weighted denoising objective:

$$L_{\mathrm{WDCE}}(\theta)=\sum_{t=1}^T \mathbb{E}\!\left[w_t\|\epsilon-\epsilon_\theta(x_t,t)\|^2\right].$$3

so the residual-space weight is constant:

$$L_{\mathrm{WDCE}}(\theta)=\sum_{t=1}^T \mathbb{E}\!\left[w_t\|\epsilon-\epsilon_\theta(x_t,t)\|^2\right].$$4

The interpretation given in the paper is that KL or WDCE assigns constant influence to all residuals regardless of size. Its effective score

$$L_{\mathrm{WDCE}}(\theta)=\sum_{t=1}^T \mathbb{E}\!\left[w_t\|\epsilon-\epsilon_\theta(x_t,t)\|^2\right].$$5

is unbounded, so large-error samples contribute disproportionately and training is non-robust under contamination (Li et al., 21 Jun 2026).

The same paper contrasts WDCE with bounded-influence divergences. For Hellinger,

$$L_{\mathrm{WDCE}}(\theta)=\sum_{t=1}^T \mathbb{E}\!\left[w_t\|\epsilon-\epsilon_\theta(x_t,t)\|^2\right].$$6

and for negative exponential divergence,

$$L_{\mathrm{WDCE}}(\theta)=\sum_{t=1}^T \mathbb{E}\!\left[w_t\|\epsilon-\epsilon_\theta(x_t,t)\|^2\right].$$7

with asymptotic decay slightly more aggressive than Hellinger because of an extra $$L_{\mathrm{WDCE}}(\theta)=\sum_{t=1}^T \mathbb{E}\!\left[w_t\|\epsilon-\epsilon_\theta(x_t,t)\|^2\right].$$8 factor. The resulting residual-space scores are bounded and redescending, making them analogous to robust M-estimation. The paper also states that all divergences agree to first order near convergence,

$$L_{\mathrm{WDCE}}(\theta)=\sum_{t=1}^T \mathbb{E}\!\left[w_t\|\epsilon-\epsilon_\theta(x_t,t)\|^2\right].$$9

so WDCE is adequate and simplest to implement on clean data and in small-mismatch regimes (Li et al., 21 Jun 2026).

This produces an important distinction. In PDNS, “proximal” refers to the path-space proximal point step and the weighting is trajectory-level importance tempering. In robust DDPM training, WDCE is the KL member of a residual-space reweighting family and the paper states that no proximal mapping is involved.

5. Scheduling, algorithmic realization, and empirical behavior

The proximal step size controls how strongly the update is tempered. PDNS defines

$\pi(x)\propto e^{-\beta V(x)},$0

so small $\pi(x)\propto e^{-\beta V(x)},$1 gives small $\pi(x)\propto e^{-\beta V(x)},$2 and strongly tempers weights, while large $\pi(x)\propto e^{-\beta V(x)},$3 weakens the proximal constraint and approaches non-proximal WDCE (Guo et al., 4 Oct 2025). Equivalently, the geometric schedule uses

$\pi(x)\propto e^{-\beta V(x)},$4

which gradually decreases to $\pi(x)\propto e^{-\beta V(x)},$5 and moves the terminal distribution from $\pi(x)\propto e^{-\beta V(x)},$6 toward $\pi(x)\propto e^{-\beta V(x)},$7. An adaptive schedule chooses $\pi(x)\propto e^{-\beta V(x)},$8, or $\pi(x)\propto e^{-\beta V(x)},$9, to satisfy a trust-region bound

$\beta>0$0

estimated from buffer weights.

The algorithmic implementation is buffer-based. At stage $\beta>0$1, one rolls out trajectories from $\beta>0$2 or $\beta>0$3, computes the proximal weights from the Radon–Nikodym derivatives and terminal reward, and stores only terminal pairs $\beta>0$4 in a replay buffer (Guo et al., 4 Oct 2025). Training then proceeds in either of two equivalent ways: a weight-based implementation that minimizes a weighted denoising loss on buffer samples, or a resampling-based implementation that resamples terminal states proportionally to weights and then uses an unweighted denoising loss. For SDEs, simulation cost per trajectory is $\beta>0$5, RN derivatives are accumulated online with negligible extra memory, and WDCE stores only terminal pairs rather than full trajectories.

Empirically, PDNS with proximal WDCE attains best or competitive results on 5 out of 7 continuous benchmarks, with notable gains on heavy-tailed MoS (50D) and rugged Lennard-Jones systems (LJ-13, LJ-55) (Guo et al., 4 Oct 2025). In discrete tasks, it matches the ground-truth 2-point correlations and magnetization on Ising and Potts near or below critical temperatures, outperforming learning baselines and matching SW while MH degrades on Potts. Ablations over fixed $\beta>0$6 show that large $\beta>0$7 collapses to incomplete modes quickly, whereas small $\beta>0$8 preserves coverage; the adaptive scheduler improves stability and convergence by increasing $\beta>0$9 over stages.

The robust DDPM study gives a complementary empirical picture for the non-proximal KL version. On CIFAR-10 under contamination $V$0, KL or WDCE yields $V$1, Hellinger gives $V$2, and NED gives $V$3 (Li et al., 21 Jun 2026). In a single-seed comparison against standard robust losses at the same contamination level, Huber $V$4 gives $V$5, clipped MSE $V$6 gives $V$7, HD gives $V$8, and NED gives $V$9. The practical overhead is also characterized: HD adds negligible cost, while NED uses per-sample quadrature with fully vectorized overhead below $P$0 of U-Net forward–backward time on the reported setup.

The limitations are correspondingly different. In PDNS, the choice of proximal step size matters: too small slows convergence, too large risks weight skew and collapse; memoryless references such as OU or VPSDE simplify bridge sampling and score evaluation; and steep energies or high-dimensional dynamics can increase weight variance (Guo et al., 4 Oct 2025). In divergence-induced weighted denoising, the local conditional objective for non-KL divergences is a surrogate justified asymptotically in the small-mismatch regime and is not proven to be a tight upper bound on the marginal divergence (Li et al., 21 Jun 2026).

6. Proximal averages and related proximal constructions

A distinct proximal construction comes from convex analysis rather than path-space optimization. The proximal average of two proper, lower semicontinuous, convex functions is

$P$1

with $P$2, $P$3, and

$P$4

It also satisfies

$P$5

(Bauschke et al., 2018). The associated path in $P$6 is described as a true homotopy.

For negative entropy $P$7 and quadratic energy $P$8, the proximal average produces explicit Lambert $P$9 formulas, and the resolvent identity yields computationally simple proximal mappings (Bauschke et al., 2018). In the supplied adaptation to WDCE, one defines

$X=(X_t)_{t\in[0,T]}$0

so that

$X=(X_t)_{t\in[0,T]}$1

For cross-entropy without simplex constraint, the scalar proximal mapping is

$X=(X_t)_{t\in[0,T]}$2

while for KL-type terms it is Lambert-$X=(X_t)_{t\in[0,T]}$3 based:

$X=(X_t)_{t\in[0,T]}$4

This proximal-average route is mathematically distinct from PDNS. The former blends convex penalties pointwise through resolvents and Moreau envelopes; the latter applies a proximal point method to probability measures on trajectories. The shared vocabulary can obscure that distinction, but the constructions act on different objects and solve different optimization problems.