Expressive Score-Based Priors for Distribution Matching with Geometry-Preserving Regularization
Authors: Ziyu Gong
Jim Lim
David I. Inouye
Venue: ICML 2025 poster (2025)
Keywords: Distribution Matching
Latent Representation
Score-Based Models
Diffusion Models
Primary area: general_machine_learning->unsupervised_and_semisupervised_learning
TL;DR: We propose a new VAE-based distribution matching method that uses (1) an expressive score-based prior and (2) a geometric-preserving regularization inspired by Gromov-Wasserstein.

Abstract: 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.

Expressive Score-Based Priors for Distribution Matching with Geometry-Preserving Regularization
Ziyu Gong Jim Lim 1
David I. Inouye'
Abstract
1. Introduction
As machine learning (ML) continues to advance, trustwor- thy ML systems not only require impressive performance but also properties such as fairesss, robustness, causality, and explainability. While scaling data and models im- prove performance (Kaplan et al., 2020), simple scaling may not address these issues. For example, historical bias or imbalanced data can cause even well-trained models to produce unfair outcomes, requiring additional constraints to mitigate such biases. Distribution matching (DM), also known as distribution alignment or domain-invariant rep- resentation learning, has emerged as a promising approach to address these challenges. By minimizing the divergence between latent representations, distribution matching can introduce additional objectives to ML systems, enabling them to learn representations that are fair, robust, and causal. This approach has been successfully applied to a wide range of problems, including domain adaptation (Ganin et al., 2016; Zhao et al., 2018), domain generalization (Muandet et al., 2013) causal discovery (Spirtes & Zhang, 2016), and fairness- aware learning (Zemel et al., 2013).
Distribution matching (DM) is a versatile domain- invariant representation learning technique that has been applied to tasks such as fair classifica- tion, domain adaptation, and domain translation. Non-parametric DM methods struggle with scala- bility 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 train. We ad- dress this limitation by introducing a novel ap- proach 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 de- noising score matching. This approach elimi- nates biases from fixed priors (e.g., in VAEs), en- abling more effective use of geometry-pressrring regularization, while avoiding the challenge of learning an explicit prior density model (e.g.. a flow-based prior). Our method also demon- strates better stability and computational effi- ciency compared to other diffusion-based priors (e.g., LSGM). Furthermore, experiments demon- strate superior performance across multiple tasks, establishing our score-based method as stable and effective approach to distribution matching. Source code available at https: /ttp///gttt////// com/ waw.oouyelabb/aae
DM methods can be broadly categorized into parametric and non-parametric approaches. Non-parametric methods, such as kernel Maximum Mcan Discrepancy et al., 2015; Zellinger et al., 2017) and Sinkhorn divergence (Feydy al., 2019), operate directly on sample distributions without assuming specific parametric form. Parametric DM methods, on the other hand, rely on modeling distri- butions with explicit parameters and can be further divided into adversarial and non-adversarial likelihoodl-asel ap- proaches. Adversarial methods, exemplified by Generative Adversarial Networks (GodslGoodfelloow et al., 2014), frame distribution matching as a minimax game between generator and discriminator. While highly expressive and capable of capturing complex data distributions, these methods suffer from well-documented issues such as train- ing instability, mode collapse, and sensitivity to hyperpa- rameters(Luce et al., 2018; Kurach et al., 2019; Farnia & Ozdaglar, 2020; Nie & Patel, 2020; Wu et al., 2020; Han et al., 2023). In contrast, tikeihhoddbbsse approaches lever- age probabilistic models such as variational autoencoders (VAEs) (Kingma et al., 2019) normalizing flows (Papa- makarios et al., 2021) to match distributions by maximiz-
(LouiiLsszzs
'Equal contribution 'Elmore Family School of Electrical and Computer Engineering, Purdue University. West Lafayette. IN, USA. Correspondence Ziyu Gong mmgil2000@mmil..mm. Jim Lim <jimlim9508188 com?, David Inouye <di- nouye mune@purdue.cdux>
Proceedings of the 42nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Expressive Score-Based Priors for Distribution Matching with Geometry-Preservinng Regularization
ing the likelihood of observed data under the model with relatively better training stability and ability. However, nor- malizing flows are restricted by the requirement that the latent dimension must have the same size as the input di- mension. This constraint limits their flexibility in modeling complex latent representations and can hinder their ability to capture lower-dimensional latent structures effectively (Cho et al.. 2022b). On the other hand, VAEs are valued for being able to capture meaningful and structured rep- resentations (Chen et al., 2019 Burgess et al., 2018) in lower dimension. Gong et al. (2024) proposed to use VA for DM task but imposed a simple learnable prior distribu- tion (e.g., Gaussian, Mixture of Gaussian), which aligned poorly with the true data distribution and consequently led to suboptimal performance. The need for a more expressive learnable prior distribution is also important when enforcing constraints, as these constraints ensure that the latent space retains the intrinsic geometry of the data (Uscidda et al., 2024; Nakagawa et al., 2023; Hahm et al., 2024; Lee et al., 2022; Horan et al., 2021; Gropp et al., 2020; Chen et al., 2020) which could facilitate disentan- gled representations in the latent space and consequently improving downstream tasks performance.
follows:
We introduce the Score Function Substitution (SFS) trick that computes exact variational encoder gradients using only the prior's score function, thereby cireum- venting the need for explicit density evaluation.
Leveraging SFS, we develop novel, stable, and effi- cient alternating optimization algorithm for likelihood- based DM with expressive score- priors.
Our method achieves strong performance across di- verse downstream tasks, including fair classification, domain adaptation, and domain translation.
We further demonstrate that our approach enables the effective application of Soometry-preserving regular- ization (Nakagawa et al., 2023), yielding additional performance improvements when a semantically rich latent space is available for the task.
2. Preliminaries
Variational Aligmment Upper Bound (VAUB) The paper by Gong et al. (2024) presents a novel approach to distri- bution matching for learning invariant representations. The author proposes non-adversarial method based on Vari- ational Autoencoders (VAEs), called the VAE Alignment Upper Bound (VAUB). Specifically, they introduce align- ment upper bounds for distribution matching that generalize the Pensen-Shannon Divergence (ISD) with VAE-like ob- jectives. The author formalizes the distribution matching problem with the following VAUB objective:
In order to have expressive prior, our key insight is that, for gradient-based -basen training, likelihood-baad based DM methods do not require computation of the prior density directly. Instead, they only require the gradient of the log probabil- ity of the prior score function. Building on this observation, we propose novel approach that models the prior density through its score function, precisely the computation needed for train- ing. The score function can be efficiently estimated using denoising score matching techniques. enabling us to bypass the challenges associated with learning explicit prior den- sities. Another crucial insight stems from recognizing that DM methods do not inherently require generation capabil- ities; instead, the prior distribution only used to form proper bound for divergence measures during training. This allows us to model the prior using score-based models, where sampling the prior is computationally expensive but score training and inference remain efficient and stable. We demonstrate through extensive experiments that our sim- ple yet effective algorithm significantly improves training stability and achieves superior DM results across various benchmarks. Finally, our framework can also integrate se- mantic information from pretrained models, such as CL.IP (Radford et al., 2021), to capture task-relevant features that reflect higher level semantics. By aligning the latent space with these semantic relationships, our method can ensure that the representations are not only geometrically sound but also contextually meaningful for downstream tasks, such as classification and domain adaptation.
referred to as the
log p(a)2, d) q(=|,, d) p(z)
VAUB(q df) = =min E Eg(z,,,d)
+C,
p(*)
(1)
where q(=|2r, d) is the probabilistic encoder, p(x) 2, d) is the decoder, p(z) is the shared prior, and C' is constant in- dependent of model parameters. The method ensures that the distribution matching loss is an upper bound of the divergence (ISD), up to a constant. This non-adversarial approach overcomes the instability of ad- versarial training, offering robust, stable alternative for distribution matching in fairness, domain adaptation, and robustness applications. Empirical results show that VAUB and its variants outperform traditional adversarial methods, particularly in cases where model invertibility and dimen- sionality reduction are required.
Score-based Models Score-based models (Song et al., 2021c) are class of diffusion models that learn gener- ate data by denoising noisy samples through iterative re- finement. Rather than directly modeling the data distribu- tion p(z), as done in many traditional generative models, score- score-based models focus on learning the gradient of the log-probability density of the target distribution, known as
We summarize our contributions in the field of DM as
2

---

