Manifold-valued learning

GeoJAX separates the geometry of observations from the statistical method applied to them. A geometry defines point validity, distance, logarithmic and exponential maps, and optional equivariant embeddings. The learning layer validates a collection once and then consumes those operations without making assumptions about whether a point is a vector, matrix, or Product pytree. This follows the metric-statistics view of random objects while retaining JAX transformations for differentiable primitives [Bronstein et al., 2021, Dubey and Müller, 2019, Fréchet, 1948].

Canonical datasets

For an array geometry with event shape M.shape, the canonical dataset layout is

\[ \mathtt{batch\_shape} + (n,) + \mathtt{M.shape}. \]

The axis immediately before the event dimensions is the sample axis. Product data use the same nested tuple, list, or dictionary as M.factors; every leaf has the same batch shape and sample count but its own event shape.

import jax

from geojax.geometry import Product, SPDLogEuclidean, Sphere
from geojax.learning import as_manifold_data

M = Product({
    "direction": Sphere(3),
    "covariance": SPDLogEuclidean((2, 2)),
})
values = M.random_point(jax.random.key(0), sample_shape=(32,))
data = as_manifold_data(M, values)

as_manifold_data defaults to membership validation. It rejects malformed, nonfinite, or off-manifold observations. Passing repair=True explicitly calls M.project; projection is never silent. sample_axis must be provided when an input does not already use the canonical layout, so an adapter never guesses which axis represents observations.

Each ManifoldData object is bound to the exact geometry instance that validated it. This matters when two geometries share an event shape but use different metrics or parameters. Reusing the object with that instance skips validation already performed; requesting a stronger level upgrades the check in the order shape, finite, belongs. To use the same canonical values with another geometry, adapt data.values explicitly under that geometry.

Alternate representations

The adapter accepts named representations only when their conversion has a defined geometric meaning. Examples include hyperspherical angles, Poincaré coordinates, Grassmann projectors, SPD Cholesky factors, covariance matrices for correlation geometries, SO(3) quaternions, SE(2)/SE(3) twists, raw Kendall landmarks, and low-rank factors. Product representations may themselves be pytrees.

from geojax.geometry import Grassmann
from geojax.learning import as_manifold_data

M = Grassmann((5, 2))
frames = as_manifold_data(M, projectors, representation="projector")

A spanning basis represents the same Grassmann point and may be orthonormalized directly. A nonorthonormal Stiefel frame represents a different ambient matrix, so it is accepted only through repair=True. User-defined representations can be added with register_manifold_data_adapter.

Exact capabilities

Distance-only methods require operation_kind("dist") == "exact". Methods that update intrinsic centers additionally require exact logarithmic and exponential maps. GeoJAX raises LearningCapabilityError when a geometry only offers a retraction proxy or numerical-local logarithm. This prevents an algorithm named after a geodesic quantity from silently changing its mathematical objective.

Required geometry	Learning methods
Exact distance	pairwise distances, neighbors, nearest-centroid and k-NN prediction, medoids, hierarchy, spectral and graph learning, kernel regression, MDS, Isomap, kernel PCA, Sammon, t-SNE, PHATE, transport, energy and kernel tests
Exact distance, logarithm, exponential	interpolation, Fréchet and robust summaries, geodesic and local Fréchet regression, k-means and mini-batch summaries, mean shift, CLRQ, PGA, barycentric coding, dictionary learning, enclosing balls, Fréchet ANOVA, paired tests
Equivariant embedding	Riemannian manifold metric learning

Differentiable primitives

geojax.learning.pairwise_distances() returns ordinary or squared exact distances. Collections broadcast over leading batch axes, Product values remain pytrees, and block_size limits temporary right-hand blocks while retaining a dense result.

from geojax.learning import pairwise_distances

squared = pairwise_distances(M, queries, prototypes, squared=True)
logits = -squared

Geodesic interpolation evaluates

\[\gamma(t)=\operatorname{Exp}_x\!\left(t\operatorname{Log}_x(y)\right),\]

and tangent_space_map composes a user transformation between source and target tangent spaces. These wrappers are compatible with jax.jit, jax.vmap, and differentiation whenever the selected geometry operations are.

Statistical algorithms

Fréchet means minimize weighted squared distance and use GeoJAX’s manifold optimizer. Medians use a guarded Riemannian Weiszfeld iteration. Clustering provides intrinsic Lloyd updates, medoids, valid metric-space linkage rules, spectral graph methods, mean shift, and competitive quantization [Karcher, 1977, Lloyd, 1982, Zelnik-Manor and Perona, 2005].

Dimension-reduction methods consume exact pairwise distances. PGA differs: it forms its covariance Gram matrix with M.inner at the Fréchet mean, so the result remains valid for Product manifolds and metrics that are not ambient Frobenius metrics [Fletcher et al., 2004]. Classical MDS and kernel methods report negative eigenvalue mass rather than pretending every manifold distance is Euclidean.

