Riemannian Metric Matching for Scalable Geometric Modeling of Distributions

Summary

The paper proposes Riemannian metric matching: a denoising-style training objective that learns a neural surrogate for the carré du champ (CDC) operator, and thereby a data-driven Riemannian metric that can be evaluated at any point with a single forward pass. The conceptual move is to rewrite the marginal CDC — an intractable local covariance under the smoothed data density — as an expectation over noisy samples $Y = X + \sqrt{\varepsilon}\,\xi$ conditioned on $Y$, giving a tractable regression target that has the same gradient as the marginal loss (Theorem 3.1). A low-rank parameterisation $\Gamma_\varepsilon^\theta = M^\top M$ with $M \in \mathbb{R}^{r \times D}$ makes the loss cheap to evaluate at high ambient dimension, and Theorem 4.4 shows that $r \geq 2d - 1$ is enough for a valid global factorisation via a Whitney-immersion + pullback-metric construction. Theorem 4.1 recovers the Coifman–Lafon convergence to the true induced metric as $\varepsilon \to 0$ under local manifold assumptions.

This bridges classical diffusion maps (k-NN kernel constructions of the Laplacian) and modern denoising diffusion models, and the payoff is a fully amortised, dataset-size-independent inference pipeline: the paper reports throughput advantages of up to $400\times$ over a k-NN CDC baseline on an 8-sphere in $\mathbb{R}^{64}$ and demonstrates the method at up to $D \approx 196\,608$ (FFHQ), where score-Jacobian based baselines OOM. Beyond throughput, the learned CDC is used to (i) estimate intrinsic dimension and tangent spaces, (ii) evaluate tangent-perturbation stability of Inception features on MNIST/CIFAR-10/CelebA/FFHQ, and (iii) produce on-manifold interpolations via Riemannian gradient descent. The overall framing — an amortised, principled Riemannian metric that plugs into standard diffusion-geometry pipelines — is genuinely useful.

The main conceptual limitation is that the evidence in the image regime is proxied by a novel Inception-feature-stability metric and by qualitative interpolation figures rather than by direct comparison with k-NN CDC on those very datasets, so the headline claim that k-NN "fails" in high dimension is asserted more than demonstrated. Several notational inconsistencies and reproducibility gaps (below) reinforce that this needs tightening.

Strengths

• Conceptual contribution. Casting the CDC as a conditional expectation under a Gaussian-perturbation kernel converts an $O(N^2)$ operator estimation problem into a per-sample regression problem with the same gradient (Theorem 3.1). This is a clean bridge between diffusion maps and denoising-diffusion training objectives.
• Amortised, dataset-size-independent inference. The neural CDC surrogate has $O(1)$ per-query cost, which is a genuine and demonstrable improvement over graph-based CDC pipelines whose cost scales with $N$ and $k$. Figure 2 (left) makes this vivid.
• Low-rank factorisation with theory. Theorem 4.4 justifies the $r \geq 2d - 1$ low-rank parameterisation via a Whitney-immersion + isometric correction construction, and the low-rank Frobenius loss avoids ever materialising a $D \times D$ matrix — essential for FFHQ-scale $D \approx 196{,}608$.
• Principled convergence result. Theorem 4.1 identifies the population limit as the induced Riemannian metric, and the drift cancellation of the CDC is exploited neatly so that the density normalisation of classical diffusion maps is not needed for the metric.
• Breadth of downstream tasks. The paper shows dimension estimation, tangent-space recovery, and Riemannian-optimisation-based interpolation on both synthetic and image data, illustrating that the learned CDC really does plug into a downstream Riemannian toolkit rather than being a stand-alone estimator.
• Scale to high-dimensional images. Metric matching runs on FFHQ ($D \approx 196{,}608$) with low-rank parameterisation while score-Jacobian OOMs even on CelebA at 80 GB — a meaningful demonstration of the scaling advantage of the parameterisation, not just of the objective.

Weaknesses