Expressive Score-Based Priors for Distribution Matching with Geometry-Preservinn Regularization
the score function. To learn the score function, (Vincent, 2011) and (Song & Ermon, 2019) propose training the Denoising Score Matching (DSM) objective. Essentially, data points z are perturbed with various levels of Gaussian noise, resulting in noisy observations ã. The score model is then trained to match the score of the perturbed distribution. The DSM objective defined as follows:
where d represents the domain Ad E [1.... DD (e.g.. dif- ferent class datasets modalities), and [3 [0, 1] acts regularizer controlling the mutual information between the latent variable 2 and the data I. (2(22z d) and P. [z|22, d) are the d-th domain probabilistic encoder and decoder, re- spectively, and Qz(()) is a prior distribution that is invariant to domains (Gong et al., 2024). For notational simplicity, we ignore the regularization loss and we assume 3 1. We can split the VAUB objective into three components: reconstruction loss, entropy loss, and cross entropy loss.
DSM = 2L E (s,(ū, o,) - V log (5)2)
(2)
2L
where 9o, (z)xz) represents the perturbed data distribution of Pdma(z)) I is the number of noise scales {o and the
expectation is over the distribution Pasta (2) ( 9o. (t)z).
AVAUB
Baga |- log (1l2,,d)
(5)
When the optimal score network 5 is found, 5 (2)
reconstruction
V log 4o (z) almost surely ((Vincent, 201 201n).(Song & Er- mon, 2019)) and approximates V log Pima when the noise is small (a R 0). Since score-based models learn the gradient of the distribution rather than the distribution itself, generating samples involves multiple iterative refinement steps. These steps typically leverage techniques such as Langevin dynamics, which iteratively updates the sample using the learned score function (Song & Ermon, 2019).
Byy |- log go (2(2,,
E log (v) cross entropy term
entropy term
The prior distribution in the cross-entropy term aligns with the encoder's posterior but often restricted to simple forms like Gaussians or Gaussian mixtures (Gong et al., 2024), which can distort the encoder's transformation function (Us- cidda et al., 2024). To address this, propose an expres- sive, learnable prior that adaptively mitigates such distor- tions, better capturing the underlying data structure.
Gromov- Wasserstein Distance The Optimal Transport (OT) problem seeks the most efficient way to transform one probability distribution into another, minimizing transport cost. Given two probability distributions μ and v over metric spaces (X, dx) and (Z, dz), the OT problem is:
Modeling an arbitrary probabilistic density function (PDF) is computationally expensive due to the intractability of the normalization constant. Therefore, instead of directly modeling the density Q(=), we propose to indirectly pa- rameterize the prior via its score function V, log ((z). While this avoids direct density estimation, the score function alone makes log-likelihood computations diffi- cult. Weighted score matching losses only approximate -ikimum--ikelihood estimation (MLE), and directly opti- mizing MLE using the flow interpretation becomes com- putationally prohibitive each step (Song et al., 2021a). Unlike VAEs, where efficient sampling from the prior is critical, we demonstrate that the distribution matching objective with a score-based prior can be optimized without costly sampling or computing log- likelihood. By reformulating the cross-entropy term as a gradient with respect to the encoder parameters 6, we de- rive an equivalent expression that retains the same gradient value. This allows from the encoder and compute gradients with single evalu- ation of the score function. We call this the Score F Function Substitution (SFS) trick.
inf B( rEll(a.))
where 11(2, v/) is the set of couplings with marginals f and v, and d(x, 2) is cost function, often the Euclidean distance. The Gromov-Wassen
(GW) distance extends OT to compare distributions on difffrrent metric spaces by preserving their relative structures, not absolute distances. For distributions and y' over spaces (X, dx) and (Z, dz), the GW distance is:
it requires solving ODE
inf (TEl(())) inf GWC mell(u.)
- (2, =11²
(4)
3. Methoodolog
decouple score function training
3.1. Training Objective for Distribution Matching with a Score-based Prior
We aim to optimize VAUB (Gong et al., 2024) as our distri- bution matching objective:
Proposition 3.1 (Score Function Substitution (SFS) Trick). If (0(2)z)) is the posterior distribution parameterized by 6, and Qz(z) is the prior distribution parameterized by v. then the gradient of the cross entropy term can be written
log (=(x)2zd) c0(2|x d) (2))
Σ Egg
LDM LVAUB =
3

---

Expressive Score-Based Priors for Distribution Matching with Geometry-Preservinng Regularization
aS: Vo Vo log (z)))
where the expectation in (8) is over the joint distribu- tion of observed and latent variables, i.e., 20(2, x; d) & Pdita (2,, dz)g(z)zz, d), and the expectation in (9) is over the marginal (noisy) posterior distribution g0(2, &, oy) & E [go (2, ())((=))),,((())) If the lower-level opti- mization in Eqn. is solved perfectly, then the upper bound of likelihood represented by SAUB in Eqn. 8 will be tight. While there are many possible approaches to bi-level opti- mization, we choose the simple alternating approach (Xiao et al., 2023; Chen et al., 2021) between the top-level and the bottom-level problems, while holding the parameters of the other optimization problem fixed. Because this simple alternating approach worked well in our experiments, we leave the exploration of more complex bi-level optimization approaches to future work.
(6)
-VoBerrnoetlo [-
V Volog (u)(3)
20
constonz word.
where the notation of za emphasizes its dependence on 0 and- -la=sa denotes that while is equal zg it is treated as a constant with respect to 6.
The full proof can be seen in Appendix A. In practice, Eqn. 6 detaches posterior samples from the computational graph, enabling efficient gradient computation without additional backpropagation dependencies. Details are provided in the next section. Following Proposition 3.1, we propose the score-based prior alignment upper bound (SAUB) objective defined as follows:
During VAE training, the score model is conditioned on the smallest noise level, To = min- to approximate the clean score function of the marginal posterior. As previously men- tioned, the output of the score model is detached to prevent gradient flow, ensuring memory-efficient optimization by focusing solely on the encoder and decoder parameters with- out tracking the score model's computational graph. After optimizing the encoder and decoder, these networks fixed while the score model is updated using Eqn. 9. Theo- retically, if the score model is sufficiently trained enough to fully capture latent distribution, it could be optimized using only small noise levels. However, extensive score model updates after each VAE step are computationally expensive. To mitigate this, we reduce score model updates and train with larger maximum noise level, enhancing stability when the latent representation becomes out-of-distribute
CSAUB
ΣV
Ea-a (z|2,d) - -logg (x/22,
+ +log (2 - (v, log Q (3)
Since our new loss does not affect terms related to sp, and by Proposition 1. we have Ve, Vo.p LSAUB though we note that VVAU and V LSAUB are not equal in general. In the next section, we show how to train all pa- rameters by approximating a bi-level optimization problem.
3.2. Deriving an Alternating Algorithm with Learnable Score-Based Priors
Leveraging SFS, we now develop a novel alternating op- timization algorithm for score-based prior distributions. Specifically, we parametrize the prior through its score function, denoted (+), instead of it's density (.).). Given fixed prior score function, the SAUB objective allows us to optimize the encoder and decoder parameters 0 and to (but cannot be used to update the prior because VVLVA A V.LSAUR). To update the prior, we can use denoising score- motorin objective to match the prior with the encoder's marginal posterior improve the DM varia- tional bound (Cho et al., 2022a; Gong et al., 2024)- indeed, when the prior matches the marginal posterior, the bound becomes tight. Thus, our global problem can be formulated as a bi-level optimization problem where the upper level is the SAUB objective and the lower level is the denoising score matching objective:
(OOD), The complete training process is outlined in Appendix B. We also listed the stabilization and optimization techniques in Appendix C.
3.3. Comparison with Latent Score-Based Generative Models
Latent Score-Based Generative Models (LSGM) (Vahdat et al., 202 provide powerful framework that integrates la- tent variable models with score-based generative modeling, leveraging diffusion processes to enhance data generation quality. A key innovation in LSGM is the introduction of learnable neural network prior, which replaces the tra- ditional cross-entropy term in the Evidence Lower Bound (ELBO) with score-based terms approximated via a diffu- sion model. This idea of incorporating a score-based prior is similar to our method, which also leverages score functions for the prior.
Σ min 8.6
Bap log (1 + log (0(2/ d)
d
(sa 00
(00,00 ~ 0) =-=(2t00))
}} })}
(8)
A crucial challenge of LSGM is the instability associated with computing the Jacobian term when backpropagating through the U-Net of the diffusion model. Computing this Jacobian term is analogous to approximating the Hessian
S.t. v0 E arg min Egg [1Sp(E, - -V& lologe go, (8|=)112
10
4

---

Expressive Score-Based Priors for Distribution Matching with Geometry-Preserving Regularization
10-1
(dam -0.011. Oear
LSGM (Ooa 0 0.001, Gru-
SPS Trick (dam 0.01. Gear
Reconstruction Loss
LSGM (dem 0.01,
(dua Onee LSGG (Dow 0.2 deue 20
1000
104
Grex
LSGG (Pvu 0.001, Aees
Negative Log- ligllikelihood
SPS Trick (dea = 0.01, One
10a
LSGM (0em 0.01, Ous 1)
10P
L.5GM 60am
+.2, Oone
10
10
Epoch
Figure 1. The reconstruction loss and negative log-likelihood are presented logarithmic scale for improved visualization. The experiment uses consistent hyperparameters (B = 0.1), identical VAE architecture, and same pretrained score model.
of the data distribution, which has been empirically shown to be unstable at low noise levels (Poole et al., 2022). Con- versely, our Score Function Substitution (SFS) trick elimi- nates the need to backpropagate through the diffusion model, enabling stable optimization without explicitly computing the Jacobian. In addition, the LSGM loss requires approxi- mating an expectation over noise levels with finite number of Monte Carlo (MC) samples (often a single one). We hypothesize that this MC approximation also contributes to the instability of LSGM. For further details of gradient comparison, please refer to Appendix F.
approximates the true score function, LSGM requires back- propagation through the score model's U-Net, which causes instability at low noise levels due inaccurate gradients. As shown in Fig. 1, when min 0.001, LSGM exhibits catastrophic instability, with diverging NLL and spikes in reconstruction loss. At m 0.1 and min 0.2, LSGM performs better in terms of both reconstruction loss and NLL than at amin = 0.01, indicating that unstable gradi- ents at lower noise levels negatively impacts prior matching. This concerning since low noise levels, like amin = 0.01, are commonly used in practice. In contrast, the SFS trick shows greater stability across noise levels. At 7 0.01, the NLL is better than at min 0.1. which outperforms min 0.2, suggesting that SFS ensures more reliable gra- dients when the score model is trained on lower noise levels. While both L.GM and SAUB degrade at O = 0.001, SFS stabilizes and achieves better NLL than LSGM at m 0.01, demonstrating its robustness in handling small noise configurations.
Comparative Stability: SFS vs. LSGM We assess sta- bility by measuring the posterior's 's negative log-liken likeliiol (NLL) under fixed Gaussian-mixx mixture prior. The prior and target distributions are illustrated Fig. Unlike standard training, which updates encoder, decoder, and prior parame- ters, our approach freezes the prior and uses a score model pre-trained on the defined prior, updating only the encoder and decoder. The same pre-trained score model is used for both SAUB and LSGM to ensure fair comparison. Per- formance is evaluated under score model trained on four minimum noise levels, mii E 0.001, 0.01, 0.1, 0.2, with Fmax = 1 fixed. While lower noise levels should improve likelihood estimation, as the score model more precisely
3.4. Semantie Preservation (SP) in Latent Representations via GW Inspired Constraint
Given expressive score-based prior, now investi- gate how incorporate the
regulariza-
5

---

Expressive Score-Based Priors for Distribution Matching with Geometry-Preservinng Regularization
tion introduced by Nakagawa al. (2023) without inducing the unnecessary biases of fixed priors on the latent distri- bution. Specifically, Nakagawa et al. (2023) introduces the GW metric Lew in an autoencoding framework, and we adopt this regularization in a similar manner:
4. Related Works
Learnable Priors Most variational autoencoders (VAEs) typically use simple Gaussian priors due to the computa- tional challenges of optimizing more expressive priors and the lack of closed-form solutions for their objectives. Early efforts to address this, such as Adversarial Autoencoders (AAEs) (Makhzani et al., 2016), employed adversarial net- works to learn flexible priors, resulting in smoother and more complete latent manifolds.
Liotal = LDM + Law (z))
(10)
Law(90(212)) & GWCost(* = Paitaa x) 90 22)
(11)
[8| [[dx(xx, dz 2') 201
(12)
Subsequent research (Hoffman & Johnson, 2016; Johnson et al., 2017) highlighted that simple priors can lead over- regularized and less informative latent spaces, while (Tom- czak & Welling, 2018 8) empirically showed that more ex- pressive priors improve generative quality, with significant gains in log- oog-likelihood. More recently, Latent Score-based Generative Models (LSGM) (Vahdat et al., 2021) introduced score-based priors, leveraging denoising score-matching objective to learn arbitrary posterior distributions. This approach enables high-quality image generation while cap- turing the majority of the data distribution.
where Pdata represents the data distribution, dx and dz are the predefined metric spaces for the observed and latent spaces, respectively, and λaw controls the importance of the structural preservation loss. LDM (go (=|ar) represents the distribution matching objective with go (z|zz) as the encoder, and Low(ge (2 [z)) is the structural preservation loss where Pdata is the data distribution, dx and dz are the metric spaces for the observed and latent spaces, respectively, and AGw controls the GW loss Lew (=|x)) LDM (ge (zr)) is the distribution matching objective with encoder go (2[xr).
Based Learning Gromov- Wasserstein (GW) distance has found numerous ap- plications in learning problems involving geometric and structural configuration of objects or distributions. Moreover, the GW metric has been adopted for mapping functions in deep neural networks. One of the key benefits of GW distance is its capacity to compare distributions with heterogeneous data and/or dimensional discrepancies.
Gromov-Wasserse
Selection of Metric Space and Distance Functions The GW framework's 's key strength lies in its ability to compare distributions across diverse metric spaces, where the choice of metric significantly impacts comparison quality. In low- dimensional datasets like Shape3D (Kim & Mnih, 2018) and dSprites (Matthey et al., 2017), Euclidean pixel-level distances align well with semantic differences, leading prior works (Nakagawa et al., 2023; Uscidda et al., 2024) to use 1.2 or cosine distances for isometric mappings. However, this breaks down in high-dimensional data, like real-world images, which lie on lower-dimensional manifolds. The curse of dimensionality causes traditional metrics, such as pixel-wise distances, to lose effectiveness as dimensionality increases. Recent advancements in vision-language models like CLIP (Radford et al., 2021) have shown their ability to learn robust and expressive image representations by train- ing on diverse data distributions (Fang et al., 2022). Studies (Yun al., 2023) demonstrate that CLIP captures meaning- ful semantic relationships, even learning primitive concepts. Therefore, we propose using the semantic embedding space of pre-trained CLIP models as a more effective metric for computing distances between datasets, which we define as the Semantic Preservation (SP) loss. For a detailed evalu- ation of the improvements from using CLIP embeddings, please refer the Appendix G, which includes demonstra- tions and additional results. In the following section, we will denote the Wasseerstin constraint as GW-EP, and GW-SP to differentiate the metric space we used for Gromov- Wasseertsnn constraint as Euclidean metric space Preservation (EP) and Semantic Structural Preservation (SP) respectively.
Prior works, such as Truong et al. (2022): Carrasco et al. (2024), although uses GW distance as part of the loss in the the objective but is focusing on calculating and minimizing the GW objective in the embedding space between domains /GW/Grrrre> Ztgt). On the other hand, Us- cidda et al. (2024); Nakagawa et al. (2023) defines the GW objective as being calculated between the data dimension and the embedding dimension.
Lor/ow
5. Experiments
In this section, we evaluate the effectiveness of our proposed VAUB with a score-based prior on several tasks. We conduct experiments synthetic data, domain adaptation, multi- domain matching, fair classification, and domain translation. For each experiment, we compare our methods to VAUB and other baselines and evaluate performance using various metrics.
5.1. Improving Latent Space Separation by Using Score-based Prior
The primary objective of this experiment is to demonstrate the performance of different prior distribution models within
6

---

Expressive Score-Based Priors for Distribution Matching with Geometry-Preservingg Regularization
Gaussian
MoG
VampPrior
LSGM
Ours
Original
DI
D2
D2_L1
D2
22 1-.00
The dataset consists of domains: Domain (left nested 'D-shaped') and Domain (right flipped 'D-shaped'). In each domain, the outer 'D' corresponds to Label I, and the inner 'D' to Label 2. The shared latent spaces are visualized for models trained with varying data sizes (n = 20, 100, 500 samples) using Gaussian( Gaussian(Kingma Sampprior[Tomezat & Welling, 2018), (Son, et al., 2021c) and our score-based model (columns). Legends follow the format /ddomain index) Ll Labell index)
Figure 2.
al., 2019), Mixture of Gaussians(Gong et al., 2024),
the VAUB framework. Additionally, we examine the ef- feet of varying the number of samples used during train- ing, specifically considering scenarios with limited dataset availability. To achieve this, we create synthetic nested D-shaped dataset consists of two domains and two labels, as illustrated in Fig. 2. The aim is to learn a shared latent representation across two domains and evaluate the degree of separation between class labels within this shared la- tent space. Since downstream tasks rely on these shared latent representations, better separation of class labels in the latent space naturally leads to improved classification performance. This setup draws an analogy to domain adap- tation tasks, where the quality of separation in the latent representation relative to the label space plays a critical role in determining downstream classification outcomes.
der transformation with score- based prior methods (LSGM and ours), resulting in well-separated latent representations across different classes. This holds consistently true for varying numbers of data points, from as low as 20 samples to higher counts. On the other hand, the Gaussian prior, MoG and Vampprior only achieves 90% of separation the latent space when the number data samples sufficiently large (n 100 for MoG and Vampprior prior and n = 20 for Gaussian prior), allowing the inner and outer classes to have classifier bound supported by enough data points as shown in Fig. 3. This finding is especially relevant for real-worke datasets, where the original data dimensionality can easily reach up to tens of thousands; while in this ex- periment, we worked with only a wwo-dimensional dataset, yet the Gaussian, MoG and Vampprior required more than hundreds of samples to achieve effective latent separation, whereas the score-based prior (LSGM and SFS) succeeded with few 20 samples.
In this experiment, we control the total number of data sam- ples generated for the dataset, and compare the model's performance using five types of priors: Gaussian prior, Mix- ture of Gaussian Prior (MoG), Vampprior, and a score-based prior trained with LSGM, and ours (SFS method). Consider- ing the strong relations between point-wise distance and the label information of the dataset, we use GW-EP to compute the constraint loss in both in the data domain and the latent domain. This helps to better visually reflect the underly- ing structure and separations in the latent space. As shown in Fig. 3, this performance improvement is evident in the latent space: the nested D structure is well-preserver un-
5.2. Improving the Tradeoff between Accuracy and Parity in Fair Classification
For this experiment, we apply our model to the well-known Adult dataset, derived from the 1994 census, which contains 30K training samples and 15K test samples. The target task is to predict whether an individual's income exceeds $50K, with gender (a binary attribute in this case) considered as the protected attribute.

---

Expressive -core-Based Priors for Distribution Matching with Geometry-Preservinng Regularization
Choices of Different Prior under Different Data Sizes
Tradeoff between Aop and Accuracy
0.25
1.0
CVIB
LAFTR-DP
KCR
0.9
0.20
VAUB
Ours w/o GW
AUROC (T)
0.8
Ours w/ GW-EP
0.35
Unfair Classifier
0.7
Gaussian Prior
MoG Prioe
Ann
0.10
0.6
VampPtior
L.GGM Prior
Prior
0.05
0.5
20
50
100 Number data points
200
500
0.00
Figure 3. This figure shows label separation in the latent space under varying sample sizes and prior configurations, quantified by AUROC scores from the prediction of support vector classitier. Higher scores indicate better separation. Details of the metric are described in the appendix.
0.72
0.74
0.76
0.78
0.80
0.82
0.84
0.86
Accuracy (T)
Figure 4. Demographic Parity gap (ADP) vs. Accuracy trade-off for UCI Adult dataset. Lower App is better, and higher Accuracy is better.
We adopt the same preprocessing steps in Zhao et al. (2020), and the encoder and classifier architectures are consistent with those in Gupta et al. (2021). We additionally adapt GW-FP as our constraint loss considering the lack of seman- tic models in tabular dataset such as Adult dataset. Please refer to Appendix I for more detailed architecture setup. For comparison, we benchmark our model against three non- adversarial models FCRI. al., 2021), CVIBeMoyer et al., 2018), VAUB(Gong al., 2024) and one adversar- ial model LAFTR-DP(Madras et al., 2018) and one extra baseline 'Unfair Classifier' which is obtained to serve as a baseline, computed by training the classifier directly on the original dataset.
Table 1. Domain adaptation accuracy (%).
Model
MNIST-+USPS
USPS MMNST
ADDA
89.4
90.1
DANN
77.1
73.0
VAUB
40.7
45.3
Ours w/o GW
88.1
85.5
Ours w/ GW-EP
91.4
92.7
Ours w/ GW-SP
96.1
97.4
without using target labels. We compare our SAUB method (with and without structure- prrnettvrppreeriinggg constraints) against baseline DA methods: ADDA (Zhao et al., 2018), DANN (Ganin et al., 2016), and VAUB (Gong et al., 2024). All methods use the same encoder and classifier architecture for faimesss with structure preserie-preserving constraints applied using L2 distance in Euclidean space (GW-EP) and CLIP embedding (GW-SP).
As illustrated in Fig. 4, our method not only retains the advantages of the SAUB method, achieving near-zero demo- graphic parity (DP) gap while maintaining accuracy, but it also improves accuracy across the board under the same DP gap comparing to other methods. We attribute this improve- ment largely to the introduction of the score- score-based based prior, which potentially allows for better semantic preservation in the latent space, enhancing both accuracy and fairness.
As shown in Table 1, our method outperforms the baselines in both directions. Unlike ADDA and DANN, which require joint classifier and encoder training, our approach allows for classifier training after the encoder is learned, simplifying domain adaptation. Additionally, the inclusion of a decoder enables model to naturally adapt to domain translation tasks, as demonstrated in Fig. 155. We additionally conduct novel experiments to assess the generalizability and robust- ness of our model with limited source-labeled data, detailed in Appendix D. Additionally. image translation results be- tween MNIST and USPS are presented in Appendix K.
We further provide comprehensive ablation studies and com- putational efficiency analyses (detailed in Appendix N) vali- date the effectiveness of our GW regularization component and demonstrate significant computational advantages over LSGM, particularly for high- nigh-dimmnsional applications.
5.3. Domain Adaptation
We evaluate our method on the MNIST-USPS domain adap- tation task, transferring knowledge from the labeled MNIST
(70,000 images) to the unlabeled USPS (9,298 images)
8

---

Expressive Score-Based Priors for Distribution Matching with Geometry-Preservinng Regularization
5.4. Domain Translation
Table 2. Retrieval and perceptual similarity metrics, Higher SSIM (1) and lower LPIPS (L) indicate better structural and perceptual similarity.
Translation
Trashation
Toastation
Recen
Reema
Recan
GT
GT
Image Retrieval (%)
Perceptual Similarity
Task/Model
Top-1
Top-5
Top-10
Top-20
SSIM LPIPS (1)
Black to Blonde Hair
Blacn-to-Bloude Hair
GW
50+ 1.4
14
24.4440
40: 3.5
0.395
0-431
GW-EP
1.
116=22
22.02
35.0
0,128
0.371
GW-SP
9.116 278881
4.22442
59.0
0.542
Monde-to-Black
Hair
No GW
3.1
108=33
19.0
$3.4
0.393
0.380
GW-EP 2.7:0.7
922188
158422
3044
0.426
0.385
GW-SP
2.3
18.8
28.6
4.1
48.2
0.532
0.282
females with blonde hair and vice versa. For each translated image, we compute the cosine similarity with all translated images in the target domain using CLIP embeddings. To ensure fairness, we use a different pretrained CLIP model for evaluation and for training GW-SP: for more informa- tion, see Appendix I. This process is repeated five times with randomly selected datasets to account for variability in the data. The experiment aims to measure how well the translated images preserve their semantic content. We compute the top-k accuracy, where the task is to retrieve the correct translated image from the set of all translated images. For SSIM and LPIPS, we randomly translate 1,000 images and compare their similarity metrics to quantify structural and perceptual consistency. This bidirectional evaluation black-to-blonde and blonde-to-black ensures ro- bustness and highlights the model's ability to maintain se- mantic consistency during translation. GW-SP in semantic space consistently improves accuracy for all metrics. No- tably, GW-EP performs worse than no GW loss for image retrieval. The domain translation images in Appendix M confirm that models with semantic space GW loss better preserve semantic features like hairstyle, smile, and facial structure, demonstrating its advantage. For additional ex- periments, we provide image translations between male and female subjects on the FairFace dataset in Appendix L for interested readers.
(b) Blande to Black Hair
w/o GW
GW-EP (Euclideun space)
GW-SP (semantic space)
Figure 5. All models use the same architecture. Refer to Ap- pendix for details on the neural network and CLIP model. Apply- ing GW loss in the CLIP semantic space shows superior semantic preservation in both (a) and (b). The samples are selectively chosen represent diverse variations; random samples are in Appendix M.
We conduct domain translation experiments on the CelebA dataset, translating images of females with blonde hair to black hair and vice versa. We compare three settings: GW loss in semantic space, GW loss in Euclidean space, and no GW loss. This comparison shows that GW loss in the semantic space better preserves semantic features, while Fu- clidean space GW loss is less effective in high-dimensional settings. We note that achieving state-of-the art image trans- lation performance is not our primary objective, since we employ relatively simple networks for both the VAE and diffusion model (see Appendix I). Instead, this experiment demonstrates our model's versatility across tasks and serves as a proof of concept. We believe that with -tate-of-the-art architecture and engineering design our approach will be competitive for domain translation and other imaging tasks, which we leave to future work.
6. Conclusion
In conclusion, we introduce score-based priors and structure- preserving constraints to address the limitations of tradi- tional distribution matching methods. Our approach uses score models to capture complex data distributions while maintaining geometric consistency. By applying Gromov- Wasserstein constraints in the semantic CLIP embedding space, we preserve meaningful relationships without the computational cost of expressive priors. Our experiments demonstrate improved performance in tasks like fair classi- fication, domain adaptation, and domain translation.
For quantitative evaluation of semantic preservation, we utilize Structural Similarity Index Metric (SSIM), Learned Perceptual Image Patch Similarity (LPIPS), and image re- trieval accuracy as our metrics. The models, trained for 1,500 epochs, for image retrieval translate images from a domain of 100 females with black hair to domain of 100
9

---

Expressive Score-Based Priors for Distribution Matching with Geometry-Preserving Regularization
Acknowledgment
Learning Research, pp. 3029-3039. PMLR, 13-18 Jul 2020. pttectts://prooceeeiie v v119/farnia20: html.
press/
Z.G., J.L. and D.I. acknowledge support from NSF (IIS- 2212097), ARL (W911NF-2020-221) and ONR (N00014- 23-C-1016).
Feydy, J., Séjoumé, T., Vialard, F.-X., Amari, S.-i., Trouvé, A., and Peyré, G. Interpolating between optimal trans- port and mmd using sinkhorn divergences. In The 22nd International Conference on Artificial Intelligence and Statistics, pp. 2681-2690. PMLR, 2019.
Impact Statement
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
Ganin, Y., Ustinova, E., Ajakan, H., Germain, P., Larochelle, H., Laviolette, F., Marchand, M., and Lempitsky, V. Domain-adversarial training of neural networks. The journal of machine learning research, 17(1):099-2030. 2016.
References
Burgess, C. P., Higgins, I., Pal, A., Matthey, L., Watters, N., Desjardins, G., and Lerchner, A. Understanding disentan- gling in 8-vae, 2018. URL attps://arxiia obs/1804 03599.
Gong. Z.. Usman, B., Zhao, H., and Inouye, D. I. To- wards practical non-adversarial distribution matching. In International Conference Artificial Intelligence and Statistics (AISTATS), May 2024.
Carrasco, X. A., Nekrashevich, M., Mokrov, P., Burnaev, E., and Korotin, A. Uncovering challenges of solving the continuous gromov- problem, 2024. URL https: /wwpvvaavvvvvvv///vvv/eeeeeeawaa abtgg//a///// 2303. 05978.
Goodfellow, I. Abadie, J., Mirza, M., Xu, B., Warde-Farley, D., Ozair, S., Courville, A., and Bengio, Y. Generative adversarial nets. Advances in neural informa- tion processing systems, 27, 2014.
Chen, N., Klushyn, A., Ferroni, F., Bayer, J., and Van Der Smagt, P. Learning flat latent manifolds with vaes. arXiv preprint arXiv:2002.04881, 2020.
Gropp, A., Atzmon, M., and Lipman, Y. Isometric autoen- coders. arXiv preprint XXiv:2006.09288) 2020.
Chen, R. Q., Li, X., Grosse, R., and Duvenaud, D. Isolating sources of disentanglement in variational toencoders, 2019, URL https: / //axxvivvvvv///////// orgoorgggoo /bs/ 1802.04942
Gupta, U., Ferber, A. M., Dilkina, B., and Ver Steeg, G. Controllable guarantees for fair outcomes via contrastive information estimation. In Proceedings of the AAAI Con- ference on Artificial Intelligence, volume 35, pp. 7610- 7619, 2021.
Chen, T., Sun, Y., and Yin, W. Closing the gap: Tighter anal- ysis of alternating stochastic gradient methods for bilevel problems. Advances in Neural Information Processing Systems, 34:25294-25535 2021.
Hahm, J., Lee, J., Kim, S., and Lee, Isometric represen- tation learning for disentangled latent space of diffusion models. arXiv preprint arXiv:2407. 1451, 2024.
Cho, W., Gong, Z., and Inouye. D. I. Cooperative distribu- tion alignment via jsd upper bound. In Neural Information Processing Systems (NeurIPS), dec 2022a.
Han, X., Chi, J., Chen, Y., Wang, Q., Zhao, H., Zou, N., and Hu, X. FFB: A Fair Fairness Benchmark for In- Processing Group Faimesss Methods, 2023.
Cho, W., Gong, Z., and Inouye, D. Cooperative Distri- bution Alignment via JSD Upper Bound, 2022b. URL https: /wwx/aiivvvv./////v www.xiv.org/eeeee www.obgg/bs... 220..220..20 02286.
Hoffman, M. D. and Johnson, M. J. Elbo surgery: yet another way carve up the variational evidence lower bound. In Workshop in Advances in Approximate Bayesian Inference, NIPS, volume 1, 2016.
Fang, A., Ilharco, G., Wortsman, M., Wan, Y., Shankar, V., Dave, A., and Schmidt, L. Data determines distri- butional robustness contrastive language image pre- training (clip), 2022. URL https://arxiv.... fbs/2205.01399
Horan, D., Richardson, E., and Weiss, Y. When is unsu- pervised disentanglement possible? Advances in Neural Information Processing Systems, 34:5150-5161, 2021.
Farnia, F. and Ozdaglar, A. Do GANs always have Nash equilibria? In III, H. D. and Singh, A. (eds.), Pro- ceedings of the 37th International Conference on Ma- chine Learning, volume 119 of Proceedings of Machine
Johnson, M. J., Duvenaud, D., Wiltschko, A. B., Datta, S. R., and Adams, R. P. Composing graphical models with neural networks for structured representations and fast inference, 2017.
10

---

Expressive Score-Based Priors for Distribution Matching with Goometry-Preserviing Regularization
Kaplan, J., McCandlish, S., Henighan, T., Brown, T. B., Chess, B., Child, R., Gray, S., Radford, A., Wu, J., and Amodei, D. Scaling laws for neural language mod- els, 2020. URL https::///// httppp////axxx..///aaabaaaaae//a 08361.
2013, PMLR. URL htwpt: httptttt//prroeeedingg .mlr.
press/v 1222//mmanet113.ht.
Nakagawa, N., Togo, R., Ogawa, T., and Haseyama, M. Gromov-wasserstenn autoencoders, 2023. URL https: //arxiv. war.ii.oog/absaa......oooa. www.://ggg/..//.o/.//./o www.aa////22.eooo 220.//2777000007..
Kim, H. and Mnih, A. Disentangling by factorising. In Dy,
J. and Krause, A. (eds.), Proceedings of the 35th Inter- national Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pp- 2649-2658. PMLR, 10-15 Jul 2018. URL https:// proceedings ngs mlr prwss vv0 / 8b html.
Nie, W. and Patel, A. B. Towards a better understanding and regularization of gan training dynamics. In Uncertainty in Artificial Intelligence, pp. 281-291. PMLR, 2020.
Papamakarios, G., Nalisnick, E. T., Rezende, J., Mo- hamed, S., and Lakshminarayanan, B. Normalizing flows for probabilistic modeling and inference. J. Much. Learn. Res., 22(57):1-64, 2021.
Kingma, D. P., Welling, M., et al. An introduction vari- ational autoencoders. Foundations and Trends@ in Ma- chine Learning, 12(4):307-392, 2019.
Poole, B., Jain, A., Barron, T., and Mildenhall, B. Dream- fusion: Text-to-3d using 2d diffusion. arXiv preprint arXiv:2209.14988 2022.
Kurach, K., Lucic, M., Zhai, X., Michalski, M., and Gelly, S. The GAN landscape: Losses, architectures, reg- ularization, and normalization, 2019. URL https: //tp///r///ree wwwenrevvewweeeeeeee wtt.i//wwwe///ooummmm f d= d
Radford, A., Kim, J. W., Hallacy, C., Ramesh, A., Goh, G., Agarwal, S., Sastry, G., Askell, A., Mishkin, P., Clark, J., Krueger, G., and Sutskever, I. Learning transferable visual models from natural language supervision, 2021, URL https: aawxiv org/ atbp//101.00022
Lee, Y., Yoon, S., Son, M., and Park, F. C. Regularized autoencoders for isometric representation learning. In International Conference on Learning Representations, 2022.
Song, J., Meng, C., and Ermon, S. Denoising diffusion implicit models, 2022. URL httpss//:/axxi abs/200 02502.
Louizos, C., Swersky K., Li, Y., Welling, M., and Zemel, R. The variational fair autoencoder. arXiv preprint 1rXiv:1511.000 2015.
Song, Y. and Ermon, S. Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems, 32, 2019.
Lucic, M., Kurach, K., Michalski, M., Gelly, S., and Bous- quet, O. Are gans created equal? a large-scale study. In Advances in neural information processing systems, pp. 700-709, 2018.
Song, Y., Durkan, C., Murray, and Ermon, S. Maxi- mum likelihood training of score-based diffusion models. Advances in neural information processing systems, 34: 1415-1428, 2021a.
Madras, D., Creager, F., Pitassi, T., and Zemel, R. Learning adversarially fair and transferable representations. In International Conference on Machine Learning, pp. 3384- 3393. PMLR, 2018.
Song, Y., Durkan, C., Murray, I,, and Ermon, S. Maxi- likelihood training of score-based diffusion mod- els, 2021b. utthttps: arxiv. org/ abs /201. 09258.
Makhzani, A., Shlens, J., Jaitly, N., Goodfellow, I,, and Frey, B. Adversarial autoencoders, 2016.
Matthey, L., Higgins, I., Hassabis, D., and Lerchner, A. dsprites: Disentanglement testing sprites dataset. 2017.
Song, Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Ermon, S., and Poole, B. Score-based generative mod- cling through stochastic differential equations. In In- ternational Conference Learning Representations, 2021c. URL https:/ URR.httpss//opeeerriie net/ forum? id-PxTIG12RRES.
Moyer, D., Gao, S., Brekelmans, R., Galstyan, A., and Ver Steeg, G. Invariant representations without adversar- ial training. Advances in Neural Information Processing Systems, 31, 2018.
Spirtes, P. and Zhang, K. Causal discovery and inference: concepts and recent methodological advances. In Applied informatics, volume 3, pp. 1-28. Springer, 2016.
Muandet, K., Balduzzi, D., and Schölkopf, B. Domain generalization via invariant feature representation. In Dasgupta, S. and McAllester, D. (eds.), Proceedings of the 30th International Conference on Machine Learn- ing, volume 28 of Proceedings of Machine Learning Re- search, pp. 10-18, Atlanta, Georgia, USA, 17-19 Jun
Tomczak, J. and Welling, M. Vae with a vampprior. In International conference on artificial intelligence and statistics, pp. 1214-1223, PMLR, 2018

---

Expressive Score-Based Priors for Distribution Matching with Geometry-Preserving Regularization
Truong, T.-D., Chappa, N. V. S. R., Nguyen, X. B., Le, N., Dowling, A., and Luu, K. Otadapt: Opti- mal transport-based approach for unsupervised domain adaptation, 2022. URL https::///// / /awpiv.xivv.////a///ee////aaall/ org/ abs/ 2205. 105..10333 10738.
adaptation. Advances in neural information processing systems, 31, 2018.
Zhao, H., Coston, A., Adel, T., and Gordon, G. J. Conditional learning of fair representations. In In- lernational Conference on Learning Representations, 2020. URL https:ttp///o..... htttps://openreiltt net/ forum?
Uscidda, T., Eyring, L., Roth, K., Theis, F., Akata, Z., and Cuturi, M. Disentangled representation learning with the gromov-monge gap, 2024. URL
Ld-HkeklONEPH
org/abs 2407, 2707707879.....
Vahdat, A., Kreis, K., and Kautz, Score-based generative modeling latent space. Advances neural information processing systems, 34:1 12s7-188--1302 2021.
Vincent, P. A connection between score matching and de- noising autoencoders. Neural Computation, 1367):1661- 1674, 2011. doi: 10.1162/NECO a 00142.
Wu, Y., Zhou, P., Wilson, A. G., Xing, E., and Hu, Z. Improving gan training with probability ratio clipping and sample reweighting. In Larochelle, H., Ranzato, M., Hadsell, R., Balcan, M. F,, and Lin, H. (eds.). Advances in Neural Information Processing Systems, volume 33, pp. 5729-5740. Curran Asso- ciates, Inc., 2020. URL https: httttp::/proccedings. neurips. hctps//cecppppe/2////1eee.. 3eb4 2ebf66555d3333553366ffl....... Paper.. pdf.
Xiao, Q., Lu, S., and Chen, T. An alternating optimization method for bilevel problems under the polyak-tojasiewiez condition. Advances in Neural Information Processing Systems, 01463847-638777 2023.
Yun, T., Bhalla, U., Pavlick, E., and Sun, C. Do vision- language pretrained models learn composable primitive concepts?, 2023. URL https: org/ abs/ 2203. 17271.
Zellinger, W., Grubinger, T., Lughofer, E., Natschläger, T., and Saminger-Platz, S. Central moment discrepancy (cmd) for Nomain-invarian representation learning. International Conference on Learning Representations, 2017.
Zemel, R., Wu, Y., Swersky, K., Pitassi, T., and Dwork, C. Learning fair representations. In Dasgupta, S. and McAllester, D. (eds.), Proceedings of the 30th Interna- tional Conference on Machine Learning, volume 28 of Proceedings of Machine Learning Research, pp. 325- 333, Atlanta, Georgia, USA, 17-19 Jun 2013. PMLR. URL wethhttps::t. :ttttps://proceo pttpp//tcccce....cpecpp mlr preas/v28. zeme 113 html.
Zhao, H., Zhang, S., Wu, G.. Moura, M., Costeira, P., and Gordon, J. Adversarial multiple source domain
12

---

Expressive Score-Based Priors for Distribution Matching with Geometry-Preserving Regularization
A. Proof of Proposition 3.1
Proposition (Score Function Substitution (SFS) Trick)
qe((2 (2(2z) is the posterior distribution parameterized by 0, and Q (z) is the prior distribution parameterized by th then the gradient of the cross entropy term be written as:
zolozzrg(( |- log (z))) Va VelEso
Z
V log Q (3) 8=20
(13)
where the notation of ze emphasizes its dependence on and 1|s=so denotes that while is equal to 26, it treated as constant with respect 0.
Proof.
go log (z))
(14)
Vple | log (c)(
(Reparameterizztioon trick: zg go
(15)
Bemple)
V (- log Q (gn)
(16)
Ogo(c) Bog (3)
Beep(()
(Chain rule: differentiating at go(e))
(17)
00
Oz
(e)
i log (3)
Bemp(e)
Vogo(c)
(Simplify notation)
(18)
Oz
[=)
Ölog (z)
Benple) V
go(e)
(Move V outside)
(19)
0z
E=go(e)
Veleep((e)
v log Qu (g)(
ga(c)
(Gradient applied to parts dependent on 0)
(20)
x=po(e))
- V= Qx (
Val (0
{=20))
(Change back to after pulling out gradient)
(21)
B. Pseudo-code for learning VAUB with Score-Based Prior
See Alg. 1.
13

---

Expressive Score-Based Priors for Distribution Matching with Geometry-Preservinng Regularization
Algorithm 1 Training VAUB with Score-based Prior (Alternatin Optimization)
Input: Data z, domain parameters (0d, iPu, ) hyperparameters: noise levels {Tmin m number of loops L for
score model update
2: Initialize: Parameters of Encoders 0, Decoders sp, and Score model vé
3: while not converged do
4: Step 1: Update Encoder and Decoder parameters {0, [p]
5: Draw z, d (F, d)
6: Draw ~ (1(2 d)
7:
Compute score using
where =" detached from the computational graph
8:
Compute the objective in Equation 8
9:
Perform gradient descent to minimize the objective and update {8, yp)
10:
Step 2: Update Score Model parameters v
11:
for loop = to L do/ // Number of loops for score model update
12:
Draw x, ~ Pdata (I, d)
13:
Draw 2 ~ (2 d)
Draw perturbed latent variable go, (3)=), where o; [oin mmax]
Compute the DSM loss for the score model Equation 9
16:
Perform gradient descent to minimize the DSM objective and update ep
17:
end for
18:
Repeat alternating optimization steps until convergence.
19:
end while
C. Stabilization and Optimization Techniques
Several factors, such interactions between the encoder, decoder, and score model, well as the iterative nature of the optimization process, can introduce instability. To mitigate these issues, we implemented stabilization and optimization techniques to ensure smooth and robust training.
Batch Normalization on Encoder Output (Without Affine Learning) Applying batch normalization to the encoder's mean output without affine transformations facilitates smooth transitions in the latent space, acting as a soft distribution matching mechanism. By centering the mean and mitigating large shifts, it prevents disjoint distributions, allowing the score model keep up with the encoder's updates. This regularization ensures the latent space remains within regions where the score model is trained, enhancing stability and reducing the risk of divergence.
Gaussian Score Function for Undefined Regions: further stabilize training, we incorporate small Gaussian score function into the score model handle regions beyond the defined domain of the score function (i.e., outside the maximum noise level, mmax). Inspired by the mixture neural score function in LSGMs (Vahdat et al., 2021), this approach blends score functions to address nuitr--distriuution latent samples. The Gaussian score ensures smooth transitions and prevents instability in poorly defined areas of the latent space, maintaining robustness even in undertrained regions of the score model.
Weight Initialization and Hyperparameter Tuning: We observed that the initialization of weights significantly impacts the stability and convergence our model. Poor initialization can lead to bad alignment. Therefore, grid search was used to find an optimal weight scale.
D. Limited Source Label for Domain Adaptation
We introduce, to the best of our knowledge, novel downstream task setup where there is limited labeled data in the source domain (i.e., 1%, 5%, 10%) and no supervision the target domain. We apply this setup to the MNIST-to-USPS domain adaptation task. The objective is to determine how well our model with and without structural preservation can generalize with limited source supervision.
14

---

Expressive Score-Based Priors for Distribution Matching with Geometry-Preserving Regularization
E. Prior and Target Distribution for Comparative Stability: SFS vs. LSGM
Prior Distribution (2D): Projected the Z-plane
Target Distribution (3D)
03
-05
-15
00
0.5
1.0
Figure 6. Target/Prior Distribution for the stability analysis in subsection 3.3. The prior distribution is the target distribution projected onto the Z-space.
E.1. Results
As shown in Figure Fig. 7, our method without the SP constraint (which is entirely unsupervised in the source domain) demonstrates remarkable sample efficiency. With as little 0.04% of the dataset (roughly two images per class), our method achieves an accuracy of around 40%. By increasing the labeled data to just 0.1% (about five images per class), the accuracy surpasses 73%. When we introduce the structural preservation constraint, which allows the model to transfer knowledge from a pretrained model, we observe a significant improvement in performance. With only 0.2% of the labeled data, the model's s accuracy approaches the performance of models trained on the full dataset. This boost in performance shows the effectiveness of incorporating semantic information into the latent space, allowing the model generalize better with minimal supervision.
The performance gap between models with and without the structural preservation (SP) constraint becomes more evident through UMAP visualizations of the latent space (Figure Fig. 7). While both methods achieve distribution matching and show label separation, the model without SP struggles to distinguish structurally similar digits, such as "4" and "9". In contrast, with the SP constraint, the latent space exhibits clearer, distinct separations, even for similar digits. The semantic structure injected by the SP constraint leads to more robust and meaningful representations, helping the model better differentiate between challenging classes. This highlights the effectiveness of the SP constraint in refining latent space organization.
F. More Detailed Discussion of Gradient Comparison Between LSGM and SFS Trick
Below, we detail the encoder and decoder optimization objectives for LSGM:
(w(t) |le cz(
min (z 0.40
[- log (z) + (zolz))) [log go (=(2)=))) + (xcc
2
where wt) is a weighting function,,p,f represents a diffusion model, and c ~ (O, 7). Similar to our loss objective (refer to Eqn. 8), LSGM substitutes the traditional cross-entrope term with a learnale neural network prior. Specifically, the final term the Evidence Lower Bound (ELBO) is replaced with a weighted denoising score matching objective.
15

---

Expressive Score-Base Priors for Distribution Matching with Ecometry-Preserve Regularization
Aocurecy
Accuracy -0
Scole Jannd Prir w/
Sc2 Janed Pree Score
Score Socod wi
Farcntaage Source Data Label Used 159
Forcrtage Source Data Used Pio
Figure 6. (a) MNIST USPS
Figure 6. (b) USPS MNIST (LSDA)
MNIST to USPS
UMAP MNIST USPS SP
Chear
Class
Cham
(hre
Clan
UNAAP Duncmios
10
Cla
Donai:
UNSAT
Class
-15
SIMAP Dimeradion
UMAP Dimonsion
Figure 6. (c) UMAP with SP
Figure 6. (d) UMAP without SP
Figure 7. (a) MNIST USPS (LSDA). (b) USPS to MNIST (LSDA). (c) UMAP with SP. (d) UMAP without SP, All labeled data randomly selected from the source dataset and tested on the target dataset, with results averaged over 10 trials. Both (a) and (b) demonstrate with SP loss, the model more robust to limited data. This further supported by the corresponding UMAP visualizations, where (c) shows larger separation between classes compared to (d), reflecting better class distinction.
We first adapt notations used in our objective for easy readability during comparison. Diffusion model can approximate the denoising score function by rewriting (2, t) = o2 oye, oe) (Song al., 2022). To streamline discussion and avoid repetition, we will refer the final term of this formulation as the LSGM objective, LSGM as below with the weighting function w(()= 9(t)²/oz /a," which maximizes the likelihood between the encoder posterior and the prior where g(-) is the diffusion coefficient typically proportional to the variance scheduling function (Song et al., (Vahdal et al., 2021).
write the cross-entropy term of
v(())
LISGM = By (o)
||e
6 (2,, 1)1
(22)
2
9(1)²
Re, (sla)
Su = z + ore,
(23)
2
ae
During encoder updates, the gradient computation for the last term with respect to the encoder parameters expressed as:
((t) (
Sp((200)) OS (2, ot) ) 0E
VOLISOM -Eg,, (=11)
(24)
o
08
00
This framework requires computing the Jacobian term intensive, To mitigate this, the Score Function Substitution (SFS) trick eliminates the need for Jacobian computation by
as (5.0).
0z which is both computationally expensive and memory-
16

---

Expressive Score-Based Priors for Distribution Matching with Geometry-Preserving Regularization
detaching the latent input =" in the function from the encoder parameters. The resulting gradient is expressed as:
(-S 0)
8z
VoLsrs (2
(25)
08
This modification provides significant advantages, reducing memory usage by bypassing the computational graph of the diffusion model's U-NET and enhancing stability. Poole et al. (2022) highlighted that the Jacobian computation approximates the Hessian of the dataset distribution, which is particularly unstable at low noise levels. Our empirical results in Fig. confirm these findings, demonstrating improved stability with our loss objective compared to LSGM.
17

---

Expressive Score-Based Priors for Distribution Matching with Geometry-Prsservingg Regularization
G. Choices of different metric spaces in different dataset
Histogram of distances within and between classes in MNIST
Histogram distances within and between classes USPS
2000
2500
AUROC: 0.856
withn-ciaass
betwoen-dlass
1000
1200
AUROC: 0.738
witin-cless
between-class
1400
1090
1902
AUROC:
between class
within class
3909
AUROC: 0.842
between cless
within-class
3000
600
1000
400
200
0.2
Cife
0.2
6-4
61
0.6
0.2
0.4
0.5
0.5
022
0.44
0-5
Ours: Sementic Space (CLIP)
Prior Works: Pixel Space (Euclidean)
Ours: Semantic Space (CLIP)
Prior Works: Pixel Space (Euclidean)
Histogram of distances within and between classes in CelebA
within-class
35
within-class between-class
35
between-class
30
30
AUROC: 0.732
AUROC: 0.479
25
25
20
20
15
15
10
10
Ours: Semantic Space (CLIP)
Prior Works: Pixel Space (Euclidea)
Figure& 8. Histogram of the pairwise distance between data samples within a class and between different classes for three datasets: MNIST, USPS. and CelebA. The amount of separation of two histogram is computed by using the AUROC score which being measured by a binary classifier to distinguish between with-in class results and between-class results. The class considered in MNIST and USPS is the digits, and in CelebA is hair color.
From the graph, we observe that for the MNIST and USPS datasets, both the Euclidean pixel space metric and the semantic space metric can effectively separate data pairs into within-class or between-class categories. However, the semantic space metric demonstrates a higher AUROC separation score, indicating that it provides a more reliable metric for distinguishing between these pair types.
In contrast, for the CelebA dataset, relying solely on pixel-based Euclidean distances struggles to differentiate whether the paired distances belong to within-class or between-class data pairs. By employing a semantic metric, such as the one derived from CLIP, a clear distinction emerges, underscoring its utility.
These observations highlight that while pixel space metrics like Euclidean distance may be useful for certain datasets, semantic distance metrics, when available, often offer superior performance and may even be essential for datasets with more complex structures features.
H. Multi-Domain Distribution Matching Setting
We train SAUB with SP on three different MNIST rotation angles: 90,30,,00 60 The top row is the ground truth image, the second row is the reconstruction, the third row is translation to MNIST 30 and last row is translation to MNIST 60 in Fig. Qualitatively most of the stylistic and semantic features are preserved with the correct rotation.
18

---

Expressive Score-Based Priors for Distribution Matching with Geometry-Preserving Regularization
I. Detailed Architecture of the model
I.1. Fairness Representation Learning
The encoder is a 3-layer MLP with hidden dimension 64, and latent dimension 8 with ReLU layers connecting in between. The classifier a 3-layer MLP with hidden dimension 64 with ReLU layers connecting in between.
1.2. Separation Metric for Synthetic Dataset
The classifier trained by a support vector where hyperparmeters are chosen from the list 'C": [0.1, 1, 10, 100), 'gamma": [1, 0.1, 0.01, 0.001] with 5-fold cross validation. Error plot generated from runs.
1.3. Domain Adaptation VAE Model
ENCODER ARCHITECTURE
The encoder compresses the input image x E R21x2xx88 into a latent representation. The architecture consists of the following layers:
Conv2D: 4, stride 2, 16 channels (input size 28 28 + 14 x 14).
Residual Block: 16 channels.
Conv2D: 4 4, stride 2, 64 channels (input size 14 14 7 7).
Residual Block: 64 channels.
Conv2D: 3 X 3, stride 2, 2 x latent size channels (input size 7 x + x
Residual Block: 2 x latent size channels.
Conv2D: 4 stride 1, latent size channels (output size 4 4
Split into two branches for μ and log o², each with latent size channels.
DECODER ARCHITECTURE
The decoder reconstructs the input image x' E RX128x28 from the latent representation. The architecture consists of the following layers:
Reshape: Latent vector reshaped to size (latent size, 1, 1).
Residual Block: latent size channels.
Con/Transpose2DD: 4 4, stride 1, 64 channels (output size 1 4 4).
Residual Block: 64 channels.
Conv Transsranppsee2D: 4 x 4, stride 2, 16 channels (output size 4 x 4 + 8 x 8).
Residual Block: 16 channels.
Conv Transsaaspsse2:::: 4 4, stride 4, channel (output size 8 - 28 28).
1.4. Domain Translation VAE Model
ENCODER ARCHITECTURE
The encoder compresses the input image x p3x61x61 into a latent representation. The architecture consists of the following layers:
Cony2D: stride 2, channels (input size 64 64 + 32 32).
19

---

Expressive Score-Based Priors for Distribution Matching with Geometry-Preserving Regularization
Residual Block: 64 channels.
Conv2D: 3 x 3, stride 2, 128 channels (input size 32 x 32 + 16 x 16).
Residual Block: 128 channels.
Conv2D: 3, stride 2, 256 channels (input size 16 x 16 - x 8).
Residual Block: 256 channels.
Conv2D: stride 2, latent size channels (input size 8 -
Residual Block: 2 x latent size channels.
Split into two branches for μ and logo?, o², each with latent size channels.
DECODER ARCHITECTRRE
The decoder reconstructs the input image x' E 13x61x6 from the latent representation. The architecture consists of the following layers:
Reshape: Latent vector reshaped to size (latent size, 4, 4).
Residual Block: latent size channels.
Con/Transpose2D1 3 x 3, stride 2, 256 channels (output size 4 x 4 + 88x 8).
Residual Block: 256 channels.
Con/Transpose22:: 3 x 3, stride 2, 128 channels (output size x 8 > 16 x 16).
Residual Block: 128 channels.
ConyTranspuse2DD: 3, stride 2, 64 channels (output size 16 x 16 - 32 32).
Residual Block: 64 channels.
TonvTranspose2D:3 3, stride 2, channels (output size 32 32 - 64 x 64).
Sigmoid Activation: To map outputs to the range [0, 1).
1.5. Domain Adaptation Classifier
Classifier consists of linear layers and ReLU activation function.
I.6. Pretrained CLIP
For this work, we utilized pretrained CLIP models from the OpenCLIP repository. Specifically:
VIT-II-144378-qqiickkelu on dfn5b dataset was employed for training the GW-SP regularizer.
ViT-L--44-quicgellu on dfn2b dataset was used for evaluation on the Image Retrieval task.
J. More Synthetic Dataset Results
20

---

Expressive Score-Based Priors for Distribution Matching with Geometry-Preserving Regularization
8
3
0
7
2
7
5
6
6
8
3
0
7
2
7
9
0
6
6
3
0
2
7
6
0
2
9
2
0
9
3
5
4
3
6
3
S
2
0
9
3
5
4
3
6
3
2
0
3
5
9
3
6
1
0
3
m
1
6
9
4
3
2
4
?
2
33
2
6
9
0
3
2
4
7
2
3
2
6
9
0
3
2
2
3
2
6
3
6
4
6
9
3
9
6
4
8
9
3
9
9
60
4
9
3
q
9
Figure Multi-domain adaptation: MNIST images rotated various angles.
21

---

Expressive Score-Based Priors for Distribution Matching with Geometry-Preservinng Regularization
Iranslated reconstructed
Space Representation
Gaussian
Dimondtin
Lanne Den'
Dimerakn recanstructor
Latani Dim
Translated
Latent Representation
MoG
Lokeni Din
Dimernion
Dimersion
Lanent Din
given
Space Representation
Vampprior
Dimomaon2
Ufmensionn Trarslated x1 given 22
Limunian reconstructed
Laidni Din
Latent Space Represeration
LSGM
Lutoon O
Translated
recanstructor
Latont Ropresentation
to
Ours
Dimersion
Cineroion
Latsri LSm
Figure 10. These figures show the translated dataset, reconstructed dataset, well the latent space under sample size 20.
22

---

Expressive Score-Based Priors for Distribution Matching with Geometry-Preservinn Regularization
tremilted xi given
x1 recorniruzted
Labent Spacee Represemation
Gaussian
Dinordann
Dlmorsien
Laiunn then
0.5
100
1,3
2.0
23
0.
20
20
2.
-2
-3
Dimeraion
Gimeraion
Trasiated given
X1 reconctructoc
Latent Space Ropresenration
MoG
Crmenssin
Latant O01
Dinerskan
Cimersion 1
Latann Dim
glvan
Space Represenntttion
DI
Vampprior
Dimensoo22
Latari 2012
Dinersion
Translated x. given
reconstructed
Latent Space Representation
LSGM
2inorsion2
Drnrrsinn
Laoent Omm
Dimernton
Dimamition reconstructed
Labent
Trandizted xL givan
Lazent Space Representation
D0
Ours
Diirension
Lulter Dm
2
Uiimrssin
Figure 11. This figures show the translated dataset, reconstructed dataset, well the latent space under sample size 50.
23

---

Expressive Score-Based Priors for Distribution Matching with Geometry-Preservinng Regularization
tramilated
recorstructec
Latert Space Representation
DLJ
Gaussian
Dimondton
Laiand
05
10
20
25
20
45
0.5
10
20
25
300
33
Demeraaion
Dimuraiio
Latant Dim Lete Representation
Transiated
reconstructed
MoG
Dimentionn
Latont 200
Dimensson on
Utmenssoo
Latera Din Space Representation
given
Vampprior
Dimerdon
Loconi D011
DiT
Dimension x1 reconstructed
Translated X1 given 22
Lstent Space Representation
LSGM
Dimersoon
Dinerston2
Dim Latont Represontation
Translated
rocanstructoo
Ours
Diiertsoo
Ciinrrsin
Latsii Lim
Figure 12. This figures show the translated dataset, reconstructed dataset, as well the latent space under sample size 100.
24

---

Expressive Score-Based Priors for Distribution Matching with Geometry-Preserving Regularization
given
Space Representation
D0.1
Gaussian
Dimatobraa
Dimeroban2
Latent Dim
Translated xc. given 22
X1 reconstructed
Lotent Space Regresentation
MoG
Dinossbrrs
Dimerodor
Vampprior
Cimundon
Dimmeittn
Ulierska
Latert Jim
given
LSGM
Dimundon
Dimmmaton
Lakent Jin
given
Space Representation
Ours
O
Dimerson
Dimercton
Dimension
Figure 13. This figures show the translated dataset, reconstructed dataset, as well the latent space under sample size 200.
25

---

Expressive Score-Based Priors for Distribution Matching with Geometry-Preserving Regularization
Translated x. given
reconstructed
Latent Space Representation
0.12
Gaussian
Dimunobn2
Dimendor
given
Representation
MoG
Latent 3in2
Dimercion given
Dinersion
Space Representation
Vampprior
Ormossion
Letent 2ud
given
Space Representation
LSGM
Dimamson
Lotten Chinr
Tronllaeed civan
reconstrucied
Latiett Space Rspreseniation
Latent Din
Ours
Figure 14, This figures show the translated dataset, reconstructed dataset, as well the latent space under sample size 500.
26

---

Expressive Score-Based Priors for Distribution Matching with Geometry-Preserving Regularization
K. Image translation between MNIST and USPS
MNIST to USPS
2
3
4
5
6
7
8
9
Source
0
1
2
3
4
5
6
7
8
9
Recon
3
4
5
6
7
8
9
Dommain tronslation
Figure 15. MNIST to USPS translated image trained with SP.
L. FairFace Image Translation
This experimental setting is conducted in fully unsupervised manner without SP loss. We compare proposed score-based prior (SAUB) with multi-Gaussian based learning prior (VAUB) to evaluate their effectiveness.
L.1. Handpicked samples
Domain
Domain
Domain
Domair
Reconobruction
VAUB
Ours
VAUB
Ours
Figure 15. (a) Male Female translation
Figure 15, Female Male translation
Figure 16. In this experiment, both models trained in unsupervised manner (i.e., SAUB trained without GW-SP loss). SAUB clearly exhibits superior semantic preservation in both (a) and (b), particularly with respect to features such as skin color, race, and age. Notably, SAUB makes minimal adjustments when altering gender, while VAUB struggles to retain the identity of the original data. (These samples are handpicked to illustrate the trend.)
27

---

Expressive Score-Based Priors for Distribution Matching with Geometry-Preserving Regularization
L.2. Random Samples
In Fig. 17, show completely random samples from the FairFace dataset.
Figure 17. Random samples from the FairFace experiment using our method. Top three rows translate from male to female and the bottom three translate from female to male. First row is original, second is reconstructed, and third is translated.
28

---

Expressive Score-Based Priors for Distribution Matching with Geometry-Preserving Regularization
M. Additional Random Image Translations on CelebA
Examples of random image translations between black hair and blonde hair are presented Fig. 18 and Fig. 19
Figure 18, Random Samples from Black to Blonde Hair Female
29

---

Expressive Score-Based Priors for Distribution Matching with Goometry-Preserviing Regularization
N. Ablation Studies and Computational Efficiency Analysis
Component-wise Ablation Studies To isolate the contributions our Monnov-Wasserstein (GW) structural preservation regularization component. we conducted comprehensive ablation studies comparing our SAUB method with established baselines. Following the same experimental setup described in subsection 5.2, we evaluated VAUB with GW regularization and LSGM (Vahdat al., 2021) both with and without GW-EP regularization the fairness dataset, which provides clear metrics for both distribution matching (DP gap) and downstream task performance (Accuracy).
The experimental results demonstrate two key findings: (1) GW regularization consistently improves performance across all methods in terms of downstream task performance, validating the effectiveness of our structural preservation approach; and (2) our SAUB method achieves comparable performance to LSGM with and without GW regularization, while offering computational advantages as detailed below.
Computational Efficiency Analysis We conducted systematic comparison of computational efficiency between our approach and LSGM across varying latent dimensions. Due to the varying dimensionalities across different tasks in experimental framework-from toy datasets (114 dimensions) to CelebA datasets (64 x 64 x 3 dimensions) performed controlled experiments with different latent dimensions (from 8 to 256) to simulate the parameter structures encountered across our diverse experimental tasks. Table presents the computational efficiency metrics comparing LSGM and our approach across different latent dimensions.
Table Computational efficiency: LSGM vs. Ours
VRAM (MB) |
Time/epoch (ms) |
Speedup
Dim
LSGM
Ours
LSGM
Ours
VRAM
Time
28.6
27.9
138.0
121.8
1.03x
1.13x
16
30.8
28.3
142.5
138.8
1.09x
1.03x
43.1
33.5
140.9
137.6
1.29x
1.02x
128
60.2
41.8
146.5
141.4
1.44x
1.04x
256
99.3
60.3
146.9
140.1
1.65x
1.05x
Our approach demonstrates consistently lower VRAM requirements compared to LSGM, with efficiency gains becoming more pronounced as latent dimensions increase. Training speed improvements range from 1.02x 1.13x across different dimensions. These efficiency gains stem primarily from avoiding the costly Jacobian computations required in LSGM, as detailed in theoretical analysis in subsection 3.3. The improved efficiency is particularly significant as the proportion of network parameters allocated for score-based prior distributions increases with dataset dimensionality, making our approach increasingly ad adntaggeouss for high-dimensional applications.
These results demonstrate that combining score-based priors with GW regularization not only maintains competitive performance but also provides computational benefits, making our approach more practical for world applications across diverse domains.
30

---

Expressive Score-Based Priors for Distribution Matching with Geometry-Preserringg Regularization
Figure 19, Random Samples from Blonde to Black Hair Female
31

---

Expressive Score-Based Priors for Distribution Matching with Geometry-Preserving Regularization
0.25
VAUB w/o GW-EP
VAUB w/ GW-EP (new)
LSGM w/o GW-EP (new)
0.20
LSGM w/ GW-EP (new)
SAUB w/o GW-EP
SAUB w/ GW-EP
0.15
Unfair Classifier
App (4)
0.10
0.05
0.00
0.74
0.76
0.78
0.80
0.82
0.84
0.86
Accuracy (t
Figure 20. Fairness-accuracy trade-off comparison across different methods with and without Gromov- Saatoo---assrttein regularization. The plot shows accuracy (x-axis) versus demographic parity gap App (y-axis) for VAUB, LSGM, and SAUB methods. Lower App values indicate better faimess performance, while higher accuracy values indicate better downstream task performance.
32


Reviews:
[{"review_id": "v97RZiLJP1", "reviewer": "Reviewer_5UvJ", "summary": "Existing DM methods face many challenges, and likelihood-based methods often impose unnecessary biases through fixed priors or require learning complex prior distributions. \nThis paper introduces a novel approach to distribution matching (DM) by leveraging score-based priors and Gromov-Wasserstein (GW) distance based structural regularization. This new approach eliminates biases from fixed priors and avoids the computational overhead of learning full prior densities through gradient of the log-probability density, preserves the geometric structure of data in the latent space by GW SP/EP structural regularization. Also experiments are conducted, which outperform baseline methods on the MNIST-USPS domain adaptation task.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your insightful review and valuable suggestions. We appreciate your careful assessment of our work, particularly regarding the Gromov-Wasserstein structural preservation regularization(GW) component.\nFollowing your recommendations, we conducted additional experiments to isolate the contributions of our GW components:\n\n**Additional Experimental Results**\nWe performed ablation studies where we **additionally add VAUB with GW, LSGM [1] with and without GW** into the existing experiment framework. Due to limited time for the rebuttal, we focused these experiments on the fairness dataset, which provides clear metrics for both distribution matching metric (DP gap) and downstream task performance (Accuracy): [*[Figure Anonymous Link]*](https://anonymous.4open.science/r/SAUB-DB57/) (note that we omit other baselines in the figure for clearer comparison)\nFrom the figure we observe that:\n    1.  GW regularization consistently improves performance across all methods in terms of downstream tasks performances.\n    2.  Our SAUB method achieves comparable performance compared to LSGM with and without GW regularization.\n\n**Runtime and Memory Efficiency**\nWe also compared computational efficiency metrics between our methods and LSGM:\n\n| LSGM/Ours               | dim=8           | dim=16          | dim=64          | dim=128         | dim=256         |\n|-------------------------|-----------------|-----------------|-----------------|-----------------|-----------------|\n| Allocated VRAM (MB)     | 28.6/**27.9**   | 30.8/**28.3**   | 43.1/**33.5**   | 60.2/**41.8**   | 99.3/**60.3**   |\n| Training per epoch (ms) | 138.0/**121.8** | 142.5/**138.8** | 140.9/**137.6** | 146.5/**141.4** | 146.9/**140.1** |\n\nDue to limited time, we could not measure extensive runtime analysis on all tasks. Instead, we found that as the dataset has more and more dimensions (e.g. from fairness dataset(114) to CelebA dataset(64 x 64 x 3)), the proportion of the network parameters allocated for score-based prior distributions is higher as well due to the necessity to have larger latent dimensions. Therefore, we varied the latent dimensions for the fairness tasks to simulate the parameter structure encountered in other tasks. \nAt dimension 128, the model represents a realistic scenario for applications in our domain adaptation experiments, and similarly at dimension 256 mimic the domain translation experiments.\n\nOur approach demonstrates lower VRAM requirements compared to LSGM and becomes more obvious when the latent dimension is larger. Training speed, on the other hand, improves about 1.1-1.2x across different dimensions. Such observations are consistent with our theoretical analysis in Section 3.2.1 and Appendix E. These efficiency gains stem primarily from avoiding the costly Jacobian computations required in LSGM, as detailed in our paper.\n\nWe believe these additional results strengthen our paper by clearly delineating the contributions of each component while confirming the complementary benefits of combining score-based priors with GW regularization.\nWould these clarifications and additional experimental results address your concerns sufficiently to warrant reconsidering your evaluation score?\n\n[1] Vahdat, A., et al. (2021). Score-based Generative Modeling in Latent Space.", "claims_and_evidence": "most evidences are convincing, but experiments against Latent Score-Based Generative Models (LSGM) are not conducted. \n This new approach is inspired by LSGM, so experiments against LSGM is important to reveal the improvement.", "methods_and_evaluation_criteria": "Mostly yes, but it seems GW SP/EP structural regularization is a common method, it is not derived from the score-based priors, so more evaluation are needed:\n- new approach without GW SP/EP vs LSGM\n- LSGM + GW SP/EP vs LSGM", "theoretical_claims": "yes, Proof of Proposition 3.1", "experimental_designs_or_analyses": "Yes，more experiments are needed, these are mentioned above.", "relation_to_broader_scientific_literature": "This paper proposes a new approach with score-based priors and Gromov-Wasserstein (GW) distance based structural regularization, it should be easy and common way for DM, also it can be applied in broad research.", "essential_references_not_discussed": "no", "other_strengths_and_weaknesses": "no, comments see above.", "other_comments_or_suggestions": "no, comments see above.", "questions_for_authors": "no", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "2CqJyRaKt7", "reviewer": "Reviewer_rQ6c", "summary": "This paper deals with the limitations of existing distribution matching (DM) methods, which often struggle with scalability, instability, mode collapse, or impose unnecessary biases through fixed priors.  To overcome these limitations, the authors builds upon the existing work VAUB, and propose a novel approach that models the prior density through its score function, using denoising score matching techniques.  The goal of the proposed approach is to avoid biases from fixed priors.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are grateful for the reviewer’s thoughtful feedback and recognition of our work. We apologize for any confusion caused by the inadequate explanation in **Section 3.1.1**, where the lack of explicit details led to misunderstandings regarding our training procedure. Below is our refined explanation:\n\n\nReviewer Concerns and Clarifications\n1. **The claim about flexible prior seems problematic...**\n    - The reviewer is correct that during the encoder and decoder updates, the VAUB loss is computed with a fixed prior model, meaning the parameter $\\psi$ is not updated at that stage. However, in the subsequent training step, we update the score model using the encoder’s posterior. After carefully re-examining the manuscript, we acknowledge that our manuscript did not clearly state that **our training algorithm alternates between updating the encoder/decoder and the diffusion (score) model**. This alternating approach—congruent to the strategy used in training Latent Score-Based Generative Models (LSGMs)[1]—ensures that while the prior remains fixed during certain updates, the subsequent score model update, driven by the evolving encoder posterior, provides the necessary flexibility and expressiveness. We plan to update our manuscript in future versions to explicitly clarify this point. We appreciate the reviewer’s attention to the details of our training algorithm. For further clarity, we would like to kindly refer the reviewer to the pseudocode provided in **Appendix B**, which offers a comprehensive explanation of our approach.\n\n\n2. **Usefulness in real world application is only weakly demonstrated..**\n    - We appreciate the reviewer's concern regarding our method's real-world applicability and experimental scope. **Our primary goal was to introduce a flexible distribution matching framework, demonstrating its promise in tasks like domain adaptation, domain translation, and fairness**. Our baseline experiments emphasize the method's versatility and effectiveness.\nWhile additional experiments with more recent baselines could enhance the evaluation, our objective was to establish a foundational framework with strong potential for future work. We believe that with further engineering, optimization, and advanced architectures (e.g., VQ-VAE [4] for the encoder/decoder and vision transformers [3] for diffusion models), our framework could achieve competitive performance. For this initial study, we used simple CNNs for the encoder/decoder and a shallow linear layer for the score model.\n\n3. **LSGM was demonstrated to be a strong image generative model...**\n    - This is an insightful question. Although our model was not initially designed to be a generative model, theoretically, **our model functions as a generative model if prior samples are generated via a reverse diffusion process and then decoded into synthetic images**. In principle, our approach should be capable of generating high-resolution images similar to those produced by LSGM. We believe with enough engineering with SOTA architecture, we could also achieve competitive image generation results.\nOne notable advantage of training SAUB as a generative model is the potential for reduced VRAM usage, as our method does not require backpropagation through the score model as LSGM does during the encoder/decoder update seen in the **Runtime and Memory Efficency table** of **Reviewer 5UvJ**. Additionally, our model may offer improved stability during the encoder/decoder training, as demonstrated in **Section 3.2.1**. This may lead to promising future works.\n\n4. **Does GW regularization based enhancement on existing method...**\n    - This arises an interesting ablation study. We further **add VAUB with GW, LSGM with and without GW** into existing experiment in the fairness experiment settings. As you can see in the figure [anonymous link](https://anonymous.4open.science/r/SAUB-DB57/), GW regularization consistently improves performance across all methods in terms of downstream tasks performances.\n\nWe sincerely appreciate the reviewer’s thoughtful and constructive questions, which have been invaluable in helping us refine and improve the clarity of our manuscript. If our responses have addressed all concerns and cleared any confusions, we kindly request consideration for an improved score. If further questions remain, we remain open to addressing them.\n\n\n[1] Vahdat, A., et al. (2021). Score-based Generative Modeling in Latent Space. \n\n[2] Tomczak, J. M., et al. VAE with a VampPrior. arXiv preprint, arXiv:1705.07120.\n\n[3] Gu, S., et al. (2022). Vector Quantized Diffusion Model for Text-to-Image Synthesis.\n\n[4] Razavi, A., et al. (2019). Generating Diverse High-Fidelity Images with VQ-VAE-2.", "claims_and_evidence": "The claim about expressive prior might be problematic. This seems to be related to the mismatch between $L_{\\rm VAUB}$ and $L_{\\rm SAUB}$. The prior parameters $\\psi$ is not optimized with respect to the actual $L_{\\rm VAUB}$ loss. Instead the prior $Q_{\\psi}$ seems to be fixed to some average distribution of $q_{\\theta}(z|x,d)$ (from Eq. (11)). If this is the case, it defeats the purpose of allowing flexible prior in $L_{\\rm VAUB}$.", "methods_and_evaluation_criteria": "Yes.", "theoretical_claims": "No, I did not check the proofs.", "experimental_designs_or_analyses": "Yes, the experiments on synthetic data, fairness representation learning, domain adaptation, and domain translation are all important applications of the proposed method.", "relation_to_broader_scientific_literature": "Distribution matching is an important problem in many unsupervised learning problems, e.g., generative modeling, domain adaptation, representation learning, domain translation, etc. The proposed method tries to advance techniques for non-adversarial distribution matching from a latent variable model perspective.", "essential_references_not_discussed": "No.", "other_strengths_and_weaknesses": "Strengths:\n1. Proposed method appears principled and technically sound.\n2. The proposed Score Function Substitution (SFS) trick seems novel and could be useful in other contexts as well.\n3. A variety of applications are considered to demostrate the efficacy of the proposed method\n\nWeaknesses and Questions:\n1. The claim about flexible prior seems problematic. Please refer to above explanation for details.\n2. Usefulness in real world application is only weakly demonstrated. For example, in the considered experiments on Domain adaptation and domain translation, recent baselines are not considered besides VAUB, and the qualitative results are not convincing.\n3. LSGM was demonstrated to be a strong image generative model. Can the proposed method be used for building such high-resolution image generative models, or are there inherent challenges ? \n4. Does GW regularization based enhancement on existing methods, such as VAUB or ADDA in domain adaptation, also improve their performance? It is of interest to know whether the GW regularization is only useful for the proposed scheme.", "other_comments_or_suggestions": "1. In page 4, Sec. 3.2, Appendix D should be Appendix E ? \n2. Some typos (spelling mistakes) in page 2 Sec. 2.", "questions_for_authors": "Please refer to Weaknesses and Questions.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "E1DCiMO6uq", "reviewer": "Reviewer_8aGm", "summary": "This paper proposes a new prior distribution for distribution matching. Specifically, the authors model the prior using denoising score matching, and they enhance this approach by incorporating the minimization of the Gromov-Wasserstein (GW) distance between different distributions as additional regularization. Experiments across various tasks, such as Fairness Representation Learning and Domain Adaptation, confirm the effectiveness of their methods.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "### Experimental Designs Or Analysis\n\nWe appreciate the reviewer’s suggestion and understand the importance of evaluating our method on more complex benchmarks. Our current domain adaptation experiments using MNIST and USPS were chosen deliberately as controlled, standard benchmarks that allow us to clearly illustrate the core capabilities and innovations of our framework. **Our primary goal in this work was to introduce a new, flexible framework for distribution matching that leverages a learned, expressive score-based prior within a variational probabilistic setting, and that incorporates innovations such as the Score Function Substitution (SFS) trick. And therefore, we not only applied to domain adaptation tasks but also other downstream tasks, such as fairness representation learning and domain translation.**\n\nNotably, our current model deploys basic CNN layers and a simple linear layer UNET. We believe that many additional engineering techniques, including added training methods for the score model (e.g., EMA) and the integration of state-of-the-art architectures such as VQ-VAE [5] for the encoder/decoder and Vision Transformers [4] for diffusion models, are needed to achieve competitive real-world performance. Our intent here was to demonstrate the barebone performance of our approach.\nWe acknowledge that incorporating more contemporary baselines (e.g., as mentioned in [1]) and extending evaluations to complex datasets like PACS [2] and Office-Home [3] would provide further insights into the robustness and scalability of our method. Moreover, it is promising that our method shows potential for saving memory and enhancing stability compared to LSGM seen in [table in reviewer 5UvJ], which is already regarded as a powerful model. With further enhancements, we are confident that our framework can be extended to yield competitive results in more demanding domain adaptation scenarios in future works.\n\nWe sincerely appreciate the reviewer’s thoughtful and constructive questions, which have been invaluable in helping us refine and improve the clarity of our manuscript. If our responses have addressed all concerns and cleared any confusions, we kindly request consideration for an improved score. If further questions remain, we remain open to addressing them.\n\n[1] Farahani, Abolfazl, et al. A brief review of domain adaptation.\n\n[2] Li, Da, et al. Deeper, broader and artier domain generalization.\n\n[3] Venkateswara, Hemanth, et al. Deep hashing network for unsupervised domain adaptation.\n\n[4] Gu, S., Chen, et al. Vector Quantized Diffusion Model for Text-to-Image Synthesis.\n\n[5] Razavi, A., et al. Generating Diverse High-Fidelity Images with VQ-VAE-2.", "claims_and_evidence": "All the claims are supported by clear and convincing evidence.", "methods_and_evaluation_criteria": "The methods and evaluation criteria are suitable for the problem.", "theoretical_claims": "Due to the limited time, I couldn't check every detail of the derivation. However, most of the theorems seem intuitive and appear correct.", "experimental_designs_or_analyses": "In the domain adaptation experiments, the authors rely solely on simplistic datasets (MNIST and USPS) and use limited baselines such as DANN. To enhance the robustness and relevance of their findings, it is advisable to incorporate more contemporary baselines as mentioned in [1], and to extend testing to more complex datasets like PACS [2] and Office-Home [3].\n\n[1] Farahani, Abolfazl, et al. \"A brief review of domain adaptation.\" Advances in data science and information engineering: proceedings from ICDATA 2020 and IKE 2020 (2021): 877-894.\n\n[2] Li, Da, et al. \"Deeper, broader and artier domain generalization.\" Proceedings of the IEEE international conference on computer vision. 2017.\n\n[3] Venkateswara, Hemanth, et al. \"Deep hashing network for unsupervised domain adaptation.\" Proceedings of the IEEE conference on computer vision and pattern recognition. 2017.", "relation_to_broader_scientific_literature": "The paper primarily draws inspiration from paper [1], which introduced a non-adversarial distribution matching technique using a Variational Autoencoder (VAE). Building on this foundational work, the current paper proposes new score-based priors.\n\n[1] Gong, Ziyu, et al. \"Towards Practical Non-Adversarial Distribution Matching.\" International Conference on Artificial Intelligence and Statistics. PMLR, 2024.", "essential_references_not_discussed": "As far as I known, all the essential references are discussed in the paper.", "other_strengths_and_weaknesses": "**Strengths:**\n\n1. The introduction of a score-based prior combined with GW distance regularization represents a novel approach.\n2. The method's effectiveness is demonstrated through experiments across different tasks, including Fairness Representation Learning and Domain Adaptation.\n3. The paper is well-organized and clear, making it easy to understand.\n\n**Other Weaknesses:** \n\nThere are no significant weaknesses outside of those mentioned in the \"Experimental Designs\" section.", "other_comments_or_suggestions": "1. On line 132 of the left column, it seems the authors missed including a citation for the Gromov-Wasserstein (GW) distance.\n2. On line 272 of the left column, it seems a pair of brackets is missing for “Section 2”.", "questions_for_authors": "None", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "LOuelMxk3k", "reviewer": "Reviewer_RKSQ", "summary": "The paper introduces a distribution matching method using VAE by leveraging expressive score-based priors instead of the fixed priors such as Gaussians. The main contribution is the Score Function Substitution (SFS) trick, which reformulates the gradient of the prior’s cross-entropy term to avoid the computationally expensive Jacobian calculation required in latent score-based generative models. By focusing on learning the score function via denoising score matching, the method circumvents the need for explicit density estimation and enables more stable training. Moreover, the paper incorporates structural regularization through a Gromov-Wasserstein-inspired loss, which preserves geometric relationships in the latent space to ensure that latent representations retain task-relevant structure. Experimental results across synthetic datasets, fairness representation learning (using the Adult dataset), domain adaptation (MNIST–USPS), and domain translation (CelebA and FairFace) demonstrate that the proposed approach improves latent space quality and downstream task performance.\n\n########################## Post Rebuttal ##########################\nI would like to thank the authors for the detailed response. Most of my concerns are properly addressed. I would like to raise my score.\n########################## Post Rebuttal ##########################", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate the reviewer’s thoughtful feedback and are pleased that our work was found both interesting and novel. However, we believe some aspects were probably **misunderstood or overlooked**. For example, the reviewer noted missing references, yet our manuscript thoroughly discusses them.\n\nGromov-Wasserstein Autoencoder (GWAE) [1] is a cornerstone of our approach. We state in our contribution (line 59, right column) that we \"adopt the Gromov-Wasserstein-based constraint from Gromov-Wasserstein Autoencoders,\" and this is reiterated in Section 3.3 (line 220, right column) and the related works section (line 310, left column), with additional references on lines 65 and 250.\nSimilarly, VampPrior [4] is used as a baseline model in Section 5.1 (line 291, right column) and is cited in the related works (line 287, left column).\nFurthermore, Latent Score-Based Generative Models (LSGM) are integral to our narrative. We compare LSGM with our SFS trick in Section 3.2 and provide a detailed analysis of the gradients in our objective functions.\n\n## Experimental Designs Or Analyses: \n\nOur synthetic nested D-shaped dataset was chosen as a preliminary experiment to isolate and clearly demonstrate how different priors affect latent space separation. This controlled setting provides an intuitive visualization, unlike high-dimensional latent spaces that require non-linear reduction methods. (Please refer to reviewer 5UvJ for section **Additional Experimental Results** for extra experiments comparing between LSGM and ours on Fairness task.)\nFinally, we remind the reviewer that our domain translation evaluation includes a CLIP-based metric for semantic preservation, along with LPIPS and SSIM scores, which were selected because they directly assess semantic preservation compared to metrics such as FID or PSNR. We welcome any suggestions for additional or alternative metrics.\n\n## Other Strengths and Weaknesses\n\nOur evaluation across fairness, domain adaptation, and domain translation is intended as a unified process to assess the efficacy of our distribution matching framework rather than a shifting narrative. All downstream tasks were evaluated both with and without the GW constraint. We present additional experiments in **Appendix D** and **Appendix F** that isolates the affects of GW-EP and GW-SP distances.\nFurthermore, as detailed in our response to reviewer 5UvJ in the **Runtime and Memory Efficiency** section, our SFS trick offers benefits in terms of reduced VRAM usage compared to LSGM.\nFinally, we respectfully note that several references identified as missing were, in fact, included in our manuscript. We welcome any specific suggestions for additional references.\n\n## Questions For Authors:\n\n1. We use GW-EP uniformly across all experiments and only GW-SP in domain adaptation and domain translation experiments. The reason we did not use GW-SP uniformly across all experiments is that, quoting from line 335, we do not have a semantic model to compute the semantic distance within tabular dataset such as Adult dataset that is used in fairness tasks. We do however have such semantic models for image dataset such as CLIP model and thus GW-SP is possible on these image dataset.\n\n3. Latent Score-Based Generative Models learn an expressive score-based prior by approximating the cross-entropy term with a weighted denoising score matching objective. Our SAUB approach also employs a variational probabilistic framework with a learned prior score model to capture expressive latent distributions. For further details, please see Section 3.2 and Appendix E.\nThe Score Distillation Sampling (SDS) loss bypasses the diffusion UNET's Jacobian computation. Although SDS and our Score Function Substitution (SFS) trick both involve a KL divergence term (i.g., the KL term in ELBO for SAUB), SDS replaces this term with a denoising score matching objective and removes the score gradient using the \"Sticking-the-Landing\" technique [2]. In contrast, our SFS trick computes the exact KL gradient without omitting any components or resorting to denoising score matching [5].\n\n4. Please refer to **Experimental Designs Or Analysis** response for reviewer 8aGm.\n\nWe sincerely appreciate the reviewer's thoughtful feedback and hope that our responses have fully addressed all concerns. If so, we kindly ask that you consider raising the score accordingly. We remain available to provide further clarifications if needed.\n\n[1] Zhang, Z., et al. Gromov-Wasserstein Distances: Entropic Regularization, Duality, and Sample Complexity.\n\n[2] Roeder, G., et al. Sticking the Landing: Simple, Lower-Variance Gradient Estimators for Variational Inference.\n\n[3] Vahdat, A., et al. Score-based Generative Modeling in Latent Space. \n\n[4] Jakub Tomczak, et al. Vae with a vampprior.\n\n[5] Poole, B., et al. DreamFusion: Text-to-3D using 2D Diffusion.", "claims_and_evidence": "Some claims are supported by experiments, while others lack sufficiently detailed evidence. For example:\n\n1. The authors include comparisons (e.g., NLL curves) that indicate improved stability over LSGM. However, a more in-depth analysis (perhaps with additional ablations or broader noise level tests) would strengthen the claim.\n\n2. The authors claim that their method avoids expensive Jacobian computations and is computationally efficient compared to other methods. However, there is no quantitative runtime analysis (e.g., training time per epoch or total convergence time) to support this claim. This makes it difficult to evaluate the true computational benefits relative to simpler methods like using Gaussian priors.\n\n3. The idea of incorporating a Gromov-Wasserstein-inspired loss to preserve latent space geometry is supported by qualitative visualizations and some quantitative metrics. However, the evidence could be more convincing with clearer ablation studies that isolate the impact of this regularization on various downstream tasks.", "methods_and_evaluation_criteria": "The proposed methods and evaluation criteria make sense for a proof-of-concept for distribution matching using expressive score-based priors. The use of benchmarks like MNIST–USPS for domain adaptation, the Adult dataset for fairness, and CelebA/FairFace for image translation demonstrates the method’s practical benefits. However, while these datasets are well-known and widely used in the literature, they are relatively small and may not capture the complexity or scale of real-world applications. Evaluating on larger, more challenging datasets (e.g., VisDA-2017 for domain adaptation) could give more insights into the scalability, robustness, and generalization of the proposed methods across more diverse conditions.", "theoretical_claims": "I have checked the proof provided for Proposition 3.1 in Appendix A. This proof derives the gradient of the cross-entropy term using techniques like the reparameterization trick and the chain rule. The steps are logically sound and follow common practices in variational inference.\n\nOne minor concern is that while the notation regarding the detachment of  z  from the gradient (i.e., treating it as a constant with respect to the encoder parameters) is correct, it could be made more clear for readers who are less familiar with such subtleties.", "experimental_designs_or_analyses": "I have checked the experimental designs and analyses in the paper, including:\n\n1. A synthetic nested D-shaped dataset to evaluate how different priors affect latent space separation. They analyze performance using AUROC scores derived from classifier performance. While this design is useful for visualizing and quantifying latent space quality, the synthetic nature of the data limits its complexity relative to real-world scenarios.\n\n2. Experiments on the Adult dataset are designed to test the trade-off between demographic parity and accuracy. The experimental setup is standard in fairness evaluation; however, the scale of the dataset and the tasks might not capture challenges encountered in more complex settings.\n\n3. The MNIST-USPS adaptation scenario is tested. Although these benchmarks are common in the literature, they represent relatively simple and low-dimensional domain shifts. This raises concerns about whether the method would generalize to more challenging, large-scale datasets.\n\n4. The qualitative evaluation on domain translation (e.g., CelebA and FairFace) provides visual evidence of the method’s ability to preserve semantic content. However, one issue is the absence of comprehensive quantitative metrics.\n\n#### Concerns:\n\n1. While the datasets used are standard, they are relatively small-scale and may not reflect the challenges of large-scale, high-dimensional data (e.g., VisDA-2017 for domain adaptation). \n\n2. There is limited discussion on isolating the impact of individual components (e.g., the structure-preserving GW loss, the SFS trick).\n\t\n3. Although the paper claims improved computational efficiency (particularly avoiding expensive Jacobian computations), there is no quantitative runtime analysis (e.g., training time per epoch, total convergence time). Such metrics are important to validate the claimed efficiency benefits.", "relation_to_broader_scientific_literature": "The paper’s main contributions build upon several topics, which include generative modeling, variational inference, and optimal transport. To be specific, they include: score-based priors and diffusion models; VAE with expressive priors; structural regularization via optimal transport; and bridging non-adversarial distribution matching and diffusion models.", "essential_references_not_discussed": "Yes, some references highly related are not included. For example:\n\n1. The work is built on score-based generative modeling, it would benefit from a more detailed discussion of prior advances in variational inference that improve gradient estimation. For example, [1] proposed lower-variance gradient estimators for variational inference, which are conceptually related to the proposed SFS trick. \n\n2. The proposed method is also closely related to latent score-based generative models like [2], which forms the basis for the comparison and motivation. However, the submission did not thoroughly cite or contrast these methods, which are important for understanding the novelty and benefits of the SFS trick.\n\n3. The importance of using learnable priors in VAE is well established in prior literature, particularly in works like VampPrior [3].\n\n4. Recent works such as [4] on Gromov-Wasserstein Autoencoders, which demonstrate how optimal transport can enforce structural constraints in latent spaces, which is directly relevant to the use of the GW-based Semantic Preserving loss.\n\n[1] Roeder, G., Wu, Y. and Duvenaud, D.K., Sticking the landing: Simple, lower-variance gradient estimators for variational inference. NeurIPS 2017.\n\n[2] Rombach, R., Blattmann, A., Lorenz, D., Esser, P. and Ommer, B. High-resolution image synthesis with latent diffusion models. CVPR 2022.\n\n[3] Tomczak, J. and Welling, M. VAE with a VampPrior. In International conference on artificial intelligence and statistics, 2018.\n\n[4] Nakagawa, N., Togo, R., Ogawa, T., and Haseyama, M. Gromov-Wasserstein Autoencoders. ICLR 2023.", "other_strengths_and_weaknesses": "#### Strengths:\n\n1. The authors propose a novel combination of the score-based generative modeling with structural regularization via Gromov-Wasserstein distances. The introduction of the SFS trick provides a way to sidestep costly Jacobian computations, and the idea of enforcing latent space geometry via GW-based loss is interesting.\n\n2. The ability to learn flexible priors without explicit density estimation is a promising direction that might alleviate common issues in VAEs with fixed priors.\n\n3. Some sections of the paper, particularly the detailed derivations and the pseudo-code in the supplementary material, are clearly presented. These sections provide sufficient details for understanding the theoretical contributions.\n\n4. The authors validate the proposed method on a wide range of tasks.\n\n#### Weaknesses:\n\n1. The overall narrative can be somewhat inconsistent, as the paper shifts between topics like distribution matching, fairness, and domain adaptation. This makes it challenging to pinpoint the central contribution. Additionally, key notations and the overall problem formulation are not introduced early enough, which affects readability.\n\n2. While the experiments are conducted on a wide range of tasks, the evaluation is limited to relatively small and simple datasets. This raises questions about the method’s scalability and applicability to real-world scenarios. Furthermore, the impact of the structural regularization (GW-SP loss) is not thoroughly isolated in ablation studies.\n\n3. My main concern is the absence of quantitative runtime analysis. While the paper claims that the SFS trick leads to improved computational efficiency by avoiding costly Jacobian computations, it does not provide measurements such as training time per epoch, total convergence time, or inference speed compared to standard VAEs with Gaussian priors. Such an analysis would be important to evaluate the practical values of the proposed method in real-world applications.\n\n4. The paper does not sufficiently cite or discuss related works that use optimal transport for domain adaptation and structural regularization, or other methods for learnable priors.", "other_comments_or_suggestions": "1. Make sure that abbreviations (e.g., GW, SFS, VAUB, etc.) are defined once and used consistently throughout the paper.", "questions_for_authors": "1. Could the authors clarify whether the Gromov-Wasserstein Semantic Preserving (GW-SP) loss is applied uniformly across all experiments, especially in Sections 5.1 and 5.2? If not, can the authors provide an ablation that shows the impact of including versus omitting this term?\n\n2. Could the authors provide detailed runtime comparisons (e.g., training time per epoch, total convergence time, and inference time) between the proposed methods and standard VAEs with Gaussian priors, as well as methods like LSGM?\n\n3. Could the authors elaborate on how the SFS trick relates to and differs from existing methods such as score distillation sampling and latent score-based generative models like LSGM?\n\n4. Could the authors evaluate the proposed method on large-scale datasets, such as VisDA-2017 for domain adaptation?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The paper proposes a VAE-based framework that incorporates expressive score-based priors and structural regularization via Gromov-Wasserstein loss. Reviewers appreciated the novelty of combining denoising score matching with structure-preserving objectives. During the discussion period, the authors addressed the initial concerns, particularly those related to the alternating training scheme and the connections to VAUB and LSGM. After the discussion, the reviewers reached a consensus. I suggest that the authors consider incorporating the additional experimental components discussed into the future version to make the evaluation more complete. Overall, I recommend acceptance.
Decision: Accept (poster)