Source-linked AI summary

Geometric Operator Learning with Optimal Transport

Xinyi Li, Zongyi Li, Nikola Kovachki, Anima Anandkumar

arXiv:2507.20065v1cs.LG

TL;DR

The paper addresses the computational cost of solving PDEs on complex and diverse geometries. It uses optimal transport to map mesh densities into uniform latent spaces and applies neural operators on lower-dimensional surface representations. OTNO reports faster processing, lower memory usage, and improved accuracy across automotive and variable-geometry flow benchmarks.

  • Problem

    PDE simulation on complex 3D geometries remains computationally burdensome, while existing machine-learning methods have resolution and geometry-generalization limitations.

  • Method

    OTNO formulates geometry embedding as optimal transport from physical mesh density functions to uniform latent densities, enabling neural operators on lower-dimensional surface manifolds.

  • Results

    OTNO achieves 2x-8x faster processing and 2x-8x smaller memory usage than current machine-learning methods, with slightly enhanced accuracy, and significantly outperforms other models on FlowBench.

  • Takeaways & Limitations

    Instance-dependent OT embeddings support efficient surface-based PDE computation and improve accuracy for datasets containing highly variable geometries.

  • Takeaways & Limitations

    OTNO(Plan) has O(n^2) complexity, whereas the lower-complexity OTNO(Map) sacrifices accuracy relative to OTNO(Plan).

Abstract

from arXiv · show

We propose integrating optimal transport (OT) into operator learning for partial differential equations (PDEs) on complex geometries. Classical geometric learning methods typically represent domains as meshes, graphs, or point clouds. Our approach generalizes discretized meshes to mesh density functions, formulating geometry embedding as an OT problem that maps these functions to a uniform density in a reference space. Compared to previous methods relying on interpolation or shared deformation, our OT-based method employs instance-dependent deformation, offering enhanced flexibility and effectiveness. For 3D simulations focused on surfaces, our OT-based neural operator embeds the surface geometry into a 2D parameterized latent space. By performing computations directly on this 2D representation of the surface manifold, it achieves significant computational efficiency gains compared to volumetric simulation. Experiments with Reynolds-averaged Navier-Stokes equations (RANS) on the ShapeNet-Car and DrivAerNet-Car datasets show that our method achieves better accuracy and also reduces computational expenses in terms of both time and memory usage compared to existing machine learning models. Additionally, our model demonstrates significantly improved accuracy on the FlowBench dataset, underscoring the benefits of employing instance-dependent deformation for datasets with highly variable geometries.

1 Introduction

The paper introduces OTNO, which embeds complex surface geometries through instance-dependent optimal transport and computes PDE operators in lower-dimensional latent spaces. This design targets the computational burden and geometry variability that limit existing neural operators, achieving improved efficiency, convergence, and accuracy across automotive flow benchmarks.

  • Motivation: Existing neural operators address complex geometries but remain limited by shared deformation, graph locality, high-dimensional latent spaces, or non-invertible geometry encodings.Neural operators are motivated as resolution-independent alternatives to methods constrained to specific resolutions.
  • Method: OTNO maps mesh density functions to uniform latent densities, generalizing geometry embedding beyond discretized mesh points.The framework formulates this mapping as optimal transport and unifies Monge and Kantorovich formulations.
  • Method: The Sub-Manifold Method embeds d-dimensional surface manifolds into (d −1)-dimensional latent spaces for efficient PDE resolution.For aerodynamic applications, 2D surface designs yield surface pressure and shear-velocity outputs related to total drag.
  • Results: 2x-8x faster processing and 2x-8x smaller memory usage, with slightly enhanced accuracy, are reported for RANS pressure-field prediction versus current machine-learning methods.The method is also approximately 7,000 times faster than traditional approaches.
  • Results: OTNO computes geometry embeddings separately for each shape, improving handling of highly variable geometries on FlowBench.This instance-dependent deformation distinguishes OTNO from GeoFNO and GINO, which learn shared deformation networks across geometries.

2 Problem Setting and Preliminaries

