Connecting Neural Models Latent Geometries with Relative Geodesic Representations

Authors: hanlin  yu
berfin  inal
marco  fumero
Venue: NeurIPS_cc_2024_Workshop_UniReps
Published: 2024-09-24 12:47:13

Abstract: Neural models learn representations of high dimensional data that lie on low dimensional manifolds. Multiple factors, including stochasticities in the training process, may induce different representations, even when learning the same task on the same data. However, when there exist a latent structure shared between different representational spaces,
it has been showed that is possible to model a transformation between them. In this work, we show how by leveraging the differential geometrical structure of latent spaces of neural models, it is possible to capture precisely the transformations between distinct latent spaces. We validate experimentally our method on autoencoder models and real pretrained foundational vision models  across  diverse architectures, initializations and tasks.

Main paper: 
 introduction: 
 Recent research reveals that neural models often develop similar internal representations when exposed to similar inputs, a phenomenon observed in both biological [24, 17] and artificial systems [26, 32, 19, 21] . Remarkably, even when models have different architectures, their internal representations can frequently be aligned through a simple transformation, for example linear [27, 25, 30] . This suggests a certain consistency in how NNs encode information, emphasizing the importance of studying these internal representations, and the transformations that relate them. One strategy to do this is to identify representations that are invariant to transformations between distinct models representational spaces. A simple and effective recipe to do this is the one of relative representations [32] , where samples are represented as a function of a fixed set of latent representations. The similarity function employed is cosine similarity, hinting at the fact that representations across distinct models are subject to angle preserving transformations. However, the choice of similarity function should not be limited to capture invariances to one class of transformations. As shown in [8] , other choices can be good as well, and there's not a clear best choice among simple class of transformations for capturing transformation across distinct latent spaces.In this paper we employ geodesic distance in the latent space as a metric for relative representations. This approach ensures that the relative space remains approximately invariant to the isometries of the data's manifold, as characterized by a Riemannian structure. Our contributions can be summarized as follows: (i) We propose a new representation that capture the isometric transformation between data manifolds learned by distinct models. (ii) We propose to employ an approximation of the underlying metric both for classification and reconstruction tasks. (iii) We test relative geodesics on retrieval and stitching tasks on autoencoders and real vision foundation models, across different seeds, architectures and training strategies, outperforming previous methods. 