• Sign inconsistency in the generator. Section 3 writes $L_\varepsilon := (I + P_\varepsilon)/\varepsilon$, then substitutes $L_\varepsilon = (I - P_\varepsilon)/\varepsilon$ two paragraphs later, and the proof in Appendix A uses the $-$ sign as well. The $+$ version cannot converge to a second-order operator (applied to $1$ it yields $2/\varepsilon$). This is confusing for such a load-bearing object.
• Gaussian perturbation should have zero mean. The Inception-feature-stability metric samples $v \sim \mathcal{N}(x, \Gamma_x)$, but the subsequent use $\Phi(x + v) - \Phi(x)$ and renormalisation $\|v\| = \alpha$ requires $v \sim \mathcal{N}(0, \Gamma_x)$. Because this metric drives every high-dimensional-image comparison (Figure 4), the typo should be fixed and the sampling from a rank-deficient covariance should be described operationally.
• 'Direction only' claim is overstated. Renormalising to a fixed magnitude does not remove direction bias; the sampling distribution is still shaped by the eigenvalue anisotropy of $\Gamma_x$. The current wording promises scale invariance the procedure does not deliver.
• Central 'k-NN fails on images' assertion is not benchmarked. Section 5.2 argues metric matching succeeds where classical k-NN CDC methods fail, but k-NN CDC is never actually run on CIFAR-10/CelebA/FFHQ — the failure claim is proxied by score-Jacobian OOMs (a different baseline) and general theoretical remarks about the curse of dimensionality.
• Interpolation quality is only judged visually. The "no ghost features" claim on MNIST rests on eyeballing Figure 5. For a paper whose headline is on-manifold interpolation, a quantitative on-manifold measure (nearest-neighbour distance in the training set, classifier confidence along the path, likelihood under a pretrained density model) would substantially strengthen the argument.
• Mean-centered conditional loss is asserted, not derived. Theorem 3.1 is proved for the uncentered case only; the mean-centered form actually used for CelebA/FFHQ combines the clean $f(X)$ with a frozen posterior-mean network $(P_\varepsilon f)^\theta(Y)$, which is not the obvious analogue of Eq. 6. The equivalence up to a constant is not established in this setting.
• Numerical consistency between text and Table 1. The stated "$\sim 62\times$ slower" figure for the score-Jacobian on CIFAR-10 does not reconcile with Table 1 (992.3/16.6 $\approx$ 59.8×), and the reader is left unsure whether the divisor is metric matching or Stein score.
• Point estimates without variance. The 400× speedup, the 82×/400× at 2M/8M, the 46% accuracy improvement, and the per-scale eigendecomposition speedups (67×/132×/194×/359×) are all reported as scalars, without seeds, confidence intervals, or protocol details (warmup, precision, batching).
• No non-CDC tangent-space baseline. Tangent-space accuracy is only compared with a k-NN CDC estimator; a local-PCA / LLE-style tangent estimator is a standard, cheap baseline whose absence makes it hard to tell how much of the win is intrinsic to the CDC formulation vs. neural amortisation.
• Section header mislabels a configuration paragraph. The FFHQ configuration paragraph is titled "CelebA configuration", duplicating the earlier CelebA header. Anyone copying hyperparameters risks conflating the two.
• Whitney-immersion presentation and vector-bundle notation. The proof of Theorem 4.4 presents Whitney's immersion theorem as if it gives the low-rank factor, when the isometric correction $A^{-1/2}$ is what does the work. The notation $T\mathbb{R}^K \cong \mathbb{R}^K \times \mathbb{R}^K$ conflates total space with fiber. The 1-dimensional case invokes tangent-bundle triviality without noting that this uses the classification of 1-manifolds and only holds componentwise for the local set $B(x,\delta) \cap \mathrm{supp}(\mu)$.
• Convergence theorem hypotheses are compressed. The reduction to a compact manifold-with-boundary and the smoothness/positivity assumptions on the density are handled implicitly; standard Coifman–Lafon is stated on closed manifolds, and boundary/near-boundary regularity is not made explicit.
• Cost of Riemannian-optimisation-based inference. The interpolation experiment uses $2^{13}$ Euler steps per path, which is not necessarily cheap and is never compared to a k-NN-based geodesic solver at moderate $N$; the amortised-inference selling point does not automatically transfer to this application.
• Hand-picked bandwidth on MNIST. All MNIST qualitative results (Fig. 3a/b and Fig. 5) use a single conditioning bandwidth $\varepsilon = 7$; there is no sweep or robustness study, even though the paper argues elsewhere that the network conditions on $\varepsilon$ across scales.
• No architecture ablation. The 15M-parameter FiLM-conditioned MLP is used across all sphere sizes; whether the FiLM conditioning, the parameter budget, or the rank $r = 16$ choice are each essential is not explored.