The paper formulates PDE solution learning across varying geometries and focuses on boundary solutions over complex surfaces. It motivates neural operators and positions OT-based geometry encoding as an alternative to interpolation, graph embeddings, and shared deformations.

  • Problem setting: RANS shape-design problems model the fluid domain as the support of a mesh density, with the solid surface forming part of the PDE boundary.The boundary consists of a far-field component and the solid automotive or airfoil surface, which may contain minor non-manifold structures but can be approximated as a manifold.
  • Sub-Manifold Solution Operator: The sub-manifold operator restricts the PDE solution to the boundary surface, targeting surface density inputs and surface solutions such as automotive pressure.This is motivated by applications where relevant outputs are concentrated on object surfaces, although the setting has stated limitations.
  • Problem setting: The target is to learn a solution operator mapping geometry-defining density functions to PDE solutions, with practical interest in boundary quantities on object surfaces.For automotive and airfoil problems, the desired output is the pressure field on the boundary manifold.
  • Neural Operator on Geometric Problems: The paper investigates optimal transport as an encoder and decoder that embeds varying geometries into a uniform latent mesh.This is presented as an alternative to interpolation and deformation-based geometry embeddings.

3 Geometry Embedding as Optimal Transport

The geometry embedding problem is recast as transporting a density-defined physical domain to a uniform measure on a simple latent domain. Transport maps or plans then support latent neural-operator learning and decoding.

  • OT geometry embedding: Complex geometric domains are represented by density functions and embedded into a uniform density on a canonical latent domain through optimal transport.The latent domain may be a simple geometry such as a unit sphere or torus.
  • Transport Map: The transport-map formulation encodes a physical measure through its optimal map from the uniform reference measure, then maps that representation to a latent solution.The approximate physical operator is composed as P ◦ G* ◦ Q.
  • Transport Plan: The transport-plan formulation uses a probability measure with latent and physical marginals, whose marginal representation is encoded before latent operator learning.This provides the plan-based counterpart to the map formulation.
  • OT geometry embedding: The overall approximation combines geometry transport learning with a latent neural operator and physical-space decoding.Numerical OT methods use squared Euclidean transportation cost for these transformations.
  • Relation to existing methods: Geo-FNO and GINO can be interpreted as special map-type and plan-type cases, while OT formulations differ in how geometry transformations are defined and optimized.The comparison contrasts Geo-FNO’s shared, end-to-end deformation with OT’s separately solved transport formulations and GINO’s local graph construction with global transport plans.

4 Optimal Transport Neural Operator (OTNO)

OTNO integrates optimal-transport geometry encoding and decoding with a latent Fourier neural operator. Its implementation organizes surface representations on a 2D grid so the FNO operates in reduced-dimensional space.

  • OTNO architecture: OTNO combines Projection pursuit Monge maps or Sinkhorn transport plans with OT encoders and decoders for neural-operator learning.The model supports both map-based and plan-based transport constructions.
  • OTNO Algorithm: Transported meshes represent physical surfaces after applying the OT map or plan, providing the geometric data used by the latent computation.The primary distinction between map and plan variants is how the latent and transported meshes are generated.
  • OTNO Algorithm: The algorithm builds index mappings, applies the OT encoder, evaluates the latent FNO, decodes solutions to the physical surface, and computes empirical loss across instances.The same overarching algorithm applies to OT maps and OT plans, with differences in latent-mesh and transported-mesh construction.
  • Latent Operator: The latent FNO receives a deformation representation containing latent coordinates, transported surface points, and normal cross-products.The representation Tj = (Ξj, Mj, Hj × Nj(E)) is used as the latent FNO input.

4.2 OTNO - Kantorovich Plan

The OTNO encoder and decoder use Sinkhorn-derived transport plans to connect physical boundary meshes with latent computational grids. The method addresses approximation ambiguity and storage costs through nearest-neighbor strategies and uses reduced-dimensional latent representations for efficient computation.

  • 4.2.1 Transported Mesh: Dense coupling matrices Pj represent Sinkhorn OT plans between latent computational meshes and boundary sampling meshes.The transported mesh is obtained discretely as X′j = PjXj.
  • 4.2.1 Transported Mesh: Directly multiplying dense plans with meshes can reduce prediction accuracy and incurs substantial storage costs.The text attributes reduced accuracy to approximation of the Kantorovich solution plan.
  • 4.2.1 Transported Mesh: The Mean strategy replaces transported points with nearest physical-mesh points, while Max and top-k neighbor strategies provide alternative encoding and decoding schemes.These strategies reduce ambiguity introduced by Sinkhorn coupling and avoid relying on a single direct dense-matrix multiplication.
  • 4.2.2 Latent Mesh: The latent computational grid is oversized relative to the boundary mesh, making the encoder an interpolator and the decoder a selective querier.Figure 5 depicts boundary sampling points against the latent computational grid.
  • 4.2.3 Sinkhorn Method: Sinkhorn entropy regularization makes large-scale OT computationally feasible while improving numerical stability and convergence speed.Increasing β reduces the influence of entropy regularization and brings the transport plan closer to the original Kantorovich solution.
  • Implementation: The implementation uses uniform latent and sampling-mesh masses, with voxel downsampling to limit excessive clustering in high-density regions.The normalization is intended to preserve a more balanced spatial density for aerodynamic prediction.
  • 4.3.1 Transported Mesh: PPMM outputs a transported mesh rather than an invertible function on arbitrary latent-mesh functions, so separate encoding and decoding indices are computed.The OT coupling formulation can support both directions, whereas PPMM itself cannot provide the inverse direction.
  • 4.3.3 Projection Pursuit Monge Map: PPMM estimates large-scale OT maps through iterative one-dimensional transports along the most informative projection directions.Each iteration selects a direction with the greatest projected discrepancy before applying a one-dimensional OT map.

