Learning to Generate Continually in Latent Space via Bures-Wasserstein Optimization

Continual learning for image generation aims to enable models to learn and synthesize images from new domains while preserving their ability to generate samples from previously learned ones, thereby avoiding the substantial computational cost of retraining modern generative models from scratch. The key challenge is to learn from new data without forgetting previously acquired knowledge. Existing approaches typically combine generative replay or distillation with regularization techniques, such as elastic weight consolidation (EWC), using generated samples to approximate previously learned distributions and constraining parameter updates to mitigate forgetting. However, while effective for pixel-space diffusion, we find that the assumptions underlying these regularization approaches do not generalize well to latent generative processes, such as flow matching. In this work, we propose **La**tent **B**ures-**W**asserstein (LaBW), a general framework for continual learning in latent image generative models with theoretical support. Instead of constraining model parameters, LaBW operates in the autoencoder latent space of generative models in a way that indirectly optimizes the generated image distribution in the Inception feature space. Specifically, our theoretical analysis establishes a local connection between perturbations of the latent endpoint distribution and changes in FID, revealing increased sensitivity to equal absolute covariance changes along low-variance latent directions. We further evaluate our approach across diffusion and flow matching models with both U-Net and DiT backbones on multiple benchmarks (e.g., CIFAR-10, LSUN Churches, and CelebA-HQ), demonstrating a substantial reduction in forgetting while maintaining or improving FID on both original and new domains compared to recent state-of-the-art approaches.