Connecting Neural Models Latent Geometries with Relative Geodesic Representations

We thank the reviewers for their feedback and suggestions. Below we address the Reviewer concerns. We remain available for any further clarification during the discussion period.

Anchors sampling strategies: We report the accuracies and symmetricized MRR based on Cosine on CUB dataset, comparing uniform sampling for anchors (as used in the main paper) with using farthest point sampling (fps) and kmeans, as described in Moschella et al 2023. We observe that the general conclusion persists. Below we report the stitching accuracies; the other results based on retrieval behaved similarly, and have not been included for space constraints: we included the full results in the manuscript.

Accuracies

Rel(Cosine)

uniform

fps

kmeans

RelGeo(Pullback)

uniform

fps

kmeans

RelGeo(Diet)

uniform

fps

kmeans

“Sensitive .. to N”: we provided an ablation study of the effect of N in Appendix A.2.6. We observe that the proposed method is not too sensitive to N.

“Multimodal”: we report the results of Rel(Cosine) and RelGeo(Diet) on aligning vision models and the text encoder of CLIP on Flickr30k. We use 2000 Diet data points, and simply use the first text string of each data point for the text part. We report the retrieval results on 10000 test data points. As shown in the result tables, RelGeo(Diet) outperforms Rel(Cosine) also in the multimodal setting. We will incorporate the full results of the experiment in the manuscript.

Symmetricized MRR Cosine

Cos

RelGeo(Diet)

“Proposition 3.1”: Proposition 3.1 states that each geodesic distance as used by relative geodesic representations is invariant under re-parameterisations, which holds regardless of the number of anchors.

“Non-isometric transformation”: the Riemannian geometry framework only provides invariance under isometries. As such, it is unclear what happens under non-isometric transformations. We remark that, suppose we view each neural network as a smooth function, up to some assumptions, invariance under isometries can provide certain identifiable guarantees; see e.g. [a].

[a] Identifying metric structures of deep latent variable models, Syrota et al., ICML 2025

We thank the reviewer for their feedback and suggestions. We address their questions in the following, and remain available for any further clarification in the discussion period.

“Why it works”: Prior work on representational alignment has shown that representations from different models can often be approximately aligned using simple transformations (e.g., linear, orthogonal, or locally linear maps). Even when models are trained independently—with different architectures, modalities, or datasets that nonetheless share an underlying structure—they tend to learn similar representations, suggesting convergence toward a shared encoding of entities (Huh et al., 2024). From a theoretical standpoint, identifiability results [a] imply that if two discriminative models learn the same likelihood function, their internal representations must be equivalent up to a linear transformation. However, this ideal scenario rarely holds exactly in practice: training dynamics, nuisance factors, and unmodeled variability can all introduce distortions. In our case, it may be too strong to assume that two models learn different parametrizations of an identical manifold. Instead, we adopt a weaker assumption—that they do so up to some bounded distortion. Recent theoretical work has begun to explore relaxations of strict identifiability to account for such bounded distortions [b]. Integrating these relaxations into our Riemannian framework presents a promising direction for future work.

[a] Roeder, Geoffrey, Luke Metz, and Durk Kingma. "On linear identifiability of learned representations." International Conference on Machine Learning. PMLR, 2021.

[b] Nielsen, Beatrix MG, et al. "When Does Closeness in Distribution Imply Representational Similarity? An Identifiability Perspective." arXiv preprint arXiv:2506.03784 (2025).

How fast is the computation of formulas (4) and (5): the computation of the relative geodesic representation for a single sample requires $N \times | \mathcal{A}_X |$ forward passes of the decoder, where $N$ is the number of steps, and $| \mathcal{A}_X |$ is the number of anchors. Notice that these are vectorizable and can be run all in parallel on GPU, allowing for fast evaluation. The number of steps we employ in the experiments is 15. In Figure 15, 16 in the Appendix, we provide an ablation on N on stitching experiments. The computations can be accelerated through utilizing the pullback geometry. Specifically, each outer product involving the velocity with the Riemannian metric over a small timeframe can be approximated with the corresponding quantity viewed in the output space. For Pullback, this implies the Euclidean distances between the function outputs; for Diet, this implies the spherical distances between the function outputs. Furthermore, the computations along the straight line approximation can be vectorized. As such, they allow reasonably fast evaluations.

Line 199-200: we agree that the text can be improved. We stated that it was “meaningful”, as the resulting method was demonstrated to achieve good performances in the experiments on vision foundation models as presented in Section 5, where it consistently yielded the most identifiable relative representations as measured by mean reciprocal ranks.

We thank the Reviewer for their answer and their engagement. We report below additional answers and results to the points readied by the Reviewer. We hope to have addressed their requests and remain available for further questions.

Follow up on "why it works”

