Understanding and Correcting Moment Shift in Continual Learning of Diffusion Models
Abstract
In this work, we investigate catastrophic forgetting in continual learning of diffusion models. Prior work shows that for a well-converged diffusion model, the learned early sampling trajectory can be approximated by an analytic Gaussian form that depends on the data only through its first two moments. We demonstrate that the same approximation holds in continual learning, where the analytic Gaussian form is parameterized by the shifted moments induced by catastrophic forgetting. Over the high-noise regime, forgetting therefore reduces to a tractable moment shift of the generated distribution. This allows us to correct the high-noise trajectory at sampling time by replacing it with the closed-form Gaussian trajectory carrying the true data moments. We call this intervention \textbf{T}arget-\textbf{M}oment \textbf{T}eleportation (TMT), which is training-free and compatible with most existing training-time continual learning methods. Experiments on task-incremental CIFAR-10, MNIST, FashionMNIST and AFHQv2 with a DDPM sampled by DDIM show that TMT uniformly improves average FID and reduces forgetting across three distinct continual learning methods, including the state-of-the-art rank-1 EWC method, while removing 26\% to 68\% of the neural function evaluations (NFEs).