Tristan Luca Saidi

The Sinkhorn divergence, a debiased form of entropic optimal transport, has been widely used in various scientific disciplines due to its computational and statistical tractability. Unlike the quadratic Wasserstein distance, however, the Sinkhorn divergence is not a metric and it does not naturally admit a formal Riemannian structure. These properties are central to many applications of optimal transport methods that utilize interpolation, barycentric averaging, parallel transport or other geometric operations. Recent work by Lavenant et al. (2024) addresses these shortcomings by constructing a formal Riemannian tangent space and a metric tensor on the space of probability measures that quadratically approximates the Sinkhorn divergence. In this work, we extend this framework with two major threads of contributions. For a suitable class of finite-smoothness cost functions, we prove Hadamard differentiability of the Sinkhorn metric tensor and we utilize the technique of Gigli (2012) to derive the Levi-Civita covariant derivative in this geometry; we then use the covariant derivative to obtain the geodesic and parallel transport equations, each of which amount to a second order ODE on a fixed Reproducing Kernel Hilbert Space (RKHS) with smooth coefficients. Our second contribution is to linearize the solutions to the geodesic and parallel transport equations using standard tools from Banach space ODE theory. We show that these linearizations imply $\sqrt{n}$-limiting Gaussian processes centered at population values for the plugin estimators of the exponential map, the logarithmic map, parallel transport and the induced Sinkhorn distance. Moreover, we show that the nonparametric bootstrap is consistent for all objects mentioned. To the best of our knowledge, these are the first $\sqrt{n}$-limit theorems and bootstrap consistency results for canonical geometric operations in a transport-based geometry on the space of probability measures.

In this work we propose and analyze a procedure for estimating a generalized single-index regression model. This generalized single-index regression framework models the conditional mean of the responses $y$ as a function of the covariates $x$ where the mapping is assumed to have the form $f \circ g$, where $g$ is $\beta$-Hölder with $\beta \geq 2$ and $f$ is a $C^2$ function. We propose a general procedure that operates under minimal assumptions about the distribution of covariates $x$, requiring only boundedness, sufficient smoothness and absolute continuity with respect to Lebesgue. Our procedure estimates the gradient vector $\hat v(x) \approx \nabla g(x)$, and then uses this to produce an estimate of the response function $f$ via projected kernel regression. We derive central limit theorems for gradient vector estimates $\hat{v}_x$ and regression estimates $\hat{f}(x)$, and we derive rates of convergence for our approach, illustrating that it indeed achieves the minimax rate pointwise when $\beta = 2$. Finally, we instantiate this model in the situation where $g = \Pi$ is a closest point projection onto a compact and regular embedded $C^3$ curve $\gamma$, a setting recently studied in the literature. We also supplement our theoretical results with simulations to demonstrate the utility of our procedure.

Many scientific systems, such as cellular populations or economic cohorts, are naturally described by probability distributions that evolve over time. Predicting how such a system would have evolved under different forces or initial conditions is fundamental to causal inference, domain adaptation, and counterfactual prediction. However, the space of distributions often lacks the vector space structure on which classical methods rely. To address this, we introduce a general notion of parallel dynamics at a distributional level. We base this principle on parallel transport of tangent dynamics along optimal transport geodesics and call it “Wasserstein Parallel Trends”. By replacing the vector subtraction of classic methods with geodesic parallel transport, we can provide counterfactual comparisons of distributional dynamics in applications such as causal inference, domain adaptation, and batch-effect correction in experimental settings. The main mathematical contribution is a novel notion of fanning scheme on the Wasserstein manifold that allows us to efficiently approximate parallel transport along geodesics while also providing the first theoretical guarantees for parallel transport in the Wasserstein space. We also show that Wasserstein Parallel Trends recovers the classic parallel trends assumption for averages as a special case and derive closed-form parallel transport for Gaussian measures. We deploy the method on synthetic data and two single-cell RNA sequencing datasets to impute gene-expression dynamics across biological systems.

The Hellinger-Kantorovich (HK) space provides a natural geometry for nonnegative measures with varying total mass, but its differential-geometric structure is less well understood than that of the closely related Wasserstein space of probability measures. In this paper, we take a step toward resolving this issue. We show that the cone representation of the HK geometry via the Wasserstein metric preserves the local Riemannian geometry along a class of lifted geodesics. Specifically, we give a constructive procedure that produces a Wasserstein geodesic on the cone along which the HK Riemannian geometry is preserved pointwise, yielding an explicit isometry of tangent spaces between HK geodesics and their Wasserstein lifts. This connection makes many Wasserstein-geometric tools available for HK computations. Concretely, we use it to approximate parallel transport on HK space by lifting to the cone and applying recently developed Wasserstein parallel transport tools, circumventing the high-dimensional PDE arising from the HK covariant derivative. We also derive closed-form expressions for the covariant derivative and parallel transport on Euclidean metric cones, using the theory of warped-product manifolds. Finally, we present simulations illustrating the behavior of parallel geodesics in HK space, which reveal that the HK geometry couples spatial and mass variation through the geometry of the cone -- a feature with nontrivial implications for applied use of the framework.

