GEXIT-weighted posterior training
Implementation & benchmark of arXiv:2607.10067 — Conservation Laws for Diffusion Models
Source paper: Conservation Laws for Diffusion Models arXiv:2607.10067 ⓘ · analyzed Aug 30, 2026
AI-generated research hypothesis, automatically tested. Not peer-reviewed.
Idea description
Use the conservation-law density to weight diffusion training examples by noise level instead of relying on uniform, cosine, or manually selected SNR weighting. This emphasizes noise regions whose local information contribution is largest while clipping the weights to prevent rare regions from destabilizing optimization.
Formulas
Mathematical statement
For clean data $X$ and noisy observation $Y(\tau)$, define the model posterior cross-entropy $\ell_\theta(\tau)=\mathbb E[-\log q_\theta(X\mid Y(\tau),\tau)]$. The optimal posterior is $p_\tau(X\mid Y)$, and the mismatch identity is $\ell_\theta(\tau)-H(X\mid Y(\tau))=\mathbb E_Y\mathrm{KL}(p_\tau(\cdot\mid Y)\|q_\theta(\cdot\mid Y,\tau))\geq0$. The conservation density is $g(\tau)=\partial_\tau H(X\mid Y(\tau))$, whose existence follows from the paper's channel regularity condition: $W_t(y\mid x)$ is absolutely continuous in $t$, has score $S_t(x,y)=\partial_tW_t(y\mid x)/W_t(y\mid x)$, and admits an integrable derivative envelope. Use $w(\tau)\propto |g(\tau)|$ to focus training on posterior errors contributing most to the integrated information quantity, with clipping for finite-sample stability.
Implementation notes
Modify the diffusion loss while keeping the network and optimizer unchanged. Maintain a calibration table over $M$ noise levels. Every 1,000 updates, evaluate the current model on a held-out batch at adjacent noise values and compute averaged posterior losses $\widehat\ell_j$. Estimate $\widehat g_j=(\widehat\ell_{j+1}-\widehat\ell_j)/(\tau_{j+1}-\tau_j)$, smooth its absolute value across neighboring bins, add a small $\epsilon$, normalize its numerical integral, and clip to $[w_{\min},w_{\max}]$. Stop gradients through the table and update it with exponential moving averages so that the weighting does not oscillate with minibatch noise. During training, sample $\tau$ from the existing distribution $r(\tau)$, generate $y\sim W_\tau(\cdot\mid x)$, compute the normal posterior or denoising cross-entropy, multiply it by the interpolated $w(\tau)$, and backpropagate. For Gaussian diffusion, estimate the density from the model's half-MSE/I-MMSE curve; for discrete channels, use posterior NLL finite differences. Test on CIFAR-10 with a 50M-parameter DDPM and on text8 with a small discrete diffusion transformer. Compare uniform, cosine, log-SNR, and GEXIT weighting at equal FLOPs. Track integrated validation NLL, per-noise loss, gradient variance, and FID or BPC. The expected signal is faster reduction of integrated NLL and fewer persistently underfit noise bands without increased gradient variance.
Verification
This idea has not been verified yet.
Verification happens in two stages: Stage 1 — a mechanism check on a toy system confirms the claimed mathematical phenomenon reproduces; Stage 2 — a benchmark implements the idea on a real (small) neural network task and compares it against a tuned baseline over 8 paired seeds with a permutation test.
Artifacts
Artifacts unavailable.