Trace Energy Regularization for Fidelity Oriented Diffusion Models

Diffusion models generate high quality synthetic samples, but maintaining fidelity in diffusion models remains challenging, as the differences between real and synthetic data distributions can degrade downstream model performance when synthetic images are used for downstream training. Standard DDPMs are trained to predict Gaussian noise using a mean squared error objective, whose optimal solution is a conditional mean that intends to reduce trace energy (expected squared $L^2$ norm per dimention of the predicted noise) due to conditional uncertainty. Since the predicted noise directly determines the denoising correction in the reverse process, this shrinkage can produce overly conservative reverse updates. We introduce EnDiff, a lightweight regularizer to match the predicted noise trace energy to that of the injected noise, thereby restoring it toward unit trace energy in an isotropic sense. This yields better calibrated reverse updates and promotes high fidelity generation. Across DDPM, Improved Diffusion, and Latent Diffusion Models on CIFAR10/100, CelebA, and CelebA-HQ datasets, EnDiff reveals a consistent improvement in fidelity oriented metrics. Precision increases from 0.590 to 0.648, and Density increases from 1.062 to 1.458 on the CIFAR100 dataset. These results show that trace energy regularization is a simple and effective strategy to shift diffusion models’ fidelity diversity balance toward fidelity.