ED2RM: Equivariant Denoising Diffusion Models based on Riemannian Morphological PDEs

El Hadji S. Diop, Thierno Fall, Mohamed Daoudi

Keywords: Diffusion models, partial differential equations, equivariance, Riemannian manifolds, symmetries

Abstract: Diffusion models have recently emerged and demonstrated remarkable capabilities in high-quality image synthesis and data generation. This work addresses two key issues in recent Denoising Diffusion Probabilistic Models (DDPMs), inspired by nonequilibrium thermodynamics: geometric feature extraction and equivariance. To tackle these challenges, we introduce a geometric approach to the prediction network of DDPMs by designing equivariant morphological partial differential equations (PDEs) for group convolutional neural networks (G-CNNs), referred to as PDE-G-CNNs. These PDEs are formulated on Riemannian manifolds to better capture nonlinearities, represent thin geometrical structures, and incorporate symmetries into the learning process. Our method achieves this by considering a system of two PDEs: a convection term and a first-order Hamilton–Jacobi-type PDE that acts as morphological multiscale dilations and erosions. Preliminary experiments on MNIST and RotoMNIST indicate significant performance gains compared to baseline DDPMs.

Primary Area: generative models

Submission Number: 14171