We agree with the Reviewer that a clearer connection between theoretical results and practical methods for representation alignment is needed. Our intention was to emphasize that the few theoretical results available (e.g., [a]) often rely on unrealistic assumptions, such as proofs in axiomatic settings, infinite-data regimes, or the requirement that two models learn exactly the same likelihood function. In our view, a meaningful first step toward bridging theory and practice is to relax these assumptions, as initiated in [b], and begin to model more realistic scenarios, e.g. including the dynamics introduced by model training.

We believe that Riemannian methods can play a key role in this direction to model to capture local alignment beyond linear global transformations of the space (as considered in [a]) and possibly accounting for distortions measured in the linear case in practice. We are also very interested in the results on overparameterized networks mentioned by the Reviewer; if a reference can be provided, we would be happy to add a discussion including this aspect in the manuscript.

[a] Roeder, Geoffrey, Luke Metz, and Durk Kingma. On linear identifiability of learned representations. ICML, 2021.

[b] Nielsen, Beatrix MG, et al. When Does Closeness in Distribution Imply Representational Similarity? An Identifiability Perspective. arXiv:2506.03784, 2025.

---

Compute Time analysis

We report the running times for two experiments:

1. Evaluation on FashionMNIST (Figure 2) using an autoencoder
2. Stitching experiments on CUB (Figure 14) using Diet decoder and a classification head.

In both cases, we vary the number of anchors used for the relative geodesic representations.

The computational complexity of evaluating relative geodesic representations scales with the number of anchors, requiring $N \times |\mathcal{A}_X|$ forward passes of the decoder per data point.

---

FashionMNIST (Figure 2)

For FashionMNIST, the decoder is a convolutional autoencoder (architecture in Appendix Table 6). Evaluation is performed on 10k test samples for both networks, using an NVIDIA 3080 Ti GPU. The table below reports the number of anchors, evaluation time in seconds, and the corresponding MRR score:

---

CUB Stitching (Figure 12)

For the CUB stitching experiments, the decoders are:

• RelGeo(Pullback): classification head ( architecture in Appendix Table 4)
• RelGeo(Diet): lightweight “diet” decoder (architecture in Appendix Table 5)

Experiments are performed on the CUB dataset (train/val/test split totaling 11178 samples) and averaged across all 16 model pairs. Note evaluation is faster here due to the shallower decoders compared to the FashionMNIST autoencoder.

RelGeo(Pullback)

RelGeo(Diet)

Takeaway: The evaluation of relative geodesics (Equations 4 and 5) is computationally efficient, as confirmed by our measurements. In the stitching experiments, the overall runtime is dominated by training the decoders (both the diet decoder and the classifier head). However, this step is also reasonably fast, since we train on only 2,000 samples with augmentations. For reference, the total runtime per experiment (including all possible models combinations) is on the order of a few hours. We will include computation times in the manuscript.

We thank the reviewer for their thoughtful comments. We address their concerns below and remain available for further questions during the discussion period.

• Clarifications about linearity: We respectfully remark that the approximation using the straight line does not correspond to a linear metric.

• “Figure 2 seems to imply that…”: Figure 2 and the companion text starting in Line 161 serve as a verification that the proposed straight-line approximation is a meaningful approximation of the ground truth geodesic distance, preserving the rank. We remark that the straight-line approximation corresponds to the Riemannian curve length (or the energy) of the line (“straight” in the Euclidean sense) under the Riemannian metrics. In other words, both the approximation and the true geodesic use the same Riemannian metric, with the difference lying in the employed curve. The geodesic energy is an upper bound to the true geodesic distance. Additionally we point out that the latter is not scalable to large models, while the approximation nicely scales.

• Comparison with linear methods: We repeat the stitching experiment of Figure 3, adding two additional baselines:

  • Linear: fitting a linear map between the two spaces using the anchors
  • Ortho: fitting an orthogonal map between the two spaces using the anchors, similarly to (Maiorca et al. [2024])

  and report the results in the tables below. We also note that additional qualitative evidence of comparison with linear methods is in Figure 4, fourth column.

Takeaways: Our method outperforms both the new baselines on all datasets (MNIST, FMNIST, CIFAR), while being training free. Additionally, we observe that, despite its simplicity, the dataset cannot be aligned with only a few anchors available, as confirmed by our results.

MNIST

---

FMNIST

---

CIFAR

• Experiment on larger datasets: We provided stitching experiments on vision foundation models in Section 5, which involved larger datasets, demonstrating the performance of our method and of the geodesic energy approximation scheme. We remark that computing true geodesics with models of this size is largely infeasible, due to computational and memory constraints, as well as optimization difficulties since the Jacobian might be ill defined at some points.

• Benefits to a training-free approach for representation alignment: For stitching application, having a non parametric method is of fundamental importance for different reasons:

  • Canonical mapping: Our method can be interpreted as mapping different spaces to a canonical one, invariant to the Riemannian metric adopted (see Figure 1 in the manuscript). This preferred to fitting maps between spaces as the number of adapters to train increases combinatorially with the number of spaces considered. This was observed as well in (Moschella 2023).
  • Sample efficiency: Our method is particularly sample efficient, requiring few anchors (aligned data) in order to work. As opposed, approaches which involve training (e.g. fitting linear maps or neural networks) require much more aligned data to learn a good map (see for example, the experiment above). This is of fundamental importance for example in multimodal applications, where paired data is costly.

