Comparing Foundation Models using Data Kernels

Before we can compare the embedding spaces of multiple foundation models, we require a model of the embedding space of a single foundation model. Consider a foundation model $f:\mathcal{X}\to\mathcal{Y}$, where $\mathcal{X}$ is the model’s input space (e.g., the space of text, images, etc.), and $\mathcal{Y}=\mathbb{R}^{d}$ is the model’s embedding space (e.g., pooled word embeddings, image embeddings, etc.). For $N$ input examples $x_{1},x_{2},...,x_{N}\in\mathcal{X}$, we compute the embeddings of these examples $f(x_{1})=y_{1},f(x_{2})=y_{2},...,f(x_{N})=y_{N}\in\mathbb{R}^{d}$. Let $X$ be the set $\{x_{1},x_{2},...,x_{N}\}$ and $Y$ be the matrix whose rows are $y_{1},y_{2},...,y_{N}$. We define the data kernel $A$ of $f$ on $X$ as

where $\text{TOP}_{k}:\mathbb{R}^{N,N}\to\{0,1\}^{N,N}$ is a function applied to every cell in its input matrix that returns $1$ if that cell is in the top k values of its row (and not on the diagonal) and $0$ otherwise. In other words, $\text{TOP}_{k}$ returns the hollow adjacency matrix of the $k$-nearest neighbor graph of its argument.

We interpret the data kernel as the adjacency matrix of a graph where each node is an element in our input data, and an edge exists between nodes $i$ and $j$ if $j$ is one of $i$’s $k$ nearest neighbors. Critically, the notion of distance used to define proximity when constructing this graph is induced by the foundation model $f$.

Modeling data geometry with a neighborhood graph is an established technique with known desirable properties in the dimensionality reduction literature. For example, it is known that that the shortest paths on the weighted $k$ neighbor graph approximate the geodesic distances between datapoints under the metric induced by the model. [22] The popular UMAP algoroithm [16] utilize the weighted $k$ neighbor graph to compute visualizations of high dimensional manifolds. The unweighted $k$ neighbor graph, as described in in Eq. (1), retains many desirable limiting statistical properties, such as consistent estimation of the true-but-unknown underlying manifold density and metric [15, 10]. These properties motivate our use of the data kernel as a measurement of the geometry induced on a set of data by a particular model.

The Random Dot Product Graph (RDPG) [27] is a flexible model for undirected, hollow, symmetric random graphs that has been applied successfully in a variety of domains [12, 7, 25]. RDPGs posit that each node in a graph has an associated latent position $z$, and that the probability of an edge existing between nodes $i$ and $j$ in any realization of the RDPG is exactly the dot product of their corresponding latent positions $\langle z_{i},z_{j}\rangle$. We write $A\sim RDPG(ZZ^{\top})$ to refer to a graph $A$ sampled from an RDPG with latent position matrix $Z$.

For analysis related to a single data kernel, methods developed under the RDPG assumption can rely on the consistent estimation of the latent positions via the adjacency spectral embedding [1]. The consistency result holds up to the non-identifiability of an orthogonal transformation of the true-but-unknown latent positions, since $\langle z,z^{\prime}\rangle=\langle Rz,Rz^{\prime}\rangle$ for any orthogonal matrix $R$.

In the context of modeling and comparing the embedding spaces of foundation models using the RDPGs, the orthogonal non-identifiability of the parameters of the RDPG is quite natural. For example, foundation models which employ cosine similarity to measure embedding proximity such as CLIP [19] and SBERT [20] will give rise to identical data kernels under orthogonal transformation of their embedding spaces. Vertex-level inference methods designed for observed graphs assumed to be an RDPG will thus be appropriate for datum-level analysis in the embedding spaces of the foundation models.

Multi-graph analysis under the RDPG avoids the estimation of the unknown orthogonal matrix $R$ by utilizing joint embedding techniques such as the omnibus embedding [18]. In particular, let $A^{(m)}$ and $A^{(m^{\prime})}$ be the symmetrized data kernels of a fixed input dataset $X\in\mathcal{X}^{N}$ embedded by foundation models $f_{m}$ and $f_{m^{\prime}}$, respectively, and let $O_{m,m^{\prime}}$ be the omnibus matrix defined as follows:

