Generative modelling with jump-diffusions

Abstract: Score-based diffusion models generate samples from an unknown target distribution using a time-reversed diffusion process. While such models represent state-of-the-art approaches in industrial applications such as artificial image generation, it has recently been noted that their performance can be further improved by considering injection noise with heavy tailed characteristics. Here, I present a generalization of generative diffusion processes to a wide class of non-Gaussian noise processes. I consider forward processes driven by standard Gaussian noise with super-imposed Poisson jumps representing a finite activity Levy process. The generative process is shown to be governed by a generalized score function that depends on the jump amplitude distribution. Both probability flow ODE and SDE formulations are derived using basic technical effort, and are implemented for jump amplitudes drawn from a multivariate Laplace distribution. Remarkably, for the problem of capturing a heavy-tailed target distribution, the jump-diffusion Laplace model outperforms models driven by alpha-stable noise despite not containing any heavy-tailed characteristics. The framework can be readily applied to other jump statistics that could further improve on the performance of standard diffusion models.

• The paper introduces a generalized score-based generative model that incorporates jump-diffusions using Poisson jumps and a multivariate Laplace distribution.
• The methodology integrates Gaussian noise with sudden jump components via an Ornstein-Uhlenbeck process to capture heavy-tailed behaviors.
• Results show enhanced performance with lower mean-square logarithmic error compared to existing models, demonstrating robustness in modeling heavy-tailed data.

Generative Modelling with Jump-Diffusions

Introduction and Background

The paper "Generative modelling with jump-diffusions" explores a significant extension to traditional score-based generative models by incorporating jump-diffusion processes. Traditional generative diffusion models, widely implemented in applications such as image synthesis, utilize Gaussian noise as the primary stochastic mechanism in describing the transformation of complex data distributions through a diffusion process. The authors note that while these models provide state-of-the-art performance, there is potential for improvement by extending the noise model to include non-Gaussian characteristics, specifically through jump-diffusions with finite activity Lévy processes. Previous studies have shown that alternative noise distributions, such as gamma or fractional Brownian motion, can yield performance improvements in specific tasks like image generation and speech synthesis.

Model Formulation

The paper introduces a generalized score-based generative model that incorporates jump-diffusion processes, common in mathematical finance for modeling discontinuous paths. This model extends the standard continuous-time SDE framework by adding stochastic processes characterized by sudden jumps, modeled using superimposed Poisson jumps on Gaussian noise. The forward process is thus governed by a noise term consisting of Gaussian noise combined with jump components drawn from an isotropic multivariate Laplace distribution.

The key conceptual advance in this framework is the introduction of a generalized score function, which adjusts the probability flow according to the underlying jump processes. The authors derive both the probability flow ODE and an SDE formulation for this generalized model, thus enabling the deployment of this model in practice through either deterministic or stochastic trajectories.

Derivation and Implementation

A critical aspect of this model is the modification of the score function for consistent time-reversal denoising. The authors provide mathematical derivations to establish that the score function under jump-diffusion processes deviates from the traditional $\nabla \log p(x,t)$ form, instead integrating the jump amplitude distributions into the calculation. This derivation enables the model to capture heavy-tailed distributions more effectively than previous models using pure Gaussian noise alone.

For practical implementation, the paper describes a detailed computational framework, noting that the jump-diffusion processes can be efficiently approximated and implemented using discretized versions of the generative processes. The authors employ an Ornstein-Uhlenbeck process with Laplace-distributed jump amplitudes, which helps maintain computational tractability while benefiting from the non-Gaussian noise characteristics.

Results and Comparisons

The paper highlights numerical experiments demonstrating the model's capability in capturing heavy-tailed distributions, validated against benchmarks such as denoising Lévy models (DLPM) and previous Levy-Ito models (LIM). The Jump-Diffusion Laplace (JDL) model shows superior performance in fitting heavy-tailed sample distributions, achieving lower mean-square logarithmic error (MSLE) metrics compared to both DLPM and LIM. Notably, this improvement holds despite the JDL model not inherently possessing heavy-tailed properties in its noise characterization, indicating the model's robustness and flexibility.

Conclusion and Future Directions

The introduction of jump-diffusion processes into generative modelling represents a significant step in developing more comprehensive and flexible generative frameworks. The results suggest that this approach is not only comparable but can indeed enhance existing state-of-the-art methods, particularly in scenarios where data exhibits heavy tails or requires capturing more complex stochastic behaviors. Future research could explore refining the efficiency of the generalized score estimation and exploring the application of this model to a broader array of practical problems in AI, including but not limited to, more complex multi-dimensional datasets and real-time generative tasks. Moreover, further analysis could explore optimizing the implementation parameters such as jump rate and jump size variance to tailor specific applications more effectively.