Research map / Learning theory and algorithms / Optimal transport
Expectile Regularization for Fast and Accurate Training of Neural Optimal Transport
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes ENOT, a neural optimal transport solver that replaces the inner c-transform optimization with an expectile regularizer on the dual Kantorovich potentials, achieving better W2 benchmark accuracy and much faster training.
Problem
Neural optimal transport solvers need near-exact c-conjugate (c-transform) estimates during dual training, which requires non-convex max-min optimization or costly inner fine-tuning, leading to instability and slow training.
Approach
Approximate the c-conjugate transform by a tau-expectile (tau close to 1) of g^T(x) - c(x,y) conditioned on y, yielding an asymmetric least-squares regularizer on g that pulls the dual potentials toward c-concavity without any inner optimization loop.
Key results
- On CelebA64 W2 benchmark, ENOT achieves L2-UV 0.32 / 0.08 / 0.04 (early/mid/late), versus best baseline 0.61 / 0.20 / 0.09 for W2OT.
- On HD Gaussian mixtures D=2 to 256, ENOT is best at most dimensions (e.g., D=4: 0.03 vs 0.22 for W2OT-Objective) and runs in about 14-21 minutes vs 157-1028 minutes for W2OT.
- On unpaired image translation, ENOT attains lowest FID in all four tasks (e.g., CelebA(f) to Anime 64: 13.12 vs 14.65 for Extremal OT), training in about 5 hours on one RTX 3090 Ti for 64x64.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| use expectile regularization (ENOT) instead of W2OT-Objective with L-BFGS conjugate optimization | Wasserstein-2 HD benchmark · D=64 | L2 UV: 2.08 ± 0.40 -> 0.56 ± 0.03 | Helped |
| use expectile regularization (ENOT) instead of W2OT-Objective with L-BFGS conjugate optimization | Wasserstein-2 HD benchmark · D=256 | L2 UV: 0.59 ± 0.04 -> 0.51 ± 0.02 | Helped |
| use expectile regularization (ENOT) instead of W2OT-Objective with L-BFGS conjugate optimization | Wasserstein-2 benchmark · CelebA64 Late Generator | L2 UV: 0.09 ± 0.00 -> 0.04 ± 0.002 | Helped |
| use ENOT instead of Extremal OT | image-to-image translation · FFHQ → Comics 128 | MSE: 20.95 -> 17.11 | Helped |
| use ENOT instead of StarGAN | image-to-image translation · CelebA(f) → Anime 64 | FID: 0.21 -> 0.26 | Hurt |
| increase expectile τ from 0.5 to 0.9 (ablation) | ENOT · Synthetic 2D | W2: 33.47 -> 0.06 | Helped |
| increase expectile τ from 0.5 to 0.9 (ablation) | ENOT · CelebA(f) → Anime | FID: 16.43 -> 13.87 | Helped |
| increase expectile τ from 0.5 to 0.9 (ablation) | ENOT · Wasserstein-2 benchmark D=256 | L2 UV: 0.55 -> 0.5 | Helped |
| set expectile τ=0.999 instead of τ=0.9 (ablation) | ENOT · CelebA(f) → Anime | MSE: 0.266 -> 0.288 | Hurt |
| remove expectile regularization by setting λ=0 (ablation) | ENOT · Wasserstein-2 benchmark D≥64 | convergence | Hurt |
Surprising findings
- Removing the expectile regularization (lambda=0) causes divergence in high dimensions (D>=64), showing the regularizer is needed for stability.
- Setting tau=1 may cause instability, as the regularization term can vanish and potentials become unbounded.
Limitations
- Introduces two hyperparameters (expectile tau and loss weight lambda) that may need re-tuning for new data.
- Bidirectional training is only applicable to strictly convex costs of form h(x-y); tested mainly on squared Euclidean cost.
- Instability at tau=1 in some cases.
Experimental setup
Datasets: Wasserstein-2 benchmark (Gaussian mixtures HD tasks) CelebA64 W2 benchmark synthetic 2D datasets CelebA(f) Anime Faces FFHQ Comic faces v2 Handbags Shoes
Benchmarks: Wasserstein-2 benchmark (CelebA64 and HD Gaussian mixtures) synthetic 2D datasets Handbags to Shoes FFHQ to Comics CelebA(f) to Anime
Metrics: unexplained variance percentage (L2-UV) FID MSE runtime
Baselines: W2-Cycle MM-Objective MM-R-Objective W2OT-Cycle W2OT-Objective W2OT-Regression Monge Gap CycleGAN StarGAN-v2 Extremal OT Kernel OT
Builds on: Wasserstein-2 benchmark W2OT Monge gap expectile regression OTT-JAX
Closest papers by meaning
- Variational Entropic Optimal Transport (ICML 2026 · similarity 0.90)
- Energy-guided Entropic Neural Optimal Transport (ICLR 2024 · similarity 0.89)
- Overcoming Spurious Solutions in Semi-Dual Neural Optimal Transport: A Smoothing Approach for Learning the Optimal Transport Plan (ICML 2025 · similarity 0.88)
- Progressive Entropic Optimal Transport Solvers (NeurIPS 2024 · similarity 0.88)
- Decreasing Entropic Regularization Averaged Gradient for Semi-Discrete Optimal Transport (NeurIPS 2025 · similarity 0.88)
- Universal Neural Optimal Transport (ICML 2025 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.