Learning sufficient low-dimensional structures through conditional optimal transport

Technical Summary: Learning Sufficient Low-Dimensional Structures through Conditional Optimal Transport

Problem Statement
Sufficient Dimension Reduction (SDR) aims to find a low-dimensional representation of covariates $X$ that preserves the entire conditional law of a response $Y$, satisfying $Y \perp\!\!\!\perp X \mid R(X)$. Classical linear SDR methods often rely on indirect targets like conditional means or inverse regression moments, which may overlook information contained in heteroscedasticity, tail behavior, or multimodality. While recent generative approaches (e.g., GenSDR, FlowSDR) model the conditional law directly, they often operate in Euclidean settings or assume specific structural properties. This paper addresses the challenge of learning sufficient reductions in a general setting where both covariates and responses may reside in separable Hilbert spaces, utilizing Conditional Optimal Transport (COT) as a geometric framework to describe the variation of the conditional response law.

Methodology
The proposed framework, SDR-COT, leverages the geometric properties of optimal transport to characterize the conditional independence structure.

1. Conditional Optimal Transport (COT) Formulation:
   The method models the transition from an independent reference response $Y_0 \sim \rho$ to the observed response $Y_1$ given covariates $X$. This is framed as a triangular transport map $T^\star(x, y) = (x, T^\star_Y(x, y))$ that pushes the product measure $\mu_0 = \eta \otimes \rho$ to the joint law $\mu_1 = P(X, Y_1)$ while preserving the covariate coordinate.

2. Dynamic COT and Velocity Fields:
   The authors utilize the Benamou–Brenier dynamic formulation. They establish that, under a quadratic cost and a product-source setting, the Monge interpolation admits a Borel current-state velocity representative on every truncated time interval $[0, 1-\tau]$, even without assuming global injectivity of the terminal map. This velocity field $v^\star_t$ drives the displacement interpolation between the source and target laws.

3. Factorization Theorem:
   A core theoretical result demonstrates that if a sufficient representation $R(X)$ exists (i.e., $Y_1 \perp\!\!\!\perp X \mid R(X)$), then both the static optimal transport map $T^\star_Y$ and the dynamic velocity field $v^\star_t$ factor through $R$. Specifically:
   $$T^\star_Y(x, y) = G^\star_Y(R(x), y)$$
   $$v^\star_t(x, y) = (0, u^\star_Y(t, R(x), y))$$
   This implies that the dependence of the transport mechanism on the covariate is entirely mediated by the sufficient reduction.

4. Estimation via Flow Matching:
   Motivated by the factorization, the authors propose a Conditional Flow Matching (CFM) criterion. Instead of estimating the full conditional map, the algorithm learns a reduction map $R_\theta$ and a velocity network $u_\phi$ parameterized by the reduced covariate. The objective minimizes the squared error between the true displacement and the predicted velocity:
   $$L(\theta_u, \theta_R) = \int_0^{T_\tau} \mathbb{E} \left[ \| V - u_{\theta_u}^Y(t, R_{\theta_R}(X), Y_t) \|^2 \right] dt$$
   where $V = Y_1 - Y_0$ and $Y_t$ is the linear interpolation. To handle the unobserved optimal coupling, the method employs a relaxed empirical coupling using an anisotropic cost function that penalizes covariate mismatch heavily.

Key Contributions

1. Dynamic COT Optimality without Global Injectivity:
   The paper proves a Benamou–Brenier type optimality statement for dynamic COT in separable Hilbert spaces with quadratic cost. Unlike prior work that assumes global injectivity of the triangular Monge map to define velocities, this work establishes the existence of a Borel current-state velocity on truncated intervals $[0, 1-\tau]$ using cyclic monotonicity, removing the need for a priori injectivity assumptions.

2. SDR-COT Factorization Theorem:
   The authors prove that under mild regularity conditions, the response component of the optimal triangular COT map and the induced velocity field factor through any sufficient representation. This structural result formally links the conditional independence assumption of SDR to the geometric properties of optimal transport.

3. Unified Estimation Framework:
   The paper develops a population-level flow-matching criterion and provides algorithms for both Euclidean and functional data settings.

   • Euclidean: Uses standard neural networks for the reduction and velocity.
   • Functional: Handles functional covariates via basis projection and functional responses via Neural Operators (specifically Fourier Neural Operators).
   • Relaxed Coupling: Introduces a relaxed empirical coupling strategy to approximate the conditional transport map from data.

4. Statistical Consistency for Linear SDR:
   The paper establishes consistency for linear SDR in both Euclidean and Hilbert settings:

   • Euclidean: Consistency is proven using Caffarelli's contraction theorem to bound the Lipschitz constant of the velocity, allowing identification of the central subspace via a zero-loss condition.
   • Hilbert (Functional): Consistency is established without ambient Lipschitz continuity by leveraging Cameron–Martin-valued Sobolev velocities and the uniqueness of the Gaussian continuity equation.

Results

• Theoretical: The factorization theorem confirms that the conditional independence structure is inherited by the optimal transport map and velocity. The consistency theorems guarantee that, under correct specification and regularity, the estimated linear reduction converges to the true central subspace in projection distance.
• Empirical: Numerical studies on synthetic nonlinear Euclidean data, linear Euclidean data, functional covariates, and function-on-function data demonstrate that SDR-COT is competitive with established benchmarks. Notably, the method performs well when sufficient information is contained in distributional features beyond the conditional mean (e.g., variance or shape), where classical moment-based methods may fail.
• Real Data: In an analysis of Capital Bikeshare demand (functional data), SDR-COT achieved lower test errors than Functional PCA (FPCA) and weak-conditional-moment-based competitors across most reduction dimensions.

Significance and Claims
The paper claims to provide a geometrically grounded approach to SDR that directly targets the conditional law rather than its moments. By utilizing Conditional Optimal Transport, the method naturally captures complex distributional dependencies (heteroscedasticity, multimodality) that are often missed by inverse regression or conditional mean-based methods.

The authors are modest regarding the general nonlinear case: while the population criterion is "exhaustive" (any representation containing the central $\sigma$-algebra attains the optimal value), the criterion alone does not guarantee the selection of a minimal representation without additional constraints (such as fixing the output dimension). Therefore, the rigorous consistency results are explicitly limited to the linear setting with a fixed structural dimension. The work bridges the gap between optimal transport theory and sufficient dimension reduction, offering a scalable, simulation-free estimation procedure that extends to infinite-dimensional functional data.