notation and background: 
 Neural networks (NNs) can be viewed as parametric functions F θ , which are composed of an encoding map and a decoding map, represented as F θ = D θ2 • E θ1 . The encoder E θ1 : X → Z generates a latent representation z = E θ1 (x), where x ∈ X to the input domain X , and the latent space Z. The decoder D θ2 is responsible for performing the task at hand, such as reconstruction or classification. For simplicity, we will omit the parameter dependence (θ) in our notation moving forward. For any single module E (or equivalently D), we will use E X to denote that the module E was trained on the domain X . In the next sections, we will provide the necessary background to introduce our method.Latent Space Communication Given a pair of domains X , Y, a pair of neural models trained on them F 1 X , F 2 Y , and a partial correspondence between the domains Γ : A X → A Y where A X ⊂ X and A Y ⊂ Y, the problem of latent space communication is the one of finding a full correspondence Λ : E 1 (X ) → E 2 (Y) between the two domains, from Γ. In a simplified setting, e.g. two models trained with different initialization or architectures on the same data X = Y and the correspondence is the identity. When X ̸ = Y the problem becomes multimodal.The relative representations framework [32] provides a straightforward approach to represent each sample in the latent space according to its similarity to a set of fixed training samples, denoted as anchors. Representing samples in the latent space as a function of the anchors corresponds to transitioning from an absolute coordinate frame into a relative one defined by the anchors and the similarity function. Given a domain X , an encoding function E X : X → Z, a set of anchors A X ⊂ X , and a similarity or distance function d : Z × Z → R, the relative representation for a sample x ∈ X is:RR(z; A X , d) = ai∈A X d(z, E X (a i )),where z = E X (x), and denotes row-wise concatenation. In [32] , d was set as cosine similarity. This choice induces a representation invariant to angle-preserving transformations. In this work, our focus is to leverage the intrinsic geometry of latent spaces to capture isometric transformations between data manifolds approximations.Latent space geometry For the latent space of a neural network, it is in general hard to reason about its Riemannian structure. However, it is often easier to assign a Riemannian structure to the output space. As such, one can define a pullback metric from the output space to the latent space, which is a standard operation in Riemannian geometry (see e.g. Ch.2.4 of [12] ).Formally, considering the decoder D : Z → X takes as input a latent representations z ∈ Z and outputs x. Given a Riemannian metric defined on x as G X (x). Then, the Riemannian metric at z can be obtained asG Z (z) = ∂x ∂z ⊤ G X (x) ∂x ∂z = J z (D) T J z (D),where J z (D) is the Jacobian of D evaluated at z. The metric tensor G X is useful to compute quantities such lengths, angles and areas on M. Given a smooth curve γ : [a, b] → M one can compute the energy E of γ as follows [34] EQUATIONwhere v(t) = γ(t). This can be approximated using finite difference approaches. Geodesics are minimizers of this energy [34] .When considering a differential geometry perspective, the problem of latent space communication can be interpreted as finding a transformation between the data manifolds M 1 , M 2 approximated by two neural models F 1 , F 2 . The relative representation framework can capture this transformation implicitly if equipped with the right metric. A natural candidate for this metric is the geodesic distance defined on M 1 , M 2 , respectively. This choice make the relative representations invariant to isometric transformation of the manifolds M 1 , M 2 . However, for high dimensional problems, the high cost of computing the geodesic renders the above methods inappropriate [34, 9] . Furthermore, one can argue against directly using the latent geometry induced by deterministic models from a theoretical perspective [16] , as it may result in undesirable properties, e.g. the resulting geodesics going outside of the data manifold.We therefore consider using the approximate curve energy / distance of the straight line (in the Euclidean sense) connecting the representations in the latent space:RR geo (z; A X ) = ai∈A X E(γ(z, E X (a i )))where γ(z 1 , bz 2 ) = (1 -α)z 1 + αz 2 is the convex combination between the points z 1 , z 2 . Further descriptions on our method for obtaining the geometric representations can be found in Section A.1. 