The exact empirical Wasserstein routine solves the finite weighted transport problem and reports its transport plan, marginal residuals, reduced costs, and duality gap. sinkhorn_divergence is a separate optional OTT-JAX operation; regularization is never labeled exact [Cuturi, 2013].

Inference separates three null hypotheses. Fréchet ANOVA compares object-valued populations through their means and variances, the Biswas–Ghosh statistic uses only interpoint distances, and the Wasserstein test compares empirical measures [Biswas and Ghosh, 2014, Dubey and Müller, 2019].

For group proportions \(\gamma_j\), Fréchet variances \(V_j\), variance estimators \(\sigma_j^2\), and

\[ F=V_{\mathrm{pooled}}-\sum_j\gamma_jV_j, \qquad U=\sum_{j<k} \frac{\gamma_j\gamma_k(V_j-V_k)^2}{\sigma_j^2\sigma_k^2}, \]

frechet_anova reports the Dubey–Müller statistic

\[ T_n= \frac{nU}{\sum_j\gamma_j/\sigma_j^2} + \frac{nF^2}{\sum_j\gamma_j^2\sigma_j^2}. \]

The asymptotic calibration requires the regularity conditions of the cited result; the permutation option instead calibrates the same statistic under exchangeable group labels.

riemannian_metric_learning forms similar- and dissimilar-pair scatter matrices after an equivariant embedding, regularizes both matrices, and uses their weighted log-Euclidean closed form. The balance parameter controls the relative contributions; the default midpoint matches the core RMML construction [Zhu et al., 2018].

Supervised prediction

Distance classifiers make the smallest geometric commitment. The nearest- centroid rule estimates one weighted Fréchet mean \(\widehat\mu_c\) per class and predicts

\[ \widehat c(x)=\arg\min_c d(x,\widehat\mu_c), \]

while knn_classifier votes among the \(k\) closest training observations. Both work with any exact-distance geometry, including nested Product data.

Tangent classifiers first compute a reference Fréchet mean \(\widehat\mu\) and an intrinsic metric Gram matrix

\[ G_{ij} =\left\langle\operatorname{Log}_{\widehat\mu}(x_i), \operatorname{Log}_{\widehat\mu}(x_j)\right\rangle_{\widehat\mu}. \]

Its positive eigenspace gives metric-orthonormal coordinates, rather than an ambient Frobenius flattening. tangent_space_logistic_regression fits a multinomial softmax model there; tangent_space_discriminant_analysis provides regularized LDA and QDA. These models are local to the selected logarithm chart and should not be interpreted across a cut locus.

Manifold-valued responses

geodesic_regression fits the one-predictor curve

\[ \widehat Y(t) =\operatorname{Exp}_{p}\!\left((t-\overline t)v\right) \]

by intrinsic least squares, profiling \(v\in T_p\mathcal M\) while optimizing the intercept \(p\) [Fletcher, 2013]. This is a parametric geodesic model, so systematic curvature away from one geodesic remains in the residuals.

local_polynomial_regression instead estimates a conditional Fréchet mean. For local-linear smoothing at \(t\), its signed weights are

\[ s_i(t)=K_h(t_i-t) \frac{S_2-(t_i-t)S_1}{S_0S_2-S_1^2}, \qquad S_r=\sum_i K_h(t_i-t)(t_i-t)^r, \]

and the prediction minimizes \(\sum_i s_i(t)d^2(p,y_i)\). Degree zero uses positive Nadaraya–Watson weights. Signed local-linear objectives need not be globally convex on a general manifold, so GeoJAX initializes them with the positive local Fréchet mean and reports no global-optimum claim [Petersen and Müller, 2019].

Uncertainty and testing

bootstrap_frechet_mean resamples the empirical measure, recomputes intrinsic means, and reports the requested quantile of replicate distances from the original estimate as a bootstrap geodesic ball. It is an approximate percentile region, not a curvature-corrected confidence set.

The energy statistic compares between-sample and within-sample distances [Székely and Rizzo, 2013]. On a completely general metric space, equality characterization requires an appropriate negative-type condition. The MMD test similarly requires a positive-semidefinite kernel [Gretton et al., 2012]: GeoJAX checks the observed Gram eigenvalues by default because an RBF of squared geodesic distance is not universally PSD. The check uses the larger of the requested tolerance and a dtype-aware eigensolver backward-error bound, so roundoff-scale negative eigenvalues in float32 are not mistaken for a mathematically indefinite kernel. paired_frechet_test applies sign flips to paired tangent displacements and therefore assumes exchangeability under those flips.

Scalable summaries

