📊 statistics

Self-Regulating Annealing in Heavy-Tailed Diffusion Models

This paper proposes a novel SDE-based sampler for heavy-tailed diffusion models that utilizes a state-dependent diffusion coefficient to induce a self-regulating annealing mechanism, which is theoretically and experimentally shown to be essential for accurately reproducing samples from heavy-tailed distributions.

Original authors: Keito Wakatsuki, Hideaki Shimazaki

Published 2026-06-02

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Keito Wakatsuki, Hideaki Shimazaki

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Picture: Fixing the "Average" Problem

Imagine you are teaching a robot to draw pictures of cats. Most modern AI models (called Diffusion Models) work like a sculptor who starts with a block of marble (pure noise) and slowly chips away pieces to reveal the cat.

In standard models, the "chipping" process assumes that mistakes are usually small and average. It's like assuming every cat is a typical house cat. This works great for normal data, but it struggles with heavy-tailed datasets.

What is a "heavy tail"?
Think of a crowd of people. Most are average height, but a few are giants or dwarfs. In a "heavy-tailed" dataset, those giants and dwarfs (outliers) are much more common than in a normal crowd. Standard AI models tend to ignore these outliers or shrink them down to look "average," losing the true nature of the data.

The Solution: Heavy-Tailed Diffusion Models (HTDMs)

The researchers at Kyoto University looked at a new type of model called HTDMs. Instead of using a standard "Gaussian" (bell curve) math rule, these models use a Student's t-distribution.

The Analogy: If the standard model is a strict teacher who only accepts answers that are close to the average, the HTDM is a more flexible teacher who expects and accepts wild, extreme answers (the "giants" and "dwarfs") as part of the normal class.

The Missing Piece: The "Self-Regulating" Sampler

While HTDMs were invented to handle these extreme data points, the researchers noticed a problem with how they generate (sample) new data.

Previous methods mostly used a smooth, predictable path (an ODE) to generate images. But the researchers asked: What if we used a path that includes some randomness (an SDE), but makes that randomness change based on where the robot is?

They proposed a new SDE-based sampler that acts like a self-regulating thermostat.

The "Self-Regulating" Mechanism Explained

Imagine you are trying to find a hidden treasure (the final clean image) in a foggy forest.

When you are far away: You are lost and don't know where the treasure is. The model says, "Okay, we need to explore!" It turns up the "noise" (the fog and wind) to help you jump around and search a wider area.
When you are close: You can see the treasure. The model says, "Great, you're close! Let's be careful." It turns down the noise so you can walk slowly and precisely to the exact spot without tripping over.

The Innovation:
In this paper, the "noise level" isn't fixed. It is state-dependent.

If the AI's current guess is very different from the target (the "distance" is large), the model automatically injects more noise to help it explore.
If the guess is close to the target, the model injects less noise to refine the details.

The authors call this "Self-Regulating Annealing." It's like the AI has an internal compass that knows when to be wild and when to be calm, all by itself.

What They Found

The researchers tested this new method on a simple, made-up dataset that was full of extreme outliers (a "heavy-tailed" distribution).

The Standard Model (Gaussian): It failed to capture the extremes. It smoothed them out, like a photo editor blurring out a giant's height.
The Old HTDM Method (ODE): It did a better job but was rigid.
The New Method (Self-Regulating SDE): This was the winner.
- It successfully reproduced the "giants" and "dwarfs" of the data.
- It proved that the "self-regulating" noise was necessary. When they turned off the feature (making the noise constant), the model failed to generate the extreme values again.

Summary

The paper introduces a smarter way for AI to generate data that has extreme outliers. By adding a "smart thermostat" to the generation process, the AI automatically knows when to shake things up (add noise) to explore wild possibilities and when to settle down to refine the result. This ensures that the final output keeps the true, wild character of the original data, rather than smoothing it out into something boring and average.

Technical Summary: Self-Regulating Annealing in Heavy-Tailed Diffusion Models

Problem Statement
Standard diffusion models rely on Gaussian formulations for both forward diffusion and reverse generative processes. While this approach offers analytical tractability and has proven successful in image generation, its suitability for heavy-tailed or outlier-prone datasets remains unclear. Recent extensions, specifically Heavy-Tailed Diffusion Models (HTDMs) based on the Student's $t$ -distribution, have been proposed to improve tail fidelity. However, while HTDMs theoretically support Stochastic Differential Equation (SDE) sampling, existing literature (e.g., Pandey et al. [1]) has primarily focused on Ordinary Differential Equation (ODE) solvers, leaving the SDE perspective largely unexplored. A critical gap exists in understanding how to construct an SDE-based sampler for HTDMs that preserves the heavy-tailed properties of the target distribution.