experiments: 
 In the following we will evaluate relative geodesic representations on the latent communication problem across models trained with different initializations, different architectures, and tasks. right) datasets respectively. In green, we plot the performance of [32] , in blue, our method. Shaded area indicates standard deviation across different random set of anchors. Relative geodesics consistently outperform the cosine baseline, obtaining peak performance.Experimental setting For the following experiment we trained pairs of convolutional autoencoders (F 1 , F 2 ) with different initializations on the MNIST [11] , FashionMNIST [37] , CIFAR10 [23] datasets. The architecture of the convolutional autoencoder is detailed in the Appendix. After training we extracted 10k samples from the test set, and map them to the latent spaces of the two models, to representationsZ 1 = E 1 (X), Z 2 = E 2 (X)respectively. Starting from a small set of anchors in correspondence A X → A Y , the objective is to evaluate how well from the relative representations is possible to recover the full correspondence between the representations Z 1 , Z 2 . As baseline we compare with relative representations using cosine similarity [32] .In Figure 1 we plot the performance in terms of MRR on MNIST,FashionMNIST, CIFAR10 datasets. To obtain the score we first compute similarity matrices between relative representations of the two spaces as D(Z 1 , Z 2 ) where D i,j = RR(Z1) T i RR(Z2)j ∥RR(Z1)i∥2∥RR(Z2)j ∥2 . Then we compute the Mean Reciprocal Rank (MRR, see Appendix A.2.1) on top of the similarity matrix. In the figure we plot MRR as a function of a random set of anchors, where the shaded areas indicate the standard deviations over 5 different set of random anchors with the same cardinality. Our method consistently performs better than Relative Representation, saturating the score with few anchors on all the domains, despite the different degree of complexity of the latent spaces. In addition, our method show way less variance in the result, being more robust to the choice of the anchor set.Takeaway: Relative geodesic representation capture almost perfectly the transformations between representational spaces of models initialized differently, outperforming [32] in terms of number of anchors needed and robustness.Experimental setting For this experiment we consider the same pairs of autoencoders trained on the MNIST, FashionMNIST, CIFAR10 datasets of section 3.1. Starting from a set of five random anchors we want to [27, 25, 32] ; the rightmost column is our method. Relative geodesics yield the best stitching results using just 5 anchors. estimate a transformation T between the model representational spaces Z 1 , Z 2 . Differently from [32] , in which zero shot stitching was achieved by training once a decoder module with relative representations and then exchanging different encoder modules, here we achieve stitching without training any decoder. We compute relative representation with respect to the set of anchors, and then compute a similarity matrix D(Z 1 , Z 2 ). Then we compute the vector c = arg max i (D) representing a correspondence between the two representations matrices Z 1 , Z 2 , and use c to fit a linear transformation T to approximate the transformation between the two domains. We perform stitching by performing the following operation for a sample x ∈ X :x = D 2 • T • E 1 (x).We visualize the results of reconstructions of random samples in Figure 2 , comparing with [32, 25, 27] . For each dataset, each column represents respectively: (i) the original autoencoding mapping for a sample x of modelF 1 , D 1 (E 1 (x)) (ii) D 2 (E 2 (x)) (iii) the mapping D 2 (E 1 (x)) (iv) the mapping D 2 (T anchors E 1 (x))where T anchors is estimated on the five available anchors, (v) the mapping D 2 (T cosine E 1 (x)) where T cosine is estimated among all 10k samples with the correspondence c obtaining in the relative space of [32] (vi) Our result D 2 (T relgeo E 1 (x)) where T relgeo is estimated from the correspondence obtained in the relative geodesic space. While the baselines do not reach a good enough reconstruction quality, reconstructions with our method are almost perfect in accordance with the results in Figure 1 .Takeaway: The relative geodesic space enables to stitch together neural modules trained on different seeds. For this experiment we follow the stitching procedure of [32] , section 5. We consider ResNet-50 [18] , Vision Transformers (ViT) [13] , with both patch 16-224 and patch 32-384, and DinoV2 [33] . These models differ in architecture and pretraining tasks (classification, self supervised contrastive learning). We mainly compare cosine relative representation (cos) [32] , relative geodesic representation of the curve length assuming Euclidean geometry on the logits with 20 discretization steps of Equation 1 (geo1) and directly using the distance of the corresponding logits (geo2) as relative representations, corresponding to 1 discretization step.The accuracies are shown in Figure 3 . We plot confusion matrices of accuracies indicating that the performance of stitching the backbone of model on each with the classification head of each column. The accuracies are shown in Figure 3 , while the MRR with respect to cosine similarity are shown in Figure 4 .Takeaway: Using geometric relative representations yields better accuracies, avoiding downgrading of performance when performing model stitching. 
conclusions: 
 In this work we explored the framework of relative representation equipped with geodesic energy to capture the trasformations occuring between neural manifold learn by distinct neural architecture. We demonstrated superior performance. As limitation we observe that the evaluation results depend on the number of discretization steps when evaluating the representations. Future steps include exploring the multimodal case, when X ̸ = Y, different formulation of the energy, and considering different architectures e.g. VAEs as in [34, 3, 2] .1. For each query, rank the list of retrieved items based on their relevance to the query.2. Determine the rank position of the first relevant item in the list. If the first relevant item for query i is found at rank position r i , then the reciprocal rank for that query is 1 ri . 3. Calculate the mean of the reciprocal ranks over all queries. If there are Q queries, the MRR is given by:EQUATIONHere, r i is the rank position of the first relevant item for the i-th query. If a query has no relevant items in the retrieved list, its reciprocal rank is considered to be zero.MRR provides a single metric that reflects the average performance of the retrieval system, with higher MRR values indicating better performance. 
 
 Appendix: We provide here the architectural details of the convolutional Autoencoders employed in experiments in Figures 1 and 2 Table 1 For the classifier experiment, in order to obtain geometric representations we need a decoder. The architecture is shown in Table 2 .Classification head input_dim LayerNorm input_dim × 500 Linear-Tanh 500 × 10 Linear Table 2 For evaluating the performances of the representations, we train a classifier with the same architecture as used by [32] .Representation alignment: There is a growing evidence that neural networks trained under different settings still tend to generate similar internal representations [7, 22, 20, 26, 6, 28, 19] , which is shown to be more evident in wide and large networks [5, 31, 35] . These aligned representations make it possible to stitch models together [14, 4, 10] , allowing the swapping of components between different networks.Latent geometries [34, 36] considered the latent space of autoencoders, proposing to use a pullback metric, assuming the output space is Euclidean. For classifiers one can obtain a Riemannian metric primarily using two approaches [15] , either by pulling back the Fisher Information Matrix [1, 2] or by assuming an Euclidean geometry on logit space and pulling back the metric.Preprint. Under review. 


