Diffusion Models as an Extension of Variational Autoencoders

Diffusion models are likelihood-based generative models that utilize time-forward and time-backward stochastic differential equations (SDEs). While conventional diffusion models (such as the score-based model) only learn the time-backward process, more flexible frameworks have been recently proposed to also learn the time-forward process by employing the Schrödinger Bridge (SB), which improves the expressive power of the model. Since the SB is described by the two nonlinear partial differential equations (PDEs), it is difficult to solve them. In the SB-type diffusion model, this difficulty is addressed with the mapping of the PDEs to SDEs by nonlinear Feynman-Kac theory, and the objective function for simultaneously leaning both the forward and backward processes is constructed. However, the mathematical structure behind this framework is too complicated to give an intuitive understanding of the objective function.

In this work, by reinterpreting the SB-type models as an extension of Variational Autoencoders (VAEs), we propose a unified scheme to construct the objective function for the diffusion models. In our scheme, the information processing inequality plays an essential role. This inequality enables us to extend the number of hidden layers in VAEs to infinite, and the transition probabilities between the layers can be described by SDEs. As a result, it is revealed that the objective function for the SB-type model is composed of the prior loss and the drift matching parts.