Rex: Reversible Solvers for Diffusion Models

Diffusion models have quickly become the state-of-the-art for numerous generation tasks across many different applications. Encoding samples from the data distribution back into the models underlying prior distribution is an important task that arises in many downstream applications. This task is often called the inversion of diffusion models. Prior approaches for solving this task, however, are often simple heuristic solvers that come with several drawbacks in practice. In this work, we propose a new family of solvers for diffusion models by exploiting the connection between this task and the broader study of algebraically reversible solvers for differential equations. In particular, we construct a family of reversible solvers using an application of Lawson methods to construct exponential Runge-Kutta methods for the diffusion models. We call this family of reversible exponential solvers Rex. In addition to a rigorous theoretical analysis of the proposed solvers we also emonstrate the utility of the methods through a variety of empirical illustrations.