5 Experiments

OTNO is evaluated on car-surface CFD datasets and a diverse-geometry FlowBench dataset. It improves accuracy while reducing time and memory on car tasks, and retains an accuracy advantage on FlowBench where sub-manifold cost savings do not apply.

  • Dataset Setup: The car datasets target surface-dependent quantities, enabling OTNO to embed 3D geometry into a 2D manifold representation.ShapeNet-Car contains 611 designs with 3.7k vertices, while DrivAerNet-Car contains 4k meshes with 200k vertices.
  • 5.1.1 ShapeNet Car Dataset: 6.70% relative error improves on GINO’s 7.21% while reducing total time eightfold and GPU memory sevenfold on ShapeNet-Car.Table 1 reports the ShapeNet-Car pressure-field comparison under matched experimental settings.
  • 5.1.2 DrivAerNet Car Dataset: OTNO(Plan) halves MSE versus RegDGCNN, runs five times faster, and uses 24 times less GPU memory on DrivAerNet-Car.Against GINO, it also achieves slightly lower MSE, marginally higher R2, and fourfold lower time and fivefold lower memory.
  • 5.1.2 DrivAerNet Car Dataset: OTNO(Map) has worse error and time than OTNO(Plan) but substantially lower memory because it does not expand the latent space.The comparison isolates a performance–memory trade-off between the two OT implementations.
  • 5.2 Showcase of Dataset with Diverse Geometries: On FlowBench, OTNO significantly outperforms in accuracy under both global M1 and boundary-focused M2 metrics, especially M1.The boundary metric covers the region conditioned on 0 ≤ SDF ≤ 0.2.
  • 5.2 Showcase of Dataset with Diverse Geometries: FlowBench provides no cost reduction because its solutions are not confined to a boundary sub-manifold, even within the boundary-focused metric.The relevant region remains a finite-width area rather than a lower-dimensional line.
  • 5.2 Showcase of Dataset with Diverse Geometries: Geo-FNO is omitted on FlowBench because its relative L2 errors consistently exceed 60%, while OTNO solves an OT plan or map separately for each shape.The paper associates OTNO’s stronger diverse-geometry performance with instance-specific deformation rather than a shared deformation map.

6 Ablation Studies

The ablations compare OT transport-plan and transport-map implementations, encoder/decoder strategies, normal features, latent mesh shapes, expansion factors, and large-scale sampling behavior. Results favor the Mean and Single strategies, cross-product normals, torus latent meshes, α = 3 expansion for OTNO(Plan), and reveal accuracy and computational limits at larger scales.

  • 6.1 Encoder & Decoder Strategy: “Mean” strategy has better performance than “Max” for using approximate OT transport plans.
  • 6.1 Encoder & Decoder Strategy: “Single” encoder/decoder outperforms “Multi-Enc” and “Multi-Dec” configurations that use multiple nearest elements.
  • 6.2 Normal Features: “Cross Product” normal features perform best among the tested approaches for augmenting the transport-map representation.
  • 6.3 Latent Mesh Shape: The torus latent mesh provides the best performance, while the double sphere yields worse accuracy than the sphere and doubles time cost.
  • 6.4 Expansion Factor: α = 3 achieves the best accuracy for OTNO(Plan) on both DrivAerNet and ShapeNet datasets.
  • 6.5 Large-Scale OT: At approximately 18k points, both OTNO variants reach their lowest accuracy, while increasing sampling size further does not improve accuracy.Large-scale OT precision limits performance; OTNO(Plan) also trains much longer than OTNO(Map), with Sinkhorn and PPMM costs exceeding linear FNO complexity.