We define $\hat{Z}^{(m)}_{m,m^{\prime}}$ as the first $N$ rows of the adjacency spectral embedding of $O_{m,m^{\prime}}$ and $\hat{Z}^{(m^{\prime})}_{m,m^{\prime}}$ as the remaining $N$ rows. When $A^{(m)}$ and $A^{(m^{\prime})}$ are instances of the same RDPG model, both $\hat{Z}^{(m)}_{m,m^{\prime}}$ and $\hat{Z}^{(m^{\prime})}_{m,m^{\prime}}$ are consistent estimates for the true-but-unknown underlying latent positions $Z$ [13]. As a result, if two data kernels are drawn from the same underlying RDPG, the estimates of $z_{i}$ after the omnibus embedding are such that:

This convergence property is the key to our ability to compare model representation spaces, which we develop as a hypothesis test in the next section. Intuitively, we can think of the omnibus embedding as a theoretically justified way to align different foundation models’ embedding spaces for subsequent comparison.

Under the null hypothesis that the two data kernels are realizations of the same underlying RDPG, the limit theory outlined in [13] and briefly introduced in Eq. (2) suggest a statistical hypothesis testing procedure for the embeddings induced by foundation models that is sensitive to changes in representations on a per-datum basis. In particular, if the two data kernels represent document $i$ similarly then the distance $d_{i}=\|\hat{z}_{i}^{(m)}-\hat{z}_{i}^{(m^{\prime})}\|_{2}$ should be “small” for large enough $N$. Conversely, if $d_{i}$ is “too large” we may reconsider that the two foundation models represent document $i$ differently.

We formalize this statement via the following hypothesis test for document $i$:

The hypothesis test requires defining a distance such that we can control for the Type 1 error and still reject the null when appropriate. To this end we propose estimating the null distribution of the distance $d_{i}$ via a RDPG bootstrap. Algorithm 1 describes the bootstrap procedure for generating the null distribution of the $d_{i}$ for each document in $X$. From the null distribution for document $i$ and the observed distance $d_{i}$ we obtain a $p$-value by subtracting the percentile at which $d_{i}$ falls in the null distribution from 1 and reject $H_{0}$ for $p$-values smaller than a pre-specified threshold.

Figure 1 visualizes several steps of the bootstrap when applied to the mean-pooled word embeddings of a pretrained BERT model. In particular, we visualize the bootstrap of a data kernel that BERT induces on $10,000$ English Wikipedia articles. The left panel shows the UMAP [16] projections of $\hat{Z}^{(1)}$, the latent positions estimated from BERT’s data kernel. The center-left panel shows the UMAP projections of the latent position estimates from of the bootstrap data kernels, $ASE(A^{(b)})$. Note that the left and center visualizations exhibit landmark level similarities (e.g. the long peninsula), but are not aligned. The center-right panel shows the UMAP projection of $O_{1,b}$, one of the omnibus latent position estimates. Note that the center-right panel has aligned the previously unaligned landmarks. This alignment allows us to build a null distribution which captures how far latent position estimates from $A^{(b)}$ typically are from $A^{(1)}$. The percentiles of selected null distributions are visualized using concentric circles in the far right panel.

Under the emerging data-centric AI paradigm [28], a model’s behavior is largely determined by the data it is trained on. If this is the case, we would expect interventions into a model’s training data to affect its embedding space, and thus, its data kernel.

To test the effect of a training data intervention on the data kernel, we trained two randomly initialized BERT models – a “baseline" model and a “plant-ablated” model – on two different versions of the DBPedia14 corpus [9]. The DBPedia14 training set consists of 40,000 examples of 14 different classes of documents such as Plants, Animals, Educational Institutions, etc. The baseline model is trained using masked language modeling (MLM) on all classes in the DBPedia14 train set. The plant-ablated model is trained using MLM on all documents in the DBPedia14 training set except those belonging to the plant class. We downsample the total size of the baseline model’s train set uniformly across classes to ensure both models see the same number of tokens during training. Both the plant-ablated model and the baseline model are trained for 3 epochs using the Adam optimizer [11] with a learning rate of $1e^{-4}$ and achieve comparable terminal losses. The data kernels we refer to below are defined on a random 10,000 article sample from the DBPedia14 evaluation set with $k=64$.

