# Expressive Score-Based Priors for Distribution Matching with Geometry-Preserving Regularization

arXiv: https://arxiv.org/abs/2506.14607 · HTML: https://mlsanity.com/paper/2506.14607/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 32 of its 50 references that are in the map.
- First posted on arXiv 17 Jun 2025.

## Builds on (most cited first, 30 of 32)

- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [Denoising Diffusion Implicit Models](https://mlsanity.com/paper/2010.02502/) (2020)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [Generative Modeling by Estimating Gradients of the Data Distribution](https://mlsanity.com/paper/1907.05600/) (2019)
- [Domain-Adversarial Training of Neural Networks](https://mlsanity.com/paper/1505.07818/) (2015)
- [DreamFusion: Text-to-3D using 2D Diffusion](https://mlsanity.com/paper/2209.14988/) (2022)
- [Normalizing Flows for Probabilistic Modeling and Inference](https://mlsanity.com/paper/1912.02762/) (2019)
- [An Introduction to Variational Autoencoders](https://mlsanity.com/paper/1906.02691/) (2019)
- [Maximum Likelihood Training of Score-Based Diffusion Models](https://mlsanity.com/paper/2101.09258/) (2021)
- [Score-based Generative Modeling in Latent Space](https://mlsanity.com/paper/2106.05931/) (2021)
- [Domain Generalization via Invariant Feature Representation](https://mlsanity.com/paper/1301.2115/) (2013)
- [Disentangling by Factorising](https://mlsanity.com/paper/1802.05983/) (2018)
- [Learning Adversarially Fair and Transferable Representations](https://mlsanity.com/paper/1802.06309/) (2018)
- [Are GANs Created Equal? A Large-Scale Study](https://mlsanity.com/paper/1711.10337/) (2017)
- [VAE with a VampPrior](https://mlsanity.com/paper/1705.07120/) (2017)
- [The Variational Fair Autoencoder](https://mlsanity.com/paper/1511.00830/) (2015)
- [Central Moment Discrepancy (CMD) for Domain-Invariant Representation Learning](https://mlsanity.com/paper/1702.08811/) (2017)
- [Data Determines Distributional Robustness in Contrastive Language Image Pre-training (CLIP)](https://mlsanity.com/paper/2205.01397/) (2022)
- [Composing graphical models with neural networks for structured representations and fast inference](https://mlsanity.com/paper/1603.06277/) (2016)
- [Invariant Representations without Adversarial Training](https://mlsanity.com/paper/1805.09458/) (2018)
- [Conditional Learning of Fair Representations](https://mlsanity.com/paper/1910.07162/) (2019)
- [Controllable Guarantees for Fair Outcomes via Contrastive Information Estimation](https://mlsanity.com/paper/2101.04108/) (2021)
- [Learning Flat Latent Manifolds with VAEs](https://mlsanity.com/paper/2002.04881/) (2020)
- [Do Vision-Language Pretrained Models Learn Composable Primitive Concepts?](https://mlsanity.com/paper/2203.17271/) (2022)
- [A Large-Scale Study on Regularization and Normalization in GANs](https://mlsanity.com/paper/1807.04720/) (2018)
- [Isometric Representation Learning for Disentangled Latent Space of Diffusion Models](https://mlsanity.com/paper/2407.11451/) (2024)
- [Improving GAN Training with Probability Ratio Clipping and Sample Reweighting](https://mlsanity.com/paper/2006.06900/) (2020)
- [Uncovering Challenges of Solving the Continuous Gromov-Wasserstein Problem](https://mlsanity.com/paper/2303.05978/) (2023)
- [Cooperative Distribution Alignment via JSD Upper Bound](https://mlsanity.com/paper/2207.02286/) (2022)

## Abstract (arXiv)

Distribution matching (DM) is a versatile domain-invariant representation learning technique that has been applied to tasks such as fair classification, domain adaptation, and domain translation. Non-parametric DM methods struggle with scalability and adversarial DM approaches suffer from instability and mode collapse. While likelihood-based methods are a promising alternative, they often impose unnecessary biases through fixed priors or require explicit density models (e.g., flows) that can be challenging to train. We address this limitation by introducing a novel approach to training likelihood-based DM using expressive score-based prior distributions. Our key insight is that gradient-based DM training only requires the prior's score function -- not its density -- allowing us to train the prior via denoising score matching. This approach eliminates biases from fixed priors (e.g., in VAEs), enabling more effective use of geometry-preserving regularization, while avoiding the challenge of learning an explicit prior density model (e.g., a flow-based prior). Our method also demonstrates better stability and computational efficiency compared to other diffusion-based priors (e.g., LSGM). Furthermore, experiments demonstrate superior performance across multiple tasks, establishing our score-based method as a stable and effective approach to distribution matching. Source code available at https://github.com/inouye-lab/SAUB.