streaming_frechet_mean performs the inductive update

\[ \mu_t =\operatorname{Exp}_{\mu_{t-1}} \!\left(\frac{w_t}{\sum_{j\leq t}w_j} \operatorname{Log}_{\mu_{t-1}}(x_t)\right). \]

It is exact for a weighted Euclidean mean and order-dependent on a curved manifold. By default the first positive-weight observation initializes the recursion. A user-provided initial_point contributes only when accompanied by an explicit positive initial_weight, so no undocumented pseudo-observation is introduced. minibatch_frechet_mean and minibatch_kmeans use shuffled batches and decaying log-map steps, following the stochastic Riemannian optimization view [Bonnabel, 2013]. They require an explicit key and return objective and update histories; they are approximations to the corresponding full-batch estimators.

Barycentric coding and dictionaries

For atoms \(D_1,\ldots,D_m\) and an observation \(x\), intrinsic barycentric coding solves

\[ \min_{w\in\Delta^{m-1}} \frac12\left\|\sum_{j=1}^m w_j\operatorname{Log}_x(D_j)\right\|_x^2 +\frac{\lambda}{2}\|w\|_2^2. \]

geodesic_barycentric_coding uses projected gradient steps on the simplex and reconstructs each observation as the weighted Fréchet mean of its atoms. manifold_dictionary_learning alternates these codes with a Product-manifold optimization of the fixed-code intrinsic residual. Backtracking accepts only atom updates that do not increase the reported reconstruction objective [Ho et al., 2013]. These are dense barycentric codes; the simplex \(\ell_1\) norm is constant, so the routine does not claim lasso sparsity. When sample weights are supplied, atom optimization, backtracking acceptance, and the returned objective all use the same normalized weighted criterion.

Robust and graph learning

The trimmed mean repeatedly retains the smallest squared geodesic residuals until the requested fraction of total sample-weight mass remains. If the trim boundary crosses an observation’s weight, that boundary weight is included partially. This reduces to ordinary count trimming under equal weights. geodesic_m_estimator uses Huber, Cauchy, or Tukey residual weights in guarded iteratively reweighted Fréchet updates. geodesic_spatial_depth evaluates

\[ D(x)=1-\left\| \sum_i w_i \frac{\operatorname{Log}_x(x_i)}{d(x,x_i)} \right\|_x, \]

with the coincident contribution set to zero; distance ranks use midranks around an intrinsic median [Fletcher et al., 2009]. A weighted trimmed fit reports its objective with weights renormalized over the retained measure.

For partially labeled data, label_propagation diffuses class scores over a geodesic-distance affinity graph while clamping known labels [Zhou et al., 2003]. manifold_regularized_regression solves the dense transductive objective

\[ \sum_{i\in L}(f_i-y_i)^2 +\lambda_A\|f\|_2^2 +\lambda_I f^\top L_G f, \]

where \(L_G\) is the graph Laplacian [Belkin et al., 2006]. The current result predicts only the supplied graph vertices; it does not claim the RKHS out-of-sample extension of the full manifold-regularization framework. For a \(k\)-nearest-neighbor graph, the diagonal is excluded before sorting, so tied duplicate observations cannot accidentally select themselves as neighbors.

Learning roadmap

The adapter-first core now covers the former priority-one roadmap, intrinsic dictionary learning, robust analysis, and graph semi-supervision. The next candidate families are deliberately narrower:

Method family	Candidate public methods	Geometric requirements
Graph embeddings	diffusion maps and Laplacian eigenmaps	exact distance or a user-supplied affinity
Supervised reduction	discriminant tangent components and supervised distance embeddings	exact logarithm and metric or an equivariant embedding
Model assessment	intrinsic silhouette, gap statistic, and resampling-based clustering stability	exact distance; randomized procedures require keys
Transport statistics	entropic Wasserstein barycenters and transport-based clustering	optional differentiable Sinkhorn backend

Curves and spatial-analysis routines remain outside this roadmap. Sparse and approximate-neighbor backends should be designed as separate scalability work rather than hidden changes to the dense methods documented in this release.

Computational limits

The learning layer is dense in this release. Pairwise data need \(O(n^2)\) storage, while Floyd-Warshall Isomap and hierarchical updates can need \(O(n^3)\) work. block_size lowers temporary pairwise memory but does not make the returned matrix sparse. High-level assignments, permutations, graph construction, and exact transport are intentionally not promised to be end-to-end differentiable.

References

[BNS06]

Mikhail Belkin, Partha Niyogi, and Vikas Sindhwani. Manifold regularization: a geometric framework for learning from labeled and unlabeled examples. Journal of Machine Learning Research, 7(85):2399–2434, 2006. URL: https://www.jmlr.org/papers/v7/belkin06a.html.

[BG14]