Figure 2 shows the result of applying the datum level hypothesis test to the baseline and plant-ablated models. The top left panel of Figure 2 shows the UMAP visualizations of the individual data kernels induced by each of the two models. Note that, while there are some structural similarities between the visualizations of the kernels, they are not directly comparable. The bottom left panel of Figure 2 shows the UMAP visualization of the aligned embeddings from both models. The top center panel of Figure 2 highlights the regions of the aligned visualization that correspond to plant documents. The bottom center panel of Figure 2 shows colors the aligned baseline embeddings by 1-minus the datum-level $p$-value from Algorithm 1. Larger values (represented here by hotter colors) indicate a higher probability that the baseline and plant ablated models’ representations differ on a given document. Note that the two peninsulas consisting predominantly of plant articles appear “hotter" than other regions.

The top right panel of Figure 2 compares the $p$-values under the null distribution calculated using Algorithm 1 to the distribution of a uniform random variable. A hypothesis test is valid (i.e., has Type I error less than a pre-specified threshold) for all pre-specified thresholds if and only if the null distribution of $p$-values is uniform. In our case, the hypothesis test is valid. The bottom right panel compares the cumulative distribution of $p$-values for different classes of DBPedia14 documents. The $p$-values corresponding to plant documents have a left shifted distribution when compared to the $p$-values corresponding to the other classes. Additionally, the $p$-value distributions corresponding to animals and natural places, classes which are semantically related to plants, are also shifted left relative to the other classes. This indicates that the baseline BERT used information contained in the articles about plants to inform its representations of articles about animals and natural places.

This ablation study highlights the ability of our datum-level test to identify systematic differences between models without a predefined metric. Further, this procedure can be used to surface exactly the set of documents whose representations were affected by a data or model intervention. These properties are useful in situations where an intervention systematically affects a class of data (e.g., data corresponding to a particular gender, race, or idea) that is neither explicitly modeled in the existing data ontology nor surfaced via a predefined metric.

The framework we have developed thus far enables datum level pairwise model comparison. However, given the sheer size and variety of the model design space, it is desirable to be able to compare multiple models at once. We can easily extend our pairwise comparison framework to allow for this multi-model comparison.

In section 3, we defined a notion of distance between datum representations based on the Euclidean distance between their aligned latent position estimates. We can extend this notion of distance between datums to a notion of distance between models by considering the spectral norm of the difference between all of the aligned latent position estimates: $\|\hat{Z}^{(m)}_{m,m^{\prime}}-\hat{Z}^{(m^{\prime})}_{m,m^{\prime}}\|_{S}$. Under the assumption that $A^{(m)}$ and $A^{(m^{\prime})}$ are realizations from the same underlying RDPG model, we have

where $\hat{Z}^{(m)}_{m,m^{\prime}}$ and $\hat{Z}^{(m^{\prime})}_{m,m^{\prime}}$ are the corresponding aligned embeddings [13].

Given foundation models $f_{1},\ldots,f_{M}$ and a document corpus $X$, we can calculate the spectral norm of the differences between aligned latent position estimates of the data kernels for each pair of models with the understanding that if $\|\hat{Z}^{(m)}_{m,m^{\prime}}-\hat{Z}^{(m^{\prime})}_{m,m^{\prime}}\|_{S}<\|\hat{Z}^{(m)}_{m,m^{\prime\prime}}-\hat{Z}^{(m^{\prime\prime})}_{m,m^{\prime\prime}}\|_{S}$ then $f_{m}$ is in some sense “closer" to $f_{m^{\prime}}$ than $f_{m^{\prime\prime}}$ (with respect to $X$). Further, we can infer the relative positions of these models on a model-manifold from their pairwise distances using classical multi-dimensional scaling.

Classical multi-dimensional scaling [23] recovers the unknown relative positions of a collection objects given a pairwise Euclidean distance matrix $D$. In particular, for objects $V^{(1)},\ldots,V^{(M)}\in\mathbb{R}^{d}$ where $D_{m,m^{\prime}}=\|V^{(m)}-V^{(m^{\prime})}\|_{2}$ is a pairwise distance matrix on the $V^{(m)}$’s, classical multi-dimensional scaling will recover the relative positions of each $V$ up to an orthogonal transformation. Recovery of the relative positions of objects with multi-dimensional scaling in more general contexts is possible when the objects and dissimilarity used to construct $D$ are Euclidean realizable [5].

