Robust Barycenter Estimation using Semi-Unbalanced Neural Optimal Transport
Authors: Milena Gazdieva, Jaemoo Choi, Alexander Kolesov, Jaewoong Choi, Petr Mokrov, Aleksandr Korotin
ICLR 2025 | Venue PDF | Archive PDF | Plain Text | LLM Run Details
| Reproducibility Variable | Result | LLM Response |
|---|---|---|
| Research Type | Experimental | We rigorously establish the theoretical underpinnings of the proposed method and demonstrate its robustness to outliers and class imbalance through a number of illustrative experiments. Our source code is publicly available at https://github.com/milenagazdieva/U-NOTBarycenters. ... We conduct a number of experiments on toy and image data setups to demonstrate the performance of our method and showcase its robustness to outliers and imbalance of classes (5). |
| Researcher Affiliation | Academia | Milena Gazdieva Skolkovo Institute of Science and Technology Artificial Intelligence Research Institute Moscow, Russia EMAIL; Jaemoo Choi Georgia Institute of Technology Atlanta, GA, USA EMAIL; Alexander Kolesov Skolkovo Institute of Science and Technology Artificial Intelligence Research Institute Moscow, Russia EMAIL; Jaewoong Choi Sungkyunkwan University Seoul, Korea EMAIL; Petr Mokrov Skolkovo Institute of Science and Technology Moscow, Russia EMAIL; Alexander Korotin Skolkovo Institute of Science and Technology Artificial Intelligence Research Institute Moscow, Russia EMAIL |
| Pseudocode | Yes | Algorithm 1: SUOT Barycenter with Semi-Unbalanced Neural Optimal Transport Input: Distributions P1:K, S accessible by samples; transport costs ck : Xk Y 7 R; maps Tk,ω and m-congruent potentials fk,θ, k K; number NT of inner iterations; batch sizes. Output: Trained (stochastic) maps T[1:K],ω approximating SUOT plans between Pk and barycenter Q . repeat Sample batches Xk Pk, k K; For each xk Xk sample an auxiliary batch S[xk] S; f(Tk(xk,s)) c(xk,Tk(xk,s)) Update θ by using c Lf θ to maximize c Lf; for n T = 1, 2, . . . , NT do Sample batches Xk Pk; For each xk Xk: sample an auxiliary batch S[xk] S; d LTk 1 |Xk| P c(xk,Tk(xk,s)) f(Tk(xk,s)) Update ω by using c LT ω to minimize c LT ; until not converged; |
| Open Source Code | Yes | Our source code is publicly available at https://github.com/milenagazdieva/U-NOTBarycenters. ... The code for our model is written using Py Torch framework and is publicly available at https://github.com/milenagazdieva/U-NOTBarycenters. |
| Open Datasets | Yes | The first marginal distribution consists of grayscaled images of digits 2 (49% of training dataset), 3 (50%) and 7 (1% outliers) of MNIST data. ... P1 and P2 as collections of images of young individuals (age between 5 to 20) and elderly individuals (age of over 50), respectively, from the FFHQ (Karras et al., 2019) dataset. |
| Dataset Splits | Yes | The marginal distributions P1, P2 consist of FFHQ images of young individuals (5 < age < 20) and elderly individuals (age > 50). We randomly partition each distribution into 90% for the training set and 10% for the test set. |
| Hardware Specification | Yes | Table 2: Training time for various experiments. ... GPU RTX 2080Ti RTX 2080Ti RTX 2080Ti RTX 3090Ti ... Table 3: Inference time for various experiments. ... GPU RTX 2080Ti RTX 2080Ti RTX 2080Ti RTX 3090Ti ... Table 8: Hyper-parameter settings and training time for FFHQ experiment. ... GPU RTX 3090Ti |
| Software Dependencies | No | The code for our model is written using Py Torch framework and is publicly available at https://github.com/milenagazdieva/U-NOTBarycenters. ... Specifically, we sample 2000 points and solve the Unbalanced Optimal Transport (UOT) problem from P1 to P2 using the Python Optimal Transport (POT) library (Flamary et al., 2021). ... To classify the gender of images, we use the open python library called Open CV. --- The paper mentions PyTorch, POT library, and Open CV but does not specify their version numbers, which are required for reproducible software dependencies. |
| Experiment Setup | Yes | Implementation Details. For the experiments in Section 5.1 and the class imbalanced experiments, we use the number of iterations of 10K. For the outlier experiments in Section 5.2, we employ the number of iterations of 20K. We use the batch size of 1024. We update m with the Adam Optimizer with learning rate of 10-3 and (β1, β2) = (0, 0.9). For other hyperparameters, we follow the experimental settings of (Kolesov et al., 2024a). ... Table 4: Hyper-parameter settings of Algorithm 1 for various experiments. ... Table 8: Hyper-parameter settings and training time for FFHQ experiment. |