Reproducibility & code

Code and data availability were inspected as part of the review; the observations below are based on reading the manuscript specification, not on running any code.

• Missing CIFAR-10 configuration. Appendix E documents synthetic-sphere, MNIST, CelebA, and (mislabelled) FFHQ configurations, but no CIFAR-10 configuration despite CIFAR-10 appearing in Figure 4 and Table 1 — a reproducer has no UNet architecture, rank, noise range, or optimiser to work from for that dataset.
• Missing VAE and latent-space MM specification for Figure 9. The latent-space interpolation experiment cites only "the smaller VAE with 2M parameters and latent dimension $d = 8$" and "the larger VAE with 12M parameters and $d = 32$"; encoder/decoder architecture, ELBO weights, training procedure, latent-space metric-matching architecture, and the pullback-metric definition are absent.
• Inception network Φ underspecified. The high-dimensional evaluation metric depends on an unspecified pretrained Inception network (which weights, which feature layer, which preprocessing). Since Fig. 4 curves depend materially on these choices, the current text does not allow quantitative reproduction.
• Baselines are pointer-cited, not self-contained. Score Jacobian and Stein score are cited to Kharitenko et al. 2025 and Azeglio & Di Bernardo 2025 respectively; the exact score model, score-Jacobian autograd protocol, and rank-1 correction formula are not stated in the paper.
• No seeds or repetition protocol. The 400×, 82×/400×, 46%, and 67×/132×/194×/359× headline numbers are reported without variance bars, and Appendix E does not specify seeds or repetition counts for training runs or timing measurements. Timing methodology (warmup, batching, precision) is likewise not documented.
• k-NN CDC implementation not identified. Appendix E.3 lists the hyperparameter grid for the k-NN CDC baseline but not the underlying library (FAISS? scikit-learn? torch?), GPU/CPU dispatch, or query-time batching — all of which materially affect the 400× throughput ratio at $N = 8\text{M}$.
• Mislabeled configuration paragraph. The header "CelebA configuration." is used for the FFHQ paragraph, which risks a reproducer applying CelebA-labelled hyperparameters to FFHQ or vice versa.
• Metric-matching vs. posterior-mean noise range mismatch on FFHQ. The metric-matching model is trained over $h \in [1.0, 5.0]$, while its accompanying posterior-mean model is trained over $h \in [0.002, 20.0]$; no rationale is given, and the implication for scale-conditioned inference is not stated.
• Interpolation hyperparameters. The stability blend $\delta = 10^{-4}$ for $(1-\delta)\Gamma^\theta + \delta I$ and the $2^{13}$-step schedule are given, but no robustness study is described, so a reader cannot judge which parts of Figure 5 are robust to these choices.

Recommended Changes

Essential