Methodology
The authors propose a novel SDE-based sampler for HTDMs that explicitly incorporates a state-dependent diffusion coefficient. The methodology proceeds as follows:

Derivation of the Reverse SDE: Starting from the discrete-time reverse kernel of HTDMs (which uses a Student's $t$ -distribution), the authors take the continuous-time limit ( $\Delta t \to 0$ ). Unlike previous approaches that omitted the state-dependent factor in the posterior covariance during training, this work retains and explicitly incorporates the factor $\frac{\nu + \Delta_t^2}{\nu + d}$ during the sampling phase.
State-Dependent Coefficient: The resulting SDE includes a diffusion term scaled by a coefficient $\alpha(\Delta_t^2(x_t), t) = \sqrt{\frac{\nu + \Delta_t^2}{\nu + d - 2}}$ . Here, $\Delta_t^2$ represents the squared Mahalanobis distance between the current noisy state $x_t$ and the denoiser's estimate of the target state $D_\theta(x_t, \sigma_t)$ .
Self-Regulating Annealing Mechanism: The authors theoretically demonstrate that this state dependence induces a "self-regulating annealing" mechanism.
- When the current state is far from the target estimate (large $\Delta_t^2$ ), the diffusion coefficient increases, injecting stronger noise to promote broader exploration.
- When the state is close to the target (small $\Delta_t^2$ ), the diffusion coefficient decreases, leading to milder noise and more gradual convergence.
- This behavior is analogous to state-dependent modulation of effective temperature in deformed associative memory models.
Sampling Algorithm: The proposed sampler (Algorithm 1) numerically solves this SDE (using the Euler–Maruyama method) from $t=T$ to $t=0$ , starting from a Student's $t$ prior.

Key Contributions

Novel Sampler: The paper proposes the first SDE-based sampler specifically designed for HTDMs that explicitly includes the state-dependent diffusion coefficient derived from the Student's $t$ posterior.
Theoretical Insight: It provides a theoretical demonstration that the state dependence of the diffusion coefficient induces a self-regulating annealing mechanism, linking HTDM dynamics to memory retrieval dynamics in statistical physics.
Empirical Validation: Through experiments on a synthetic one-dimensional Student's $t$ dataset, the authors show that this state dependence is necessary for reproducing heavy-tailed distributions.

Experimental Results
The authors evaluated four samplers on a synthetic dataset ( $St_1(x|0, 1, 3)$ ):

VE-SDE: A standard Gaussian baseline.
t-ODE: An ODE-based sampler for HTDMs (from prior work).
t-SDE: The proposed SDE-based sampler with the state-dependent coefficient.
Ablated t-SDE: The proposed SDE sampler with the state-dependent coefficient set to 1 (constant).

Findings:

Tail Fidelity: The t-SDE and t-ODE closely matched the ground truth in the extreme tails (quantile-quantile plots and extreme-tail probabilities). In contrast, the Gaussian VE-SDE severely underestimated tail probabilities, and the Ablated t-SDE failed to generate any samples in the extreme tails (probability reduced to zero).
Overall Quality: While the ablated t-SDE showed poor overall quality (high 1-Wasserstein distance), the proposed t-SDE achieved performance comparable to the t-ODE and superior to the Gaussian baseline.
Conclusion: The state-dependent diffusion coefficient is essential for faithfully reproducing heavy-tailed behavior.

Significance and Claims
The paper claims that the proposed SDE-based sampler bridges the gap between the theoretical potential of HTDMs and practical sampling. By explicitly incorporating the state-dependent diffusion coefficient, the sampler naturally achieves a self-regulating annealing mechanism that adapts the effective noise scale based on the distance to the target. The authors assert that this mechanism is not merely a theoretical curiosity but a practical necessity for generating samples from heavy-tailed distributions, as evidenced by the failure of the ablated variant.

Limitations
The authors note that their current evaluation is limited to a synthetic one-dimensional setting. The selection of the degrees-of-freedom parameter ( $\nu$ ) was performed via a small sweep, and the paper identifies extending this analysis to high-dimensional benchmarks (e.g., images) and developing principled methods for selecting $\nu$ as directions for future work.