Decision: In light of the positive reviewers' feedback and relevancy of the submission, we are pleased to accept this paper for presentation at UniReps 2024. We kindly ask the authors to incorporate the reviewers' suggestions and feedback in the final camera-ready version of the manuscript.

Reviews: review: title: title: incremental improvement to relative representations alignment with impressive results
main_review: main_review: Authors propose a new geodesic representation that leverages the differential geometrical structure of latent spaces to model the transformations between them, demonstrating improved performance compared to previous methods. This approach is shown to be effective with experiments with aligning and stitching models, and the results are impressive. I did not completely follow the derivation of the the curve energy used to form the relative representations and would like to have seen some discussion of the runtime complexity of the implementation used. But I am familiar with relative representation research and found the background well covered and experiments convincing.

Rating provides is relative to the abstracts track as this appears to be an excellent fit - providing a short and insightful finding very relevant to the workshop audience.
title: title: Conditions under which the results hold unclear
main_review: main_review: ### **Paper summary**

The paper proposes comparing different models’ latent representations (relative representations, RR) according to Riemannian geodesic distances, presumably a more natural metric compared to cosine similarity that has mostly been prevalent so far. Due to intractability, the paper resorts to using the Euclidean distance. The paper empirically finds superior RR, both in terms of a MRR metric and autoencoder reconstructions based on stitching RR across models.

### **Strengths**

It is important to point out possible limitations in using cosine similarity as a metric, and proposing Riemannian geometry as a natural alternative, even if it ultimately is not used in the paper. The empirical findings on improvements are surprising and potentially valuable if persistent.

### **Weaknesses, questions and suggestions**

- Why is MRR used? What are alternatives that have been rejected?  
- Cosine similarity a priori is a natural metric for neural network’s latent spaces, as latent representations tend to be linear, because they will be processed by a linear readout at the end of the network. However, it discards information about the norm of the vector. Would the authors expect a dot-product metric $z\_1 \\cdot z\_2$ that incorporates both angle and norm to be superior to Euclidean distance?  
- Would the same results have been achieved if one just had proposed using Euclidean distance $|z\_1 \- z\_2|$ instead of $\text{cosine  similarity}(z\_1, z\_2)$? I.e., is the Riemannian framework strictly needed for the paper’s empirical results?  
- It is unclear to me in how far the results on superiority are contingent on data preprocessing. E.g., is it possible to state whether superiority persist if data is centered and normalized? Making code available would be helpful to this end.   
- Limitations are not discussed.

### **Recommendation**

The core weakness is to fully comprehend the processing conducted in the paper, which is why I right now recommend rejection.
title: title: Using Geodesics for Relative Representations
main_review: main_review: **Summary:** 

The work uses the pull-back metric from the output space as a distance metric for relative representations. 
They show that this better preserves accuracy when stitching together CIFAR-10 classifiers compared to using cosine similarity, 
and also that this allows for stiching of latent spaces with a linear transformation. 


**Questions and Suggestions:**

Q1: If you have decided to call your kind of relative representations "relative geodesics", then I suggest introducing the name in (i) in the introduction, e.g. 
"(i) We propose a new representation that capture the isometric transformation between data manifolds learned by distinct models, called relative geodesics."
This is to avoid confusion when you use the term later.