We thank the Reviewer for their detailed comments and feedback. Please find our responses below. We remain available for any additional questions during the discussion period.

• Typos line 122 and 147: Thanks for noticing these typos. The convex combination should be: \tilde{\gamma}(z\_1,z\_2; \alpha)= (1- \alpha) z\_1 + \alpha z\_2\. We fixed the typos in the manuscript.

• Evidence of fidelity of the straight line approximation: In Figure 2 and the companion text starting in Line 161, we empirically verify that the straight line approximation is generally reasonable by reporting pairwise matrices and showing that the rank is preserved on random pairs of the CIFAR and MNIST dataset. Theoretically this provides an upper bound to the geodesic distance.

• Pullback and Diet: We propose the general methodology of using Riemannian distances to enhance relative representations, pulling back a metric $G_{\mathcal{X}}$ from the output space to the latent space. For classifier networks we employ two different choices for $G_{\mathcal{X}}$ and its ambient space: pulling back the Euclidean metric from the logits space (denoted as pullback) and pulling back the spherical metric from the layer before in Diet decoders, assuming a vMF distribution in the space (denoted as Diet). We will clarify the text and improve the flow between Section 3.1 and Section 3.2.

• “Section 3.2 is hard to read”: We skipped the technical details to keep the main text concise. We provide a concise summary below, and will provide a more detailed description in the paper. Further technical details can be found in [1]:

  • Suppose there is a finite set of vectors $v_{c}$ on a unit sphere, representing each class, such that each instance belongs to exactly one class. The latent variables $z$ are drawn from a vMF distribution with concentration parameter $\kappa$ from the corresponding $v_{c}$, and the data $x$ is generated using a function $g$ as $x=g(z)$. Then, with a Diet decoder, the representations before the final linear head identify the corresponding ground truth $z$ up to a transformation.
  • The recovery of latent variable $z$ can be rigorously proven by expanding on the theoretical framework of non-linear ICA. The theory holds under reasonable assumptions on the data generation process and the embeddings and weights in Diet decoder’s final layer. For a formal treatment of the assumptions, we refer the reviewer to Assumptions 1C and Theorem 1C in [1].
  • Theory suggests that a trained Diet decoder is capable of identifying the ground truth $z$ up to a transformation. As such, we directly assign a spherical geometry to $z$, and compute the distances accordingly.

• Reconstruction experiment (stitching): In the experiment in Figure 4 we consider a setting where the number of anchors is limited, i.e., 5 for MNIST and FMNIST datasets and 10 for CIFAR, which leads to the difference: in the corresponding Figure 4 in (Moschella et al. 2023), the number of anchors was set to 500 (cfr. Appendix A.5.5 in Moschella et al. 2023). This highlights the huge increase in sample efficiency of our approach, as confirmed by our results in Figure 3.

• Performance on Imagenet (stitching): We currently do not have a definitive explanation of this, and it remains an interesting avenue for future work.

[1] Cross-Entropy Is All You Need To Invert the Data Generating Process, Reizinger et al., ICLR 2025

We thank the Reviewer for their response. We address below their additional questions and remain available for any further clarification.

Diet vs. Pullback

We would like to gently clarify that both “Pullback” and “Diet” are proper pullback metrics:

"Pullback" uses the Euclidean metric from the classifier’s logit space (or from the autoencoder output space), as the Reviewer correctly noted.

Diet uses the spherical (cosine) metric pulled back from the penultimate layer of the Diet decoder.

Both metrics compute the Riemannian arc length (or energy) by numerically integrating along linear interpolations in the latent space. In practice, Diet also relies on a discretization along the interpolation path.

In our experiments, we typically use 15–30 dicretization steps, so the method effectively integrates the pullback metric along the path.

In summary, the two metrics differ in:

1. The metric being pulled back (Euclidean vs. cosine).

2. The output space considered (logit/output space vs. Diet penultimate layer).

These choices impact downstream tasks: for instance, Diet is more suited for retrieval tasks, while Pullback performs better for stitching, as the Reviewer observes as well. It is in general not clear if a single metric can be optimal for all tasks.

Finally, we note that different pullback metrics can be combined, either by constructing product metrics (via concatenation) or by interpolating between metrics, which we see as a promising future direction to get good perfomance on multiple diverse tasks with a unique representation.

We thank the reviewer for their answer and engagement. We would like to gently remark that stitching and retrieval are to some extent orthogonal tasks, as the former depends on identifiability up to class and the later up to instance, as such it makes sense that there is not a single metric outperforming on every downstream tasks. We also gently remark that:

• Diet still performs well and remains competitive on stitching, though may not be excelling,
• For practitioners it is generally known beforehand whether the task concerns stitching or retrieval. As such, choosing a right metric for the downstream task generally does not affect the generality of the method.

We remain available for any further clarifications.