Connecting Neural Models Latent Geometries with Relative Geodesic Representations

TL;DR

This paper introduces a method to analyze and compare the intrinsic geometric structures of neural network latent spaces, enabling precise understanding of transformations between models trained on similar data.

Contribution

It leverages differential geometry and pullback metrics to capture latent space structures, facilitating model comparison and transformation analysis.

Findings

Effective in model stitching and retrieval tasks

Applicable across various architectures and datasets

Scales efficiently to large models

Abstract

Neural models learn representations of high-dimensional data on low-dimensional manifolds. Multiple factors, including stochasticities in the training process, model architectures, and additional inductive biases, may induce different representations, even when learning the same task on the same data. However, it has recently been shown that when a latent structure is shared between distinct latent spaces, relative distances between representations can be preserved, up to distortions. Building on this idea, we demonstrate that exploiting the differential-geometric structure of latent spaces of neural models, it is possible to capture precisely the transformations between representational spaces trained on similar data distributions. Specifically, we assume that distinct neural models parametrize approximately the same underlying manifold, and introduce a representation based on the pullback metric that captures the intrinsic structure of the latent space, while scaling efficiently to large models. We validate experimentally our method on model stitching and retrieval tasks, covering autoencoders and vision foundation discriminative models, across diverse architectures, datasets, and pretraining schemes.