7 Discussion

The discussion contrasts OT with conformal and diffeomorphic mappings, emphasizing continuous distribution transport, relaxed smoothness and topology requirements, and flexible latent geometry. Experiments report that OTNO outperforms conformal mapping, torus representations outperform sphere representations, and topology consistency does not improve elasticity results.

  • 7.1 Conformal Mapping: OTNO significantly outperforms the conformal-mapping neural operator on ShapeNet-Car pressure prediction.
  • 7.2 Diffeomorphic Transformations: Continuous OT is presented as better suited than Hungarian discrete assignment for distribution transformations that encode and decode geometric information.
  • 7.2 Diffeomorphic Transformations: Piecewise continuity can suffice: the torus transport is only piecewise continuous yet outperforms the globally continuous sphere representation.
  • 7.2 Diffeomorphic Transformations: Expanding the latent mesh improves performance within a range, showing that invertibility is not always beneficial for this embedding.
  • 7.2 Diffeomorphic Transformations: OT does not require topological consistency: the torus performs best despite a topology change, and enforcing consistency does not improve the 2D elasticity results.

8 Conclusion

The paper introduces OTNO, which embeds physical geometric densities into uniform latent densities and supports both transport-plan and transport-map formulations. Surface-focused automotive and aerospace experiments use 2D latent representations, while FlowBench results show strong accuracy across diverse geometries; scalability and latent-manifold selection remain limitations.

  • OTNO maps physical-domain density functions to uniform densities on regular latent geometries using Kantorovich-based Plan and Monge-based Map implementations.
  • 2D parameterized latent meshes enable lower-dimensional computation for 3D surface inputs in automotive and aerospace applications.
  • OTNO significantly outperforms other models in accuracy on FlowBench, which contains a wider variety of shapes.
  • OTNO(Plan) has O(n^2) complexity and higher accuracy than OTNO(Map), whose O(n^3/2 log(n)) complexity reduces computational cost but diminishes accuracy.
  • Selecting the optimal latent sub-manifold remains open because experiments identify a non-canonical optimum rather than the intuitively favored spherical topology.

Appendix A. Convergence Study

The convergence study compares latent resolutions and dimensions across models. GeoFNO performs better with a 3D sphere latent mesh than with a 2D square, while the corresponding latent-resolution settings use method-specific optimal expansion factors.

  • Appendix A. Convergence Study: GeoFNO(3D) with a sphere latent mesh clearly outperforms GeoFNO(2D) with a square latent mesh.
  • Appendix A. Convergence Study: Table 12 reports latent-space resolution configurations at different sampling rates with optimal expansion factors applied for convergence.

Appendix B. DrivAerNet: Drag Coefficient

The DrivAerNet task predicts car drag coefficient, using pressure and wall-shear terms defined over the car surface. Because the dataset has high Reynolds number, the model loss retains only the pressure-drag term.

  • The target is the drag coefficient, a metric for evaluating car design on the DrivAerNet dataset.The coefficient is defined for the car surface using pressure, wall shear stress, inlet-flow direction and speed, and the frontal enclosing area.
  • The drag coefficient combines pressure drag and shear drag contributions over the car surface.
  • At Reynolds number roughly 9.39 × 10^6, the model loss considers only the pressure-drag term because pressure drag is significantly higher than shear drag.

Appendix C. FlowBench

FlowBench evaluates geometry operator learning across diverse shapes using flow-field inputs and outputs under full-space and boundary metrics. The experiments use the Easy split and report visualized results for three shape-diversity groups.

  • FlowBench setup: The FlowBench inputs are Reynolds number, a binary geometry mask, and signed distance function, while outputs are x- and y-velocity and pressure.
  • FlowBench setup: The evaluation uses the Easy case with an 80-20 random split rather than the Hard case split by Reynolds number.
  • Shape groups: The dataset is evaluated across G1 NURBS shapes, G2 spherical-harmonic shapes, and G3 non-parametric grayscale-sampled shapes, whose diversity exceeds that of the car designs.
  • Results visualization: Figures 10 and 11 visualize FlowBench results under the M1 full-space metric and M2 boundary metric, respectively.
Loading 2507.20065v1…