Munmun Biswas and Anil K. Ghosh. A nonparametric two-sample test applicable to high dimensional data. Journal of Multivariate Analysis, 123:160–171, 2014. doi:10.1016/j.jmva.2013.09.004.

[Bon13]

Silvère Bonnabel. Stochastic gradient descent on Riemannian manifolds. IEEE Transactions on Automatic Control, 58(9):2217–2229, 2013. doi:10.1109/TAC.2013.2254619.

[BBCVelivckovic21]

Michael M. Bronstein, Joan Bruna, Taco Cohen, and Petar Veličković. Geometric deep learning: grids, groups, graphs, geodesics, and gauges. arXiv preprint arXiv:2104.13478, 2021. doi:10.48550/arXiv.2104.13478.

[Cut13]

Marco Cuturi. Sinkhorn distances: lightspeed computation of optimal transport. In Advances in Neural Information Processing Systems, volume 26, 2292–2300. 2013.

[DMuller19] (1,2)

Paromita Dubey and Hans-Georg Müller. Fréchet analysis of variance for random objects. Biometrika, 106(4):803–821, 2019. doi:10.1093/biomet/asz052.

[Fle13]

P. Thomas Fletcher. Geodesic regression and the theory of least squares on Riemannian manifolds. International Journal of Computer Vision, 105(2):171–185, 2013. doi:10.1007/s11263-012-0591-y.

[FLPJ04]

P. Thomas Fletcher, Conglin Lu, Stephen M. Pizer, and Sarang Joshi. Principal geodesic analysis for the study of nonlinear statistics of shape. IEEE Transactions on Medical Imaging, 23(8):995–1005, 2004. doi:10.1109/TMI.2004.831793.

[FVJ09]

P. Thomas Fletcher, Suresh Venkatasubramanian, and Sarang Joshi. The geometric median on Riemannian manifolds with application to robust atlas estimation. NeuroImage, 45(1, Supplement 1):S143–S152, 2009. doi:10.1016/j.neuroimage.2008.10.052.

[Frechet48]

Maurice Fréchet. Les éléments aléatoires de nature quelconque dans un espace distancié. Annales de l'Institut Henri Poincaré, 10(4):215–310, 1948. URL: https://www.numdam.org/item/AIHP_1948__10_4_215_0/.

[GBR+12]

Arthur Gretton, Karsten M. Borgwardt, Malte J. Rasch, Bernhard Schölkopf, and Alexander Smola. A kernel two-sample test. Journal of Machine Learning Research, 13(25):723–773, 2012. URL: https://www.jmlr.org/papers/v13/gretton12a.html.

[HXV13]

Jeffrey Ho, Yuchen Xie, and Baba Vemuri. On a nonlinear generalization of sparse coding and dictionary learning. In Proceedings of the 30th International Conference on Machine Learning, volume 28 of Proceedings of Machine Learning Research, 1480–1488. 2013. URL: https://proceedings.mlr.press/v28/ho13a.html.

[Kar77]

Hermann Karcher. Riemannian center of mass and mollifier smoothing. Communications on Pure and Applied Mathematics, 30(5):509–541, 1977. doi:10.1002/cpa.3160300502.

[Llo82]

Stuart P. Lloyd. Least squares quantization in PCM. IEEE Transactions on Information Theory, 28(2):129–137, 1982. doi:10.1109/TIT.1982.1056489.

[PMuller19]

Alexander Petersen and Hans-Georg Müller. Fréchet regression for random objects with Euclidean predictors. The Annals of Statistics, 47(2):691–719, 2019. doi:10.1214/17-AOS1624.

[SzekelyR13]

Gábor J. Székely and Maria L. Rizzo. Energy statistics: a class of statistics based on distances. Journal of Statistical Planning and Inference, 143(8):1249–1272, 2013. doi:10.1016/j.jspi.2013.03.018.

[ZMP05]

Lihi Zelnik-Manor and Pietro Perona. Self-tuning spectral clustering. In Advances in Neural Information Processing Systems, volume 17, 1601–1608. 2005.

[ZBL+03]

Dengyong Zhou, Olivier Bousquet, Thomas Navin Lal, Jason Weston, and Bernhard Schölkopf. Learning with local and global consistency. In Advances in Neural Information Processing Systems, volume 16. 2003. URL: https://proceedings.neurips.cc/paper/2003/hash/87682805257e619d49b8e0dfdc14affa-Abstract.html.

[ZCH+18]

Pengfei Zhu, Hao Cheng, Qinghua Hu, Qilong Wang, and Changqing Zhang. Towards generalized and efficient metric learning on Riemannian manifold. In Proceedings of the Twenty-Seventh International Joint Conference on Artificial Intelligence, 3235–3241. 2018. doi:10.24963/ijcai.2018/449.