---
title: Proximal WDCE in Diffusion Neural Samplers
url: https://www.emergentmind.com/topics/proximal-weighted-denoising-cross-entropy-wdce
type: topic
---

# Proximal WDCE in Diffusion Neural Samplers

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" [2510.03824], 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" [2606.22521], 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 [2510.03824]. 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 [2606.22521].

## 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]}$ [2510.03824]. If the reference path measure is memoryless and has terminal marginal $\nu$, the terminal reward is
$$
r(x):=-\beta V(x)-\log \nu(x),
$$
and the optimal path measure satisfies
$$
P^*(X)\propto P(X)e^{r(X_T)}.
$$
Under memorylessness, $P_T^*=\pi$.

PDNS applies a proximal step with step size $\eta_k>0$:
$$
P^{\theta_k^*}=\arg\min_{P^\theta}\left\{-\mathbb{E}_{P^\theta}[r(X_T)] + \mathrm{KL}(P^\theta\|P) + \frac{1}{\eta_k}\mathrm{KL}(P^\theta\|P^{\theta_{k-1}})\right\}.
$$
The solution is a geometric interpolation in path space,
$$
P^{\theta_k^*}\propto (P^{\theta_{k-1}})^{\frac{1}{\eta_k+1}}(P^*)^{\frac{\eta_k}{\eta_k+1}},
$$
with density ratio
$$
\frac{dP^{\theta_k^*}}{dP^{\theta_{k-1}}}\propto
\left(\frac{dP^*}{dP^{\theta_{k-1}}}\right)^{\frac{\eta_k}{\eta_k+1}}.
$$
If one initializes with $P^{\theta_0}=P$ and solves each subproblem to optimality, then
$$
P^k \propto P^{\lambda_k}(P^*)^{1-\lambda_k},\qquad
\lambda_k:=\prod_{i=1}^k\frac{1}{\eta_i+1}\to 0,
$$
and the terminal marginal obeys
$$
P_T^k \propto \pi^{1-\lambda_k}\nu^{\lambda_k}.
$$
This gives a staged path from the reference terminal law $\nu$ to the target $\pi$, 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 [2510.03824].

## 3. Continuous and discrete proximal WDCE objectives

In continuous state space, the controlled dynamics are
$$
dX_t=\bigl(b_t(X_t)+\sigma_t u_t^\theta(X_t)\bigr)\,dt+\sigma_t\,dW_t,\qquad X_0\sim\mu,
$$
with base drift $b_t$, noise scale $\sigma_t$, and control $u_t^\theta$ [2510.03824]. The practical proximal WDCE objective replaces pathwise log-likelihood with a denoising or bridge-matching term that depends only on $(X_t,X_T)$:
$$
F(P^\theta;P^{\theta_k^*})
=
\mathbb{E}_{t\sim \mathrm{Unif}[0,T],\,X\sim P^{\theta_k^*}}
\left[
\frac12
\left\|u_t^\theta(X_t)-\sigma_t\nabla \log P_{T\mid t}(X_T\mid X_t)\right\|^2
\right].
$$
Using importance sampling from the previous iterate,
$$
F(P^\theta;P^{\theta_k^*})
=
\mathbb{E}_{t\sim \mathrm{Unif}[0,T],\,X\sim P^{\bar\theta}_{k-1}}
\left[
\left(
e^{r(X_T)}\frac{dP}{dP^{\bar\theta}_{k-1}}(X)
\right)^{\frac{\eta_k}{\eta_k+1}}
\frac12
\left\|u_t^\theta(X_t)-\sigma_t\nabla \log P_{T\mid t}(X_T\mid X_t)\right\|^2
\right].
$$
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 $-\log P^\theta(X)$ 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 $s_\theta(x)_{i,n}$ over replacements [2510.03824]. The proximal WDCE becomes a weighted denoising cross-entropy:
$$
\mathbb{E}_{X\sim P^{\bar\theta}_{k-1}}
\left[
\left(\frac{dP^{\theta_k^*}}{dP^{\bar\theta}_{k-1}}(X)\right)
\mathbb{E}_{\lambda\sim\mathrm{Unif}(0,1)}
\left[
\frac{1}{\lambda}
\mathbb{E}_{x\sim\mu_\lambda(\cdot\mid X_T)}
\sum_{d:x^d=M}
\bigl(-\log s_\theta(x)_{d,X_T^d}\bigr)
\right]
\right],
$$
with
$$
\frac{dP^{\theta_k^*}}{dP^{\bar\theta}_{k-1}}(X)\propto
\left(
e^{r(X_T)}\frac{dP}{dP^{\bar\theta}_{k-1}}(X)
\right)^{\frac{\eta_k}{\eta_k+1}}.
$$
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,
$$
\frac{dP}{dP^{\bar\theta}_{k-1}}(X)=
\exp\left(
-\frac12\int_0^T \|u^{\bar\theta}_{k-1,t}(X_t)\|^2\,dt
+
\int_0^T u^{\bar\theta}_{k-1,t}(X_t)\cdot dW_t
\right),
$$
whereas in discrete masked diffusion they are computed from jumpwise log-ratios between the reference and learned generators [2510.03824].

