Conditional Wasserstein Distances with Applications in Bayesian OT Flow Matching
Authors: Jannis Chemseddine, Paul Hagemann, Gabriele Steidl, Christian Wald
JMLR 2025 | Venue PDF | Archive PDF | Plain Text | LLM Run Details
| Reproducibility Variable | Result | LLM Response |
|---|---|---|
| Research Type | Experimental | In this paper, we introduce a conditional Wasserstein distance via a set of restricted couplings that equals the expected Wasserstein distance of the posteriors. [...] Based on this, we propose an extension of OT Flow Matching for solving Bayesian inverse problems and demonstrate its numerical advantages on an inverse problem and class-conditional image generation. [...] Finally, in Section 8, we present numerical results: we verify a convergence result for an approximation of the conditional Wasserstein distance using particle flows to MNIST, and demonstrate the advantages of our Bayesian OT flow matching procedure on a Gaussian mixture model toy example and on CIFAR10 class-conditional image generation. |
| Researcher Affiliation | Academia | Jannis Chemseddine EMAIL Paul Hagemann EMAIL Gabriele Steidl EMAIL Christian Wald EMAIL Institute of Mathematics Technical University of Berlin 10623 Berlin, Germany |
| Pseudocode | No | The paper primarily presents mathematical derivations, theoretical properties, and describes algorithms in paragraph form. It does not include any clearly labeled 'Pseudocode' or 'Algorithm' blocks or figures with structured, code-like steps. |
| Open Source Code | Yes | The code is written in Py Torch Paszke et al. (2019) and is available online1. 1. https://github.com/JChemseddine/Conditional_Wasserstein_Distances |
| Open Datasets | Yes | We use our proposed Bayesian OT flow matching and verify its advantages on a Gaussian mixture toy problem and on class conditional image generation on the CIFAR10 dataset. [...] First, we verify that the convergence result for an increasing parameter β given in Proposition 10 for particle flows to MNIST (Deng, 2012). [...] Class Conditional Image Generation. [...] on CIFAR10 (Krizhevsky et al., 2009) class-conditional image generation. |
| Dataset Splits | Yes | We train a random Bayesian flow matching model according to LY,FM(12) and our OT Bayesian flow matching according to LY,OT(13) with the python package POT (Flamary et al., 2021) on a fixed dataset of size 10000, where we choose the best model according to a validation set of size 2000 until the validation loss converges. [...] We compute the distance on 50k training samples, for which we generate 50k samples given the same labels as the training samples. |
| Hardware Specification | No | The paper discusses model training and inference but does not specify any particular hardware components like GPU models, CPU types, or memory specifications. It mentions general terms like 'Euler steps' for sampling but not the underlying hardware. |
| Software Dependencies | No | We train a random Bayesian flow matching model [...] with the python package POT (Flamary et al., 2021) [...] with the package Geom Loss (Feydy et al., 2019) [...]. The code is written in Py Torch Paszke et al. (2019) and for inference simulate the corresponding ODEs using the torchiffeq (Chen, 2018) package. The paper mentions several software packages but does not provide specific version numbers for them. |
| Experiment Setup | Yes | We train a random Bayesian flow matching model [...] on a fixed dataset of size 10000, where we choose the best model according to a validation set of size 2000 until the validation loss converges. [...] The sampling is done via an explicit Euler discretization of 10 steps. [...] using the time dependent U-Net architecture from (Nichol and Dhariwal, 2021; Dhariwal and Nichol, 2021) which are trained using Adam (Kingma and Ba, 2015). As in (Tong et al., 2023) we clip the gradient norm to 1 and use exponential moving averaging with a decay of 0.9999. The differences are we use a constant learning rate of 2e-4, 256 model channels and no dropout. We train using 50k target samples for 500 epochs using a batch size of 500 for the minibatch OT couplings and a batch size of 100 for training the networks. [...] For inference we simulate the flow ODE Chen (2018) using an adaptive step size solver (Runge-Kutta of order 5). |