Stochastic Neighbor Embedding (SNE) algorithms like UMAP and tSNE often produce visualizations that do not preserve the geometry of noisy and high dimensional data. In particular, they can spuriously separate connected components of the underlying data submanifold and can fail to find clusters in well-clusterable data. To address these limitations, we propose EmbedOR, a SNE algorithm that incorporates discrete graph curvature. Our algorithm stochastically embeds the data using a curvature-enhanced distance metric that emphasizes underlying cluster structure. Critically, we prove that the EmbedOR distance metric extends consistency results for tSNE to a much broader class of datasets. We also describe extensive experiments on synthetic and real data that demonstrate the visualization and geometry-preservation capabilities of EmbedOR. We find that, unlike other SNE algorithms and UMAP, EmbedOR is much less likely to fragment continuous, high-density regions of the data. Finally, we demonstrate that the EmbedOR distance metric can be used as a tool to annotate existing visualizations to identify fragmentation and provide deeper insight into the underlying geometry of the data.

We introduce ORC-ManL, a new algorithm to prune spurious edges from nearest neighbor graphs using a criterion based on Ollivier-Ricci curvature and estimated metric distortion. Our motivation comes from manifold learning: we show that when the data generating the nearest-neighbor graph consists of noisy samples from a low-dimensional manifold, edges that shortcut through the ambient space have more negative Ollivier-Ricci curvature than edges that lie along the data manifold. We demonstrate that our method outperforms alternative pruning methods and that it significantly improves performance on many downstream geometric data analysis tasks that use nearest neighbor graphs as input. Specifically, we evaluate on manifold learning, persistent homology, dimension estimation, and others. We also show that ORC-ManL can be used to improve clustering and manifold learning of single-cell RNA sequencing data. Finally, we provide empirical convergence experiments that support our theoretical findings.

A major challenge in materials science is the determination of the structure of nanometre-sized objects. Here we present an approach that uses a generative machine learning model based on diffusion processes that are trained on 45,229 known structures. The model factors measured the diffraction pattern as well as the relevant statistical priors on the unit cell of atomic cluster structures. Conditioned only on the chemical formula and the information-scarce finite-sized broadened powder diffraction pattern, we find that our model, PXRDnet, can successfully solve the simulated nanocrystals as small as 10 Å across 200 materials of varying symmetries and complexities, including structures from all seven crystal systems. We show that our model can successfully and verifiably determine structural candidates four out of five times, with an average error among these candidates being only 7% (as measured by the post-Rietveld refinement R-factor). Furthermore, PXRDnet is capable of solving structures from noisy diffraction patterns gathered in real-world experiments. We suggest that data-driven approaches, bootstrapped from theoretical simulation, will ultimately provide a path towards determining the structure of previously unsolved nanomaterials.

We present a method for enabling Reinforcement Learning of motor control policies for complex skills such as dexterous manipulation. We posit that a key difficulty for training such policies is the difficulty of exploring the problem state space, as the accessible and useful regions of this space form a complex structure along manifolds of the original high-dimensional state space. This work presents a method to enable and support exploration with Sampling-based Planning. We use a generally applicable non-holonomic Rapidly-exploring Random Trees algorithm and present multiple methods to use the resulting structure to bootstrap model-free Reinforcement Learning. Our method is effective at learning various challenging dexterous motor control skills of higher difficulty than previously shown. In particular, we achieve dexterous in-hand manipulation of complex objects while simultaneously securing the object without the use of passive support surfaces. These policies also transfer effectively to real robots. A number of example videos can also be found on the project website: sbrl.cs.columbia.edu

In this paper, we present a novel method for achieving dexterous manipulation of complex objects, while simultaneously securing the object without the use of passive support surfaces. We posit that a key difficulty for training such policies in a Reinforcement Learning framework is the difficulty of exploring the problem state space, as the accessible regions of this space form a complex structure along manifolds of a high-dimensional space. To address this challenge, we use two versions of the non-holonomic Rapidly-Exploring Random Trees algorithm; one version is more general, but requires explicit use of the environment's transition function, while the second version uses manipulation-specific kinematic constraints to attain better sample efficiency. In both cases, we use states found via sampling-based exploration to generate reset distributions that enable training control policies under full dynamic constraints via model-free Reinforcement Learning. We show that these policies are effective at manipulation problems of higher difficulty than previously shown, and also transfer effectively to real robots. Videos of the real-hand demonstrations can be found on the project website: https://sbrl.cs.columbia.edu/