Fixed-Point Neural Optimal Transport without Implicit Differentiation

Technical Summary: Fixed-Point Neural Optimal Transport without Implicit Differentiation

Problem Statement
Optimal Transport (OT) seeks a mapping between probability distributions that minimizes a prescribed transportation cost. While recent deep learning approaches have enabled scalable OT solvers, they frequently rely on adversarial min–max optimization (e.g., WGANs) or multi-network architectures involving transport maps and dual potentials. These methods often suffer from training instability, sensitivity to hyperparameters, and convergence difficulties, particularly in high-dimensional settings. Alternative approaches based on continuous flows require solving ordinary differential equations (ODEs), incurring significant computational overhead, while Hamilton–Jacobi–Bellman (HJB) formulations struggle with scalability in high dimensions. Furthermore, many existing methods require auxiliary networks or alternating optimization procedures to enforce dual feasibility.

Methodology
The authors propose a single-network framework for quadratic optimal transport that eliminates adversarial training and implicit differentiation. The core of the method is an implicit neural formulation derived from the Kantorovich dual problem.

1. Single Potential Parameterization: The dual potential $\phi$ is parameterized as the negative of a convex function $g_\theta$ represented by a single neural network ($\phi_\theta(x) = -g_\theta(x)$). This leverages Brenier's theorem, which states that the optimal transport map is the gradient of a convex potential.
2. Proximal Fixed-Point Reformulation: Under this parameterization, the $c$-transform (a key component of the dual objective) reduces to a strongly convex proximal optimization problem:
   $$(-g_\theta)^c(z) = \inf_{y \in \mathbb{R}^d} \left\{ \frac{1}{2}\|y - z\|^2 + g_\theta(y) \right\}$$
   The minimizer $y^*_\theta(z)$ of this problem defines the backward transport map ($T_{\nu \to \mu}(z) = y^*_\theta(z)$), while the forward map is given by $T_{\mu \to \nu}(x) = x + \nabla g_\theta(x)$.
3. Training via Fixed-Point Iteration: The training objective involves minimizing the negative dual loss. Crucially, the inner minimization (the proximal operator) is solved using standard fixed-point iterations (e.g., gradient descent steps) rather than unrolled optimization or implicit differentiation.
4. Gradient Computation without Implicit Differentiation: A key theoretical insight is that the gradient of the loss with respect to the network parameters $\theta$ does not require backpropagation through the fixed-point iterations. Due to the first-order optimality conditions of the proximal operator, the terms involving the derivative of the fixed point with respect to $\theta$ ($dy^*/d\theta$) cancel out exactly. Consequently, the gradient depends only on the explicit derivatives of $g_\theta$ evaluated at the input samples and the approximate fixed-point solutions. This allows for efficient training using standard backpropagation.
5. Convergence: The authors establish that Stochastic Gradient Descent (SGD) converges to a stationary point even when the inner fixed-point problem is solved approximately (within a specified tolerance), provided the approximation error is bounded.
6. Class-Conditional Extension: The framework naturally extends to class-conditional OT by conditioning the potential $g_\theta(x, k)$ on class labels $k$, allowing a single network to handle multiple transport tasks simultaneously.

Key Contributions

• Single-Network Reformulation: The authors derive a formulation of quadratic OT that requires only a single neural network to recover both forward and backward transport maps, removing the need for adversarial min–max setups or auxiliary conjugate networks.
• Implicit Differentiation-Free Training: They demonstrate that the objective can be differentiated without implicit differentiation, relying on the exact cancellation of terms induced by the proximal operator's optimality conditions.
• Convergence Guarantees: The paper provides theoretical proofs for the convergence of SGD to a stationary point under inexact fixed-point evaluations, addressing the practical reality that the inner loop is not solved to machine precision.
• Unified Framework: The method simultaneously recovers forward and backward maps and extends naturally to class-conditional settings without architectural changes.

Experimental Results
The proposed method was evaluated on high-dimensional Gaussian benchmarks, real-world physics datasets (UCI), and image translation tasks (MNIST to Fashion-MNIST).

• High-Dimensional Gaussians: On Gaussian benchmarks with dimensions up to $d=64$, the method achieved significantly lower Unexplained Variance Percentage (UVP) compared to baselines like NOT, MM-v1, MM:R, NCF, and WGAN-QC. It maintained stable accuracy as dimensionality increased, whereas baselines showed degradation.
• Computational Efficiency: The method demonstrated favorable computational and memory efficiency, requiring fewer parameters and less training time than continuous flow-based methods (FFJORD, RNODE) and other neural OT approaches.
• Real-World Data: On diverse UCI datasets (Power, Gas, HEPMASS, MINIBOONE) and BSDS300 image patches, the method achieved competitive or superior distribution matching performance (measured by Maximum Mean Discrepancy) with substantially smaller model sizes.
• Image Translation: In class-conditional image translation tasks, the method achieved the lowest Fréchet Inception Distance (FID) and highest classification accuracy among baselines, indicating superior preservation of class structure and semantic content.
• Ablation Studies: Experiments showed that the method is robust to the stopping tolerance of the fixed-point solver; moderate inexactness in the inner loop yields a favorable trade-off between accuracy and speed. Additionally, while enforcing strict convexity via weight projection stabilized fixed-point iterations, it degraded overall transport accuracy, suggesting that a non-convex parameterization (without hard projection) offers a better balance of expressivity and stability.

Significance and Claims
The paper claims that this framework offers a stable, efficient, and scalable alternative to existing neural OT methods. By eliminating adversarial training and the computational overhead of ODE solvers or HJB equation solving over the entire domain, the method provides a "causality-free" mechanism for enforcing optimality conditions. The authors assert that their approach is broadly applicable, capable of handling high-dimensional and class-conditional problems while recovering accurate transport maps in both directions. The work positions itself as a step toward simplifying the optimization landscape of OT, moving away from complex min–max formulations toward a single-objective minimization problem that is theoretically grounded and practically efficient.