## 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 [2606.22521]. With per-step residual
$$
e_t:=\epsilon_\theta(x_t,t)-\epsilon
$$
and mismatch
$$
m_t=w_t\|e_t\|^2,
$$
the per-step likelihood ratio is lognormal, $\log r_t\sim \mathcal{N}(-m_t,2m_t)$, so the conditional $f$-divergence becomes a scalar function $h_G(m_t)$. Summing local divergences gives
$$
L_f(\theta)=\sum_{t=1}^T \mathbb{E}[h_G(m_t)],
$$
and the gradient takes the form
$$
\nabla_\theta L_f(\theta)
=
\sum_{t=1}^T h_G'(m_t)\,w_t\,\nabla_\theta
\|\epsilon-\epsilon_\theta(x_t,t)\|^2.
$$
This yields divergence-induced weighted denoising in residual space, with influence weight
$$
w_f(e_t):=h_G'(w_t\|e_t\|^2).
$$

For the KL generator, WDCE is exactly the standard weighted denoising objective:
$$
h_{\mathrm{KL}}(m)=m,\qquad h'_{\mathrm{KL}}(m)=1,
$$
so the residual-space weight is constant:
$$
w_{\mathrm{KL}}(e_t)=1.
$$
The interpretation given in the paper is that KL or WDCE assigns constant influence to all residuals regardless of size. Its effective score
$$
\psi_{\mathrm{KL}}(r)=2w_t r
$$
is unbounded, so large-error samples contribute disproportionately and training is non-robust under contamination [2606.22521].

The same paper contrasts WDCE with bounded-influence divergences. For Hellinger,
$$
h_{\mathrm{HD}}(m)=4(1-e^{-m/4}),\qquad h'_{\mathrm{HD}}(m)=e^{-m/4},
$$
and for negative exponential divergence,
$$
h'_{\mathrm{NED}}(m)\to 0\quad\text{as }m\to\infty,
$$
with asymptotic decay slightly more aggressive than Hellinger because of an extra $1/\sqrt{m}$ 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,
$$
h_G(m)=m+O(m^2),
$$
so WDCE is adequate and simplest to implement on clean data and in small-mismatch regimes [2606.22521].

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
$$
\gamma_k:=\frac{\eta_k}{1+\eta_k},
$$
so small $\eta_k$ gives small $\gamma_k$ and strongly tempers weights, while large $\eta_k$ weakens the proximal constraint and approaches non-proximal WDCE [2510.03824]. Equivalently, the geometric schedule uses
$$
\lambda_k=\prod_i \frac{1}{\eta_i+1},
$$
which gradually decreases to $0$ and moves the terminal distribution from $\nu$ toward $\pi$. An adaptive schedule chooses $\eta_k$, or $\gamma_k$, to satisfy a trust-region bound
$$
\widehat{\mathrm{KL}}(P^{\theta_{k-1}}\|P^{\theta_k^*})\le \epsilon,
$$
estimated from buffer weights.