• Fix the $L_\varepsilon$ sign inconsistency. Standardise on $L_\varepsilon = (I - P_\varepsilon)/\varepsilon$ throughout Section 3 to match both the substitution used later in the same section and the Appendix A proof.
• Correct the perturbation sampler. In the Inception-feature-stability metric, change $v \sim \mathcal{N}(x, \Gamma_x)$ to $v \sim \mathcal{N}(0, \Gamma_x)$ and describe how a degenerate Gaussian with rank-deficient covariance is realised in practice via the low-rank factor $M_\varepsilon^\theta$.
• Add a CIFAR-10 configuration to Appendix E. Provide the UNet backbone, rank, noise range, optimiser, and training schedule that produced the CIFAR-10 rows of Figure 4 and Table 1.
• Add a VAE / latent-space MM specification for Figure 9. Document both VAE architectures (2M and 12M parameters), their training procedure, the metric-matching architecture in latent space, and the pullback-metric formula.
• Rename "CelebA configuration" describing FFHQ. Change the header for the FFHQ block to "FFHQ configuration" to remove the duplicated label.
• Reconcile the "$\sim 62\times$" claim with Table 1. Either revise the factor to match the table numbers or state which two entries are being compared.
• Directly benchmark k-NN CDC on at least one of CelebA/FFHQ (or soften the "k-NN fails" language). Even a small-subsample or reduced-resolution k-NN CDC run would substantiate the headline argument of Section 5.2; alternatively, present the claim as a hypothesis and give the theoretical curse-of-dimensionality argument explicitly.
• Specify the Inception network Φ. State the weights, feature layer, input preprocessing, and the $\alpha$ grid used to produce Figure 4.
• Provide self-contained baseline descriptions. For Stein score and score Jacobian, give the explicit formulas and the autograd/vmap protocol used, so results can be replicated without pulling in each cited paper's code.

Suggested

• Report seeds and per-seed variance for headline numbers. Attach error bars (or at least min/median/max over 3 seeds) to the 400×, 82×, 46%, and 67×/132×/194×/359× claims, and document timing methodology (warmup, precision, batching).
• Add a non-CDC tangent-space baseline. Compare Frobenius $\|UU^\top - \hat U \hat U^\top\|_F$ against a local-PCA estimator on the synthetic sphere to isolate the amortisation benefit from the CDC-formulation benefit.
• Add a quantitative on-manifold check for the MNIST interpolation. For example, mean distance from each interpolated $x_t$ to its nearest training image, classifier confidence along the path, or likelihood under a pretrained density model.
• Sensitivity study for MNIST bandwidth $\varepsilon = 7$. Show that eigenvector interpretability and interpolation smoothness are stable across a small $\varepsilon$ range at inference time.
• Clarify the mean-centered conditional loss derivation. Add a short derivation (or an appendix analogue of Theorem 3.1) covering the case where $(P_\varepsilon f)^\theta$ is a separately trained network.
• Tighten the Whitney-immersion presentation. Clarify that the immersion supplies only the smooth bundle map, the isometric correction $A^{-1/2}$ does the metric-preserving work, and that the $d = 1$ case relies on the classification of connected 1-manifolds.
• Fix the vector-bundle notation. Replace $T\mathbb{R}^K \cong \mathbb{R}^K \times \mathbb{R}^K$ with $d\phi_p : T_pM \to \mathbb{R}^K$ everywhere the target is meant to be the fiber, and identify the total space only where it is genuinely needed.
• Ablate architecture choices on the synthetic sphere. A brief study of the FiLM conditioning vs. plain concatenation and of the parameter budget (15M is large for an 8-sphere in $\mathbb{R}^{64}$) would help disentangle the contribution of the objective from that of the specific parameterisation.
• Report interpolation runtime. Include a per-path wall-clock comparison between mm Riemannian optimisation ($2^{13}$ Euler steps) and a k-NN CDC geodesic solver at moderate $N$, so the amortisation story extends to this downstream use.
• Motivate the FFHQ noise range choice. Explain why the metric-matching model uses $h \in [1.0, 5.0]$ while the posterior-mean model uses $h \in [0.002, 20.0]$, and whether the narrow band constrains the range over which the CDC surrogate is valid at inference.
• Add a $\delta$-sweep for the interpolation. Report which qualitative properties of Figure 5 are robust to the Tikhonov mixing weight $\delta$ used to invert the rank-deficient learned metric.