Q2: On line 77 and 78 it says: "We therefore consider using the approximate curve energy / distance of the straight line (in the Euclidean sense)".
I assume you mean "length" of the straight line. I also assume that what you mean is that you approximate curve lengths using straight lines on the output space, 
then pull it back to the latent space and use this as the distance measure. However, I find this is very unclear in the text. Especially since you often write 
energy/length in the citation above and in all of A.1. Writing it in this way causes confusion, since the energy and length of the curve are two different 
things. So when you write you use the energy/length, it sounds like you use both. Please clarify. 

Q3: There is a typo in the equation on line 79: \gamma(z_1, bz_2) should be \gamma(z_1, z_2).

Q4: The procedure in experiment in 3.1 is unclear. It says: "Starting from a small set of anchors in correspondence A_X -> A_Y, the objective is to evaluate how well from the relative 
representations is possible to recover the full correspondence between the representations". What do you mean by "recover"? 
Since you are using the Mean Reciprocal Rank (MRR), you must be using something as queries. Since you write (94-95) "Then we compute the Mean Reciprocal Rank
(MRR, see Appendix A.2.1) on top of the similarity matrix." It sounds like the similarity matrix is the query? What do you retrieve? 
Please also clarify what you use as relevance measure to rank the items and how you decide which is the "first relevant item in the list". 

Q5: Please write the full name the first time you introduce a new concept. I suggest writing the full name of Mean Reciprocal Rank in both the 
caption of figure 1 and in line 92. You can then use the abbriviation in line 94. 

Q6: In line 97, it says: "Our method consistently performs better than Relative Representation...", you mean "Using curve energy consistently 
performs better than using cosine similarity". They are both relative representations. 

Q7: In A.4 MRR seems to be quite low for both curve energy and cosine similarity. Do you have any idea why? 
Especially why it is so different from the experiment in 3.1? Please also add details about how you calculated the MRR in this case.  


**Strengths:**

S1: I like the idea of using geodesics for the distance measure 

S2: Using geodesics for the relative representations, they show that the representations can then be connected using a linear transformation.

S3: They average over multiple choices of anchors in 3.1 to account for the variability incurred by random choice of anchors. 

S4: The structure in section 3 with "Experimental setting", "Analysis of results" and "Takeaway" is very nice. 


**Weaknesses:** 

W1: The discription of how they construct the distance measure is unclear, see Q2. 

W2: The description of their experiment 3.1 is unclear, see Q4. 


**Justification:**

I recommend this paper to be accepted to the workshop with the added clarifications.

The idea is good even though some integral parts need to be rewritten.
title: title: Promising novel development in relative representations
main_review: main_review: The paper is well written, easy to follow and understand. The introduction into the framework and the mathematical background is very clear and well motivated, the use of geodesics for relative representations seems convincing. While relative representations and all experiments presented are not novel, the use of geodesics provides a novel approach to relative representations which improves those quite significantly. The experiments are relevant and well-motivated, however the discussion of results is lacking and should be extended/included. The take-away message after each result is a very nice touch!

**Questions:**

3.1
- How are anchors chosen? 
- Are the different models just initialized with a different random seed? Do all models achieve the same performance? 
- Why do you use MRR, which kind of information does this measure and what does it show wrt relative representations? 
- Can you also include a linear baseline?

3.2
- Do you only use 5 anchors? How do the results look with more anchors (RR based on cos seem to need more anchors based on Fig.1)?

3.3
- Is geo1 or geo2 better, you mention geo1 in the figure cpation but based on the figure geo2 (center) seems to perform slightly better?
- What is geo2 exactly? Is this also based on geodesics?
- What is the take-away from the MRR results?

**Suggestions for improvement:**
- Increase font size on all figures and include titles for Fig.3 to increase readability
- Use different color scale for Fig.3 to make it easier to spot differences between methods
- Describe the actual results in “analysis of results” and not the what you plotted, description of the plot should be moved to figure captions and the analysis should not be referred to the reader
- L.140 spelling: trasformations -> transformations
- Appendix: include captions for tables and expand on A.4 

This work is of quite high significance and interest for the community and I recommend acceptance.