The algorithmic implementation is buffer-based. At stage $k$, one rolls out trajectories from $P^{\theta_{k-1}}$ or $P^{k-1}$, computes the proximal weights from the Radon–Nikodym derivatives and terminal reward, and stores only terminal pairs $(X_T^{(i)},w^{(i)})$ in a replay buffer [2510.03824]. 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 $O(\mathrm{NFE})$, 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) [2510.03824]. 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 $\gamma_k$ show that large $\gamma_k$ collapses to incomplete modes quickly, whereas small $\gamma_k$ preserves coverage; the adaptive scheduler improves stability and convergence by increasing $\gamma_k$ over stages.

The robust DDPM study gives a complementary empirical picture for the non-proximal KL version. On CIFAR-10 under contamination $\epsilon=0.30$, KL or WDCE yields $\mathrm{FID}=93.0\pm 29.6$, Hellinger gives $91.9\pm 33.4$, and NED gives $77.5\pm 19.7$ [2606.22521]. In a single-seed comparison against standard robust losses at the same contamination level, Huber $(\delta=1.0)$ gives $\mathrm{FID}=94.0$, clipped MSE $(\tau=10)$ gives $90.5$, HD gives $56.2$, and NED gives $55.1$. The practical overhead is also characterized: HD adds negligible cost, while NED uses per-sample quadrature with fully vectorized overhead below $1\%$ 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 [2510.03824]. 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 [2606.22521].

## 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_\lambda(f,g)(x)
=
\left(
(1-\lambda)\left(f+\tfrac12\|\cdot\|^2\right)^*
+
\lambda\left(g+\tfrac12\|\cdot\|^2\right)^*
\right)^*(x)
-
\tfrac12\|x\|^2,
$$
with $p_0(f,g)=f$, $p_1(f,g)=g$, and
$$
(p_\lambda(f,g))^*=p_\lambda(f^*,g^*).
$$
It also satisfies
$$
e_1 p_\lambda(f,g)=(1-\lambda)e_1 f+\lambda e_1 g,\qquad
\mathrm{prox}_{p_\lambda(f,g)}=(1-\lambda)\mathrm{prox}_f+\lambda \mathrm{prox}_g
$$
[1807.08878]. The associated path in $\lambda$ is described as a true homotopy.

For negative entropy $f(x)=\sum_i(x_i\log x_i-x_i)$ and quadratic energy $g(x)=\tfrac12\|x\|^2$, the proximal average produces explicit Lambert $W$ formulas, and the resolvent identity yields computationally simple proximal mappings [1807.08878]. In the supplied adaptation to WDCE, one defines
$$
WDCE_\lambda(x):=p_\lambda(f_{\mathrm{entropy}}(x),g_y(x)),
\qquad
g_y(x)=\tfrac12\|x-y\|^2,
$$
so that
$$
\mathrm{prox}_{WDCE_\lambda}(u)
=
(1-\lambda)\mathrm{prox}_{f_{\mathrm{entropy}}}(u)
+
\lambda \mathrm{prox}_{g_y}(u).
$$
For cross-entropy without simplex constraint, the scalar proximal mapping is
$$
\mathrm{prox}_{\tau(-p\log x)}(y)
=
\frac{y+\sqrt{y^2+4\tau p}}{2},
$$
while for KL-type terms it is Lambert-$W$ based:
$$
\mathrm{prox}_{\tau[x\log(x/q)-x]}(y)
=
\tau W\!\left(\frac{q}{\tau}e^{y/\tau}\right).
$$

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.

Source: https://www.emergentmind.com/topics/proximal-weighted-denoising-cross-entropy-wdce