In our case, we assume that there is a true-but-unknown vector representation $V^{(m)}\in\mathcal{V}$ of the the foundational model $f_{m}$ in model space with respect to $X$ and that $\hat{Z}^{(m)}$ is a matrix-valued representation of $V^{(m)}$. If we consider the space $\mathcal{V}$ as a parameterization of the space of RDPGs, the dissimilarity matrix with entries $D_{m,m^{\prime}}=\|\hat{Z}^{(m)}_{m,m^{\prime}}-\hat{Z}^{(m^{\prime})}_{m,m^{\prime}}\|_{S}$ is Euclidean realizable given some regularity conditions on $\mathcal{V}$ [2]. The multi-dimensional scaling of $D$ yields vectors $\hat{V}^{(1)},\ldots,\hat{V}^{(M)}\in\mathbb{R}^{d}$ that are Euclidean approximations of the foundation models $f^{(1)},\ldots,f^{(M)}$ relative positions in model space with respect to $X$. We refer to the space spanned by $\hat{V}^{(1)},\ldots,\hat{V}^{(M)}$ as the “model manifold" with respect to $X$. This process – with inputs of a collection of foundation models and a document corpus $X$ – is described in Algorithm 2.

To investigate the properties of the model manifold, we build on the experiment described in Section 3.2. Instead of considering just a plant-ablated model and a baseline model, we now consider models that are trained on various ablated subsets of DBPedia14. We find that the relative positions of models on the manifold induced by our procedure are approximately parameterized by the relative concentrations of DBPedia14 classes in the model training sets. Further, we find that the distance between models on the manifold correlates strongly with the difference in the models’ performance on several downstream metrics. This indicates that proximity on the model manifold may be used as an effective proxy measure for similarity in model performance.

Again, we trained all models on ablated subsets of DBPedia14 using the masked language modeling objective for 3 epochs using the Adam optimizer and a learning rate of $1e-4$. This time, each model had the training data corresponding to some combination of the Plant, Artist, and Educational Institution classes ablated. Recall that each class in the DBPedia14 dataset has 40,000 total documents in it. To determine the particular ablation profile for a model $m$, we first selected an element of the two-dimensional simplex $(p_{1}^{(m)},p_{2}^{(m)},p_{3}^{(m)})\in\Delta^{2}$ (i.e., elements of $\mathbb{R}^{3}$ whose elements are all non-negative and that sum to 1). The model then had access to $p_{1}^{(m)}\cdot 40,000$ documents from the Educational Institution class, $p_{2}^{(m)}\cdot 40,000$, documents from the Artist class, and $p_{3}^{(m)}\cdot 40,000$ documents from the Plant class. The documents from the remaining classes in the training data were then uniformly downsampled to ensure that all models had access to the same number of training tokens. Hence, each model’s training set is naturally parameterized by its corresponding element of the 2-d simplex $(p^{(m)}_{1},p^{(m)}_{2},p^{(m)}_{3})$.

The left panel in Figure 3 shows the theoretical manifold (the two-dimensional simplex which defines the class concentrations in each model’s training set) and the empirical manifold found when applying Algorithm 2 with $k=64$ and aligning the empirical manifold with the theoretical manifold via the Procrustes algorithm. Here, models parameterized by $(1,0,0),(0,1,0),(0,0,1),(0.5,0.5,0),(0,0.5,0.5),(0.5,0,0.5),$ and $(0.33,0.33,0.33)$ are used to estimate the and visualize model manifold. The empirical manifold is visually similar to the theoretical manifold, and we interpret this as evidence that the models, through their data kernels, exist on a low dimensional structure that (in this case) is parameterized by the concentrations of classes in each model’s training set. More succinctly, $\hat{V}^{(m)}\approx(p^{(m)}_{1},p^{(m)}_{2},p^{(m)}_{3})$ in this experiment.

To understand how different downstream metrics correlate with distance on the recovered model manifold, we compute a metric relevant to document classification and a metric relevant to language modeling and regress them with respect to the manifold distance to a landmark model. For both tasks we select the model parameterized by $(1,0,0)$ (i.e. the model trained on all of the Educational Institution articles, but none of the Artist or Plant Articles) as the landmark model.

For the classification task we consider a uniform subsample (500 documents per class) of the DBPedia14 evaluation set. We train a linear support vector machine (SVM) on the mean-pooled embeddings from each model using a random 80% of the selected evaluation data, and then infer the class membership of the remaining 20%. ion for a document agrees with the prediction from the classifier trained on embeddings from the landmark model. We then compute the correlation between the predictions of each model’s SVM and the landmark model’s SVM. In the right panel of Figure 3, we show that manifold distance between models is a strong indicator of their prediction correlation.

