Radial Tail Repair: Post-Training Correction for Diffusion and Flow Models
cs.LG, stat.ML
Submitted: 2026-10-12
Updated: 2026-10-12
Comments: 30 pages
License: http://creativecommons.org/licenses/by/4.0/
The gist: Diffusion and flow-matching models may struggle to accurately represent distributional tails, especially for heavy-tailed data.
Abstract
Diffusion and flow-matching models may struggle to accurately represent distributional tails, especially for heavy-tailed data. Existing approaches often use heavy-tailed noise or modify the generative dynamics, requiring dedicated retraining and more difficult estimation. We take a simpler approach: rather than relearning the full distribution, we correct only the model's behavior in the extreme regime. Our method builds on the radial-angular decomposition of multivariate heavy tails, where radius and direction become asymptotically independent. This allows us to preserve the directional structure learned by the pretrained model and correct only its radial tail. We estimate the tail onset from reference and generated samples, then rescale generated extremes so that their radial survival matches the target distribution. The resulting method is a lightweight post-training correction applied to a frozen diffusion or flow-matching model, requiring neither retraining nor changes to its sampling dynamics. Experiments on challenging synthetic heavy-tailed distributions and real weather and wildfire datasets show that this simple correction substantially improves extreme-tail coverage using only a small number of fitting samples.
Sources
- Generative modelling with jump-diffusions
- The Reward Was in Your Data All Along: Correcting Flow Matching with Discriminator-Guided RL
- Inference-Time Alignment for Diffusion Models via Variationally Stable Doob's Matching
- Do Heavy Tails Help Diffusion? On the Subtle Trade-off Between Initialization and Training
- Generalizing Score-based generative models for Heavy-tailed Distributions
- DOHF: Online Diffusion Fine-tuning with Doob's h-transform Guidance
- Alignment of Diffusion Model and Flow Matching for Text-to-Image Generation
- Lipschitz regularity in Flow Matching and Diffusion Models: sharp sampling rates and functional inequalities
- On Forgetting and Stability of Score-based Generative models
- Theoretical Guarantees for SMC-Guided Diffusion Sampling
- WEIRDO: WEak resIdual Regularized DOob's h-transform diffusion alignment
Related papers
- Polynomial-Augmented Neural Networks (PANNs) with Weak Orthogonality Constraints for Enhanced Function and PDE Approximation
- AIRL-S: Unifying Reinforcement Learning and Search-Based Test-Time Scaling via Adversarial Inverse Reinforcement Learning
- Transformers as Bayesian In-Context Experimenters: Smoothness-Adaptive Efficient ATE Estimation
- Convergence issues in Relational Concept Analysis based on AOC-posets
- Beliefs Beyond Posteriors: The Predictive Cost of ELBO Binding
- Understanding Diffusion Models via Ratio-Based Function Approximation with SignReLU Networks