For the language modeling task, we randomly sampled $1,000$ documents in the Educational Institution class from the DBPedia14 evaluation set. We compute the average pseudo-perplexity [21] of each document in our evaluation set, and report the average pseudo-perplexity across all evaluation documents for each model in the center panel of figure 3. We find that models that were trained with a higher concentration of Educational Institution documents have both a lower average pseudo-perplexity and are closer on the manifold We find that the concentration of Educational Institution documents in the training set correlates negatively with both average pseudo-perplexity and manifold distance to the landmark model

The linear goodness-of-fit ($R^{2}$) for both the relationship between manifold distance and prediction correlations as well as the relationship between manifold distance and pseudo-perplexity are above 0.8, and the $p$-values corresponding to Kendall’s tau rank correlation test are less than 0.01. We interpret these results as evidence that the manifold distance and metrics related to important downstream tasks – classification and language modeling – are strongly correlated.

While we are excited about the potential of data kernel based foundation model comparison, there are several important limitations of our current work.

The most salient limitation is the scope of our experiments – we have opted to evaluate the behavior of our comparison framework in several highly controlled situations. The benefit of this choice is that we can ensure that the framework behaves appropriately under known conditions. The limitation of this choice is that it provides limited evidence that the framework will generalize to less controlled, empirical situations. We aim to address this question of generalization in subsequent work on the empirical model manifold.

An additional limitation of our work is the computational burden required for data kernel based comparison. Each iteration of our bootstrap procedure requires the computation of a large singular value decomposition as part of the adjacency spectral embedding. This presents a significant challenge to practitioners who wish to access our comparison methodology, but have a limited compute budget. We encourage future work on more efficient bootstrap methods that can help to alleviate this burden, and aim to address it as part of our future research agenda.

While we view the development of higher fidelity model evaluation procedures as generally positive from a broader impact point of view, there are a few important caveats to note. As is the case for any evaluation, there is a risk that our framework fails to surface a systematic bias present in a model. This risk is particularly present in our work, as a failures of our framework may lead practitioners to have a false sense of security regarding the fairness of their models. We emphasize that model evaluation is a nuanced and complex question, and that the best approaches combine several model evaluation strategies to paint a maximally holistic picture of model behavior.

Moreover, the computationally intensive nature of the work has a direct environmental impact. In particular, we estimate that this work produced at most 3 metric tons of carbon dioxide emissions, which is roughly the equivalent of burning 350 gallons of gasoline. Nomic is committed to tracking and mitigating the environmental impact of its research, and has offset 3 metric tons of emissions for this project.

Given the large and growing variety of foundation model variants, we are interested in using our procedure to build a manifold of empirical models. Technologies such as the Hugging Face Hub [26] enable us to build model manifolds on large collections of models and datasets. Having access to a large empirical model manifold will allow researchers to focus their analysis on models that exist in key locations on the manifold (e.g. centers of mass, branch points, etc.). Additionally, having a notion of distance between empirical models will allow researchers to generalize their findings regarding particular models to families of models that are proximal on the manifold.

The model manifold comes naturally equipped with a distance that captures information pertinent to relative performances of the models. We can leverage this fact for a variety of model selection tasks. In particular, the model manifold allows model selection to be framed as a search task over the model manifold itself. For example, users can select the best model for a particular task with far fewer evaluations than an exhaustive search by employing a locally sensitive optimizer on the model manifold. Additionally, the task of selecting the best model for a task subject to a constraint (e.g., a compute constraint) given a known good model that does not fulfill that constraint is akin to projecting the known good model onto a constrained subset of the model manifold.

The model manifold can be employed to study populations of users communicating on a platform without requiring direct access to the content of their communications. One such procedure would involve fine tuning an edge language model to an individual user’s communications, and relaying only the embeddings of the fine tuned model on a public corpus of interest back to the central server. A model manifold where each model corresponds to a user could then be constructed and utilized for subsequent inference.

One of the fundamental challenges of scaling up our work is the computationally intensive nature of our bootstrap procedure. In particular, the singular value decomposition required for each iteration of our bootstrap presents a significant computational bottleneck. There is a rich and active literature on efficient singular value decomposition algorithms, and we hope to evaluate the viability of ideas from this literature in future versions of our bootstrap.