1 Introduction↩︎

Large Language Models (LLMs) have been utilized for a wide range of tasks [1]–[3], including but not limited to question answering, code generation, reasoning, and domain-specific assistants [4]. Concurrently, the number of available LLMs has also grown exponentially, with hundreds of proprietary models and thousands of open-source models.¹ These models vary greatly in size, architecture, training data, and expertise [5], offering opportunities to match inference queries to the best-suited model [6], [7]. Utilizing and managing a vast pool of language models raises many practical challenges [8], [9]. How can we compare models across sizes and architectures [10]? How can we map model abilities to find similarities or differences among them [11]? How can we minimize redundancy in a model pool by retaining a small subset of models while preserving their aggregated abilities [12]? How can we maintain a dynamic pool [13] where new models are frequently added and older models are deprecated? Addressing these questions necessitates a representation that makes model capabilities easy to summarize and compare at scale.

We propose a simple yet powerful idea: represent each model by a fixed-dimensional vector that summarizes its capabilities across queries. We refer to this representation as a model embedding.

1.0.0.1 Why are model embeddings useful?

By embedding models in a shared vector space, we can compare and organize models using simple approaches like nearest-neighbor lookup [14] and clustering [15]. This supports model discovery and organization (e.g., finding similar models, grouping model families, and understanding joint capability profile and redundancy), portfolio selection (choosing a smaller set of models while covering all capabilities), and monitoring (for security purposes and tracking model drift). Model embeddings are also useful for decision making at inference time: they can support query routing by estimating which models are likely to perform well on a given query, allowing selection of models best suited for each query.

1.0.0.2 Information sources for model embeddings.

To capture capabilities of language models, we require an information source that reflects it. In the full information setting, model parameters can be directly used for generating such embeddings, but it is infeasible due to heterogeneous model sizes and architectures [9] which are often proprietary [1], [16], [17] and accessible only via APIs. An alternative is to distill model capabilities [18] from output logit distributions of generated tokens, which can be problematic when there are different tokenizations used across models [19]. Instead, a practical way of capturing the capabilities of diverse models is by observing their generated responses to queries, which can be assigned performance scores (e.g., correctness value) by comparison with ground truth answer [20] or by using LLM-as-a-judge [21]. Concretely, for a text query \(x\) and model \(m\), we observe a text response and compute a score \(y^{(m)}(x)\), often simply the binary correctness value. For each model \(m\), we aim to map the observed set of evaluations into a vector \(z_m\) which summarizes its capabilities.

1.0.0.3 Desiderata.

The aforementioned use cases inform key requisites of an embedding generation approach. Concretely, we seek embeddings with the following properties:

• Black-box compatible: rely only on queries and evaluation scores, without requiring access to model weights, internal activations, or logit distributions.

• Compatible with varying evaluations: do not require all models to be evaluated on the same fixed query set to obtain comparable embeddings.

• Sample efficient: require only a few query evaluations to produce informative embeddings, which improve as more evaluations are collected.

• Training-free onboarding: adding new models should not require training model-specific parameters, nor affect existing model embeddings.

• Informative of capabilities: support accurate performance prediction on unseen queries.

• Geometrically meaningful: proximity under distance measures (e.g., cosine or Euclidean) should reflect model similarity.

1.1 Related works.↩︎

Prior work on finding vector representations of language models can be broadly grouped into parametric and nonparametric approaches. 1 provides a conceptual comparison of existing work with our method.

1.1.0.1 Parametric approaches.

A common parametric design learns per-model embeddings as trainable parameters along with a correctness predictor trained on model–query evaluations. EmbedLLM [22] follows this general design, while IRT-Net [23] and JE-IRT [24] use an item response theory formulation for interpretable predictions [25]. Although effective for correctness prediction and routing, these embeddings are learned via stochastic gradient descent (SGD) and are therefore not uniquely determined by the evaluation data: different training runs on the same evaluation data can produce varying embeddings, an issue exacerbated by overparameterization of the predictor and embeddings. This renders distance-based analyses (e.g., nearest neighbors, clustering, or similarity search) on embedding geometry brittle. Moreover, adding new models requires training new embeddings. A related parametric approach [26] utilizes matrix factorization to represent models in query encoding space. See 10 for additional remarks and experiments demonstrating these limitations.

1.1.0.2 Nonparametric approaches.

Nonparametric methods utilize common evaluation queries to directly compute model embeddings using geometric operations, without learning any parameters. LLM-DNA [27] concatenates query and response encodings and downprojects them using a Gaussian random matrix to obtain model embeddings. All models are required to be evaluated on the same queries, which does not generalize well to partial evaluations. Additional evaluations of queries cannot be utilized for refining the embeddings, and there is no correctness predictor that can be directly used for query routing.

1.2 Approach Overview and Contributions↩︎

3pt

Motivated by strengths of both parametric and nonparametric embedding generation approaches, we propose LOw-dimensional embeddings for Model Capability UnderStanding (LOCUS), an attention-based method that processes a collection of queries evaluated on a model to generate a fixed-dimensional model embedding via an encoder \(F_\theta\). After training the encoder once, generating an embedding for a model (including newly added models) is just a deterministic forward pass over its available evaluations, and the embedding can be further refined by adding more evaluations. In addition, we train a correctness predictor \(G_\psi\) which processes a model embedding and a query encoding to predict the correctness of model on that query.

The attention layers in LOCUS jointly process the available query evaluations, allowing them to extract a strong indicator of model capabilities even from limited data. We validate this insight in 1, where LOCUS achieves higher routing accuracy across a wide range of training-set sizes, with up to \(\mathbf{4.8\times}\) better sample efficiency than prior baselines. The embeddings are less prone to overfitting to a particular model’s evaluation set and remain stable as evaluation queries vary. Finally, the rich geometry of the embedding space enables several downstream tasks, including identifying model families, selecting model portfolios, and allowing resilient fallback decisions.

The remainder of the paper is organized as follows: 2 introduces our approach; 3 reports results and analyses; and 4 concludes our work.

2 Using Attention for Model Capability Representation↩︎

In this section, we describe an attention-based encoder (\(F_\theta\)) for generating model embeddings from scored query evaluations, together with a query correctness predictor (\(G_\psi\)) that estimates a model’s probability of correctness on a test query. Each query \(x\) is represented by a fixed dimensional encoding/embedding \(\phi(x)\in\mathbb{R}^{d_\phi}\) (e.g., generated by a pretrained sentence encoder [28], [29]). Let \(\mathcal{Q}_m\) denote the set of queries model \(m\!\in\!\mathcal{M}\) is evaluated on. Let normalized evaluation scores be \(y^{(m)}(x_i)\!\in[0,1]\) (typically binary correctness \(\in\!\{0,1\}\)) for all queries \(x_i\!\in\!\mathcal{Q}_m\). The evaluation set for model \(m\) is thus \[\begin{align} S_m \;=\; \{(\phi(x_i),\, y^{(m)}(x_i))\}_{x_i\in\mathcal{Q}_m} \end{align}\] with \(n_m := |S_m|\) denoting the number of evaluations for model \(m\).

In 2.1, we describe a tokenization scheme for individual evaluations, present the standard and an efficient variant of multi-head attention over evaluation tokens, and a final aggregation layer which produces model embeddings. 2.2 describes query correctness predictor, and 2.3 elaborates on the training procedure for (\(F_\theta\),\(G_\psi\)).

2.1 Attention-based model capability encoder \(F_\theta\)↩︎

To process the individual query evaluations in \(S_m\) for generating model embeddings, we need to convert each evaluation pair of query encoding and correctness score into a systematic representation for facilitating calculations.

2.1.0.1 Continuous tokenization of evaluations.

We convert each evaluation pair \((\phi(x_i), y^{(m)}(x_i))\) into a \(d\)-dimensional token \(t^{(m)}_i\) using a trained MLP \(h_\omega\!:\!\mathbb{R}^{d_\phi}\!\times\!\mathbb{R}\!\rightarrow\!\mathbb{R}^{1\times d}\) as: \[\begin{align} \fbox{\;t_i^{(m)} \;=\; h_\omega\!\left(\phi(x_i),\; y^{(m)}(x_i)\right)\in\mathbb{R}^{1\times d}, \;\;} \end{align}\] and stack tokens as rows to form \[\begin{align} \label{eq:tokenized95evaluations} X_m^{(0)} \;=\; \begin{bmatrix} t_1^{(m)} \\ \vdots \\ t_{n_m}^{(m)}\end{bmatrix}\in\mathbb{R}^{n_m\times d}. \end{align}\tag{1}\] Intuitively, \(h_\omega\) acts as a continuous tokenizer mapping heterogeneous inputs (query features + score) into a common token space suitable for attention operations. We view the tokenizer \(h_\omega\) as the first stage of the encoder and define \(X_m^{(0)}=h_\omega(S_m)\). The subsequent attention layers then map \(X_m^{(0)} \mapsto z_m\), and by slight abuse of notation we write \(F_\theta(S_m)\equiv F_\theta(X_m^{(0)})\). The encoder \(F_\theta\) ultimately generates a single model embedding \(z_m\in\mathbb{R}^{d}\) from the query evaluations.

We require the model embeddings to depend only on the query evaluation set \(S_m\), and not on the particular order in which they are listed in \(X_m^{(0)}\). To achieve this, we apply bidirectional attention without positional encodings. Under this design choice, shuffling the rows of \(X_m^{(0)}\) only shuffles the intermediate token representations, implying permutation equivariance. We finally apply a learned-query attention aggregation step (4 ) which maps the last layer’s output tokens into a single model embedding vector \(z_m\). This operation performs aggregation without any position specific weights, resulting in a final embedding that is permutation invariant to token ordering in \(X_m^{(0)}\) (1 ).

2.1.0.2 Attention over evaluation tokens.

We denote a Transformer block [30] by \(\mathrm{TBlock}(\cdot,\cdot,\cdot)\), consisting of (i) a Multi-Head Attention (MHA) sublayer with a residual connection and layer normalization (\(\mathrm{LN}_0\)), followed by (ii) a position-wise feed-forward network (\(\mathrm{FFN}\), a single hidden layer MLP \(d\!\to\! d_{\mathrm{ff}}\!\to\! d\)) applied independently to each token along with another residual connection and layer normalization (\(\mathrm{LN}_1\)): \[\bar{Q} = \mathrm{LN}_0\Big(Q + \mathrm{MHA}(Q,K,V)\Big)\in\mathbb{R}^{n_q\times d},\] \[\mathrm{TBlock}(Q,K,V) = \mathrm{LN}_1\Big(\bar{Q} + \mathrm{FFN}(\bar{Q})\Big)\in\mathbb{R}^{n_q\times d}.\]

Let \(X_m^{(\ell)}\in\mathbb{R}^{n_m\times d}\) denote the output tokens after \(\ell\) layers of Transformer blocks (\(X_m^{(0)}\in\mathbb{R}^{n_m\times d}\) are input evaluation tokens in 1 ). Our self-attention block at layer \(\ell\) over evaluation tokens is thus: \[\begin{align} X_m^{(\ell)} \;=\; \mathrm{TBlock}\big(X_m^{(\ell-1)},\,X_m^{(\ell-1)},\,X_m^{(\ell-1)}\big)\in\mathbb{R}^{n_m\times d}. \end{align}\]

2.1.0.3 Latent bottleneck attention block.

Direct self-attention over \(n_m\) tokens incurs \(\mathcal{O}(n_m^2)\) computational cost, which can be restrictive with large number of evaluations \(n_m\). We reduce this by introducing \(r\ll n_m\) learned latent vectors \(U^{(\ell)}\in\mathbb{R}^{r\times d}\) which act as an attention bottleneck in each layer. Concretely, one latent bottleneck attention block consists of two Transformer blocks: a compression block where \(U^{(\ell)}\) attends to the evaluation tokens \(X_m^{(\ell-1)}\), followed by a broadcast block in which the evaluation tokens attend to the latent outputs \(H_m^{(\ell)}\): \[\begin{align} { \setlength{\fboxsep}{3pt} \fbox{ \begin{figure}\includegraphics[width=0.8\textwidth]{_pdflatex/wqlnbfvx.png}\tag{2}\end{figure} } } \tag{3} \end{align}\] This reduces the overall attention computation cost to \(\mathcal{O}(n_m\!\times\!r)\) while still enabling full information flow through the attention layers [31]–[33].

2.1.0.4 Learned-query attention aggregation.

After \(L\) latent-bottleneck attention blocks, we aggregate the \(n_m\) output token features into a single model embedding \(z_m\) using a trained query parameter \(s\in\mathbb{R}^{1\times d}\), which is learned with other model parameters during training: \[\begin{align} \label{eq:learned95query95aggregation} \fbox{\;z_m \;=\; \mathrm{TBlock}(s,\; X_m^{(L)},\; X_m^{(L)}) \in \mathbb{R}^{1\times d}.\;} \end{align}\tag{4}\] This operation performs attention-based aggregation, where a fixed, learnable query attends over a variable number of token features [31], [34]–[36], generating a single output vector which we denote as the model embedding. As this operation depends only on attention over the token features and does not use positional encodings, the resulting \(z_m\) is invariant to permutations of the input evaluations \(X^{(0)}_m\) in 1 .

2.2 Query correctness predictor \(G_\psi\)↩︎

Generated model embeddings \(\{z_m\}_{m\in\mathcal{M}}\) in 4 are actionable for model selection only if we have a mechanism for translating the encoded information into correctness prediction of models for unseen queries. Such query-wise correctness prediction scores for language models can be used for a multitude of applications such as routing, confidence-based majority averaging of responses, ranking model performances etc. To this end, we design a correctness predictor (decoder) that processes query encodings along with model embedding vectors to return a probability of correctness value. Formally, given model embedding \(z_m\) and query encoding \(\phi(x)\), the decoder predicts \[\begin{align} \label{eq:correctness95probability} \widehat{p}_\psi\left(y^{(m)}(x)\!=\!1 \Big| z_m, \phi(x)\right)=\sigma\big(G_\psi\left(z_m,\phi(x)\right)\!\big), \end{align}\tag{5}\] where \(G_\psi\) is a lightweight MLP producing a scalar logit and \(\sigma(\cdot)\) is the sigmoid function.

The overall architecture thus consists of two components: the encoder \(F_\theta\) is responsible for producing a capability representation \(z_m\) that is useful across queries, while the predictor \(G_\psi\) combines \(z_m\) with query features to output a correctness probability value.

2.3 Training and usage↩︎

We train the model encoder \(F_\theta\) (including tokenizer \(h_\omega\)) and decoder \(G_\psi\) (correctness predictor) jointly from a supervised dataset containing model-query evaluation pairs and observed correctness values. Each step samples a mini-batch of models; for each model \(m\), we form (i) an encoder input subset \(S_m^{\mathrm{enc}}\) of scored evaluations to build \(z_m\), and (ii) a decoder query batch \(S_m^{\mathrm{dec}}\) on which we apply the correctness prediction loss. Additional details provided in 12.

2.3.0.1 Batch structure.

For any model \(m\), let \(S_m^{\mathrm{enc}}\!\subseteq\!S_m\) be a randomly chosen subset of query evaluations, and let \(z_m\!=\!F_\theta(S_m^{\mathrm{enc}})\) denote the corresponding model embedding generated. Independently, we also randomly sample decoder queries \(S_m^\mathrm{dec}\!\subseteq\!S_m\) (with target labels \(y^{(m)}(x)\) for \(x\!\in\!S_m^\mathrm{dec}\)) and compute correctness probabilities in 5 as: \[\begin{align} \widehat{p}^{(m)}(x) \triangleq \widehat{p}_\psi\!\left(y^{(m)}(x)\!=\!1 \Big| z_m, \phi(x)\right). \end{align}\]

2.3.0.2 Training Loss.

We minimize binary cross-entropy loss averaged over the models in a mini-batch and decoder queries: \[\begin{align} \min_{\omega,\theta,\psi}\; \mathbb{E}_{m}\; \mathbb{E}_{x\sim S_m}\; \mathrm{BCE}\!\left(\widehat{p}^{(m)}(x),\, y^{(m)}(x)\right). \end{align}\]

2.3.0.3 Generating embeddings for new models.

To embed a new model \(m_\mathrm{new}\), we collect an evaluation set \(S_{m_\mathrm{new}}^\text{enc}\), and perform a forward pass over the embedding generator to obtain \(z_{m_\mathrm{new}}=F_\theta(S_{m_\mathrm{new}}^\text{enc})\). No retraining or learning any additional parameters is required at embedding time.

2.3.0.4 Correctness prediction and routing.

We estimate correctness probabilities \(\widehat{p}^{(m)}(x)\) for candidate models and query \(x\) by using the correctness predictor \(G_\psi\). These probabilities can be used directly for ranking models, routing to a suitable model for response generation, etc.

3 Experiments↩︎

We evaluate our attention-based model embedding generator and query correctness predictor on five distinct aspects:

1. Correctness prediction and routing accuracy

2. Sample Efficiency during training and test time

3. Robustness of generated embeddings to changes in the evaluation set queries, both in terms of numerical performance and embedding geometry

4. Correlation between embedding distance and correctness similarity of models

5. Practical utilities enabled by model embeddings, including fallback routing and model portfolio selection.

3.0.0.1 Models, tasks, and evaluation sets.

In our experiments, we consider a set of \(112\) language models [22] spanning a wide variety of sizes and capabilities, including both general and finetuned models. The queries in our analysis are sampled from \(10\) public benchmarks, and represent a real-world corpus across a variety of tasks such as math, general knowledge, reasoning etc. Additional details on language models used and query datasets are provided in 5. We select the sentence encoder all-mpnet-base-v2 [37] as \(\phi(x)\) for generating query representations in our analysis, and defer ablations on other sentence encoders to 7.

Our attention-based encoder² has two latent bottleneck attention blocks (3 ) (\(L\!=\!2\)), \(4\)-headed MHA (Multi-Head Attention) sublayers, \(r\!=\!64\) learned latent vectors in each block, followed by attention-based aggregation layer (4 ). The correctness predictor (5 ) is a single hidden-layer MLP (\(d_\mathrm{hidden}\!=\!64\)). The model embedding dimension is set to \(d\!=\!128\) (results on varying embedding dimensions deferred to 8). We do not utilize any test queries for generating model embeddings; they are only used for evaluating correctness prediction and performance agreement rate of models. Further details in 12.

3.0.0.2 Baselines.

We compare against two baselines that obtain model embeddings via gradient-based training from query evaluation data: EmbedLLM [22], and IRT-Net [23].

3.0pt

3.1 Correctness Prediction and Routing Accuracy↩︎

For each model \(m\) and query \(x\), the correctness predictor estimates \(\widehat{p}^{(m)}(x)\!=\!\sigma(G_\psi(z_m,\phi(x)))\!\in\![0,1]\). We convert the probability estimate into a binary prediction \(\widehat{y}^{(m)}(x)\!=\!\mathbb{I}\{\widehat{p}^{(m)}(x)\!\geq\!0.5\}\) by thresholding, and then average the agreement indicator \(\mathbb{I}\{\widehat{y}^{(m)}(x)\!=\!y^{(m)}(x)\}\) over all model-query pairs to calculate correctness prediction accuracy. Additionally, we judge whether the language model with highest predicted correctness probability \(\widehat{p}^{(m)}(x)\) actually generated a correct response to query \(x\), and report an average over all test queries as the overall routing accuracy.

Across training set sizes, LOCUS outperforms examined baselines in overall routing accuracy and correctness prediction accuracy as presented in 1 ¿tbl:tab:concise95routing95correctness?. We attribute this to the ability of attention blocks in extracting model capability patterns through query evaluations, instead of using backpropagation to obtain model embeddings that describes observed correctness patterns. We extend further examination on numerical performance in [app:routing_correctness_prediction] [app:encoder_ablation] [app:embedding_dim_ablation], and wall clock time complexity in 11.

3.2 Sample Efficiency↩︎

A key motivation for model embeddings is sample-efficient learning and new model onboarding: we aim to train a general-purpose model capability encoder from limited evaluations, and efficiently embed previously unseen models using small number of query evaluations.

3.2.0.1 Training-time sample efficiency.

1 reports routing accuracy as a function of the number of query evaluations per model used to train \((F_\theta,G_\psi)\) (i.e., \(n_m=|S_m|\)). Our method reaches strong correctness prediction and routing performance with substantially fewer training evaluations, resulting in upto \(\mathbf{4.8 \times}\) sample efficiency as compared to existing baselines. Notably, our training procedure does not require all models to share the same evaluation queries, nor to be evaluated on the same number of queries, further increasing practicality of our approach.

3.2.0.2 Test-time embedding efficiency (new model onboarding).

To measure sample efficiency for generating embeddings for unseen models, we train \((F_\theta,G_\psi)\) on query evaluation data from # models \(\in\{8,16,\ldots,92\}\) (\(1024\) evaluations each) and then embed \(16\) held-out models using # query evaluations \(\in\{16,32,\ldots,1024\}\) (details in 5). 3 shows that as few as \(\approx\) \(128\) evaluations are sufficient to obtain informative embeddings. We also observe negligible gap (\(<1\%\)) between correctness prediction accuracy using partial data trained versus full data trained \((F_\theta,G_\psi)\), highlighting generalization to unseen models.

3.3 Robustness to Varying Query Evaluations↩︎

We test how sensitive model embeddings are to the query evaluation set used as input to the encoder. Using a pretrained encoder \(F_\theta\) trained from evaluations on the full model pool, we first compute reference embeddings \(z_m^{\mathrm{ref}}\) for each model \(m\) from a fixed reference query evaluation set of size \(n_{\mathrm{ref}}\) (e.g., \(4096\)). We then perturb the input evaluations in two ways: changing which queries are used (fixed number but different queries) and changing how many queries are used (subsampling from reference evaluation set), and then analyze the stability of embeddings via routing/correctness prediction and embedding geometry.

3.3.0.1 Robustness to the choice of evaluation queries.

Keeping the number of queries in the evaluation set fixed at \(n_{\mathrm{ref}}\), we resample alternate evaluation sets whose overlap with the reference set is \(\alpha\in[0,1]\). We then recompute embeddings for all models, and calculate numerical metrics using the new embeddings. 4 shows that correctness prediction and routing accuracy remain stable across overlap fractions, and 5 illustrates that the regenerated embeddings stay close to originally generated embeddings \(z_m^{\mathrm{ref}}\) (downprojection via t-SNE [38]).

3.3.0.2 Robustness to the number of evaluation queries.

We also subsample \(n\) evaluations from the reference set of query evaluations (e.g., \(n\in\{64,128,256,\ldots,n_{\mathrm{ref}}\}\)) and recompute embeddings. Numerical performance saturates quickly: embeddings formed from only \(\approx128\)–\(256\) evaluations typically match the routing and correctness prediction accuracy obtained with larger sets (4). The corresponding t-SNE projections of embeddings in 5 further suggest geometric stability under subsampling evaluation queries.

3.3.0.3 Implications for model fingerprinting (security).

Robustness to query evaluations provides a preliminary fingerprinting signal: if a new model’s embedding repeatedly maps near the same reference model across independently sampled evaluation sets, it suggests potential model duplication and can be flagged for audit (see 9.3 for additional experiments) [39], [40].

3.4 Embedding Geometry reflects Model Similarity↩︎

We next test whether distances in the learned embedding space track differences between model capabilities. Concretely, for every model pair \((m_1,m_2)\in \mathcal{M}\!\times\!\mathcal{M}\) we compute (i) an embedding-space distance \(d_z(z_{m_1},z_{m_2})\) (cosine and Euclidean), and (ii) their correctness disagreement rate on a common test query set, defined as the fraction of queries on which their binary correctness labels differ (i.e., the normalized Hamming distance, or \(1-\)agreement). [fig:performance_distance_correlation,tab:performance_distance_correlation] shows a strong positive relationship between embedding distance and disagreement across all correlation measures: for example, Pearson correlation is \(0.845\) (cosine) and \(0.887\) (Euclidean), with similarly high rank correlations (Spearman \(\approx0.88\) and Kendall \(\approx0.71\)). These results indicate that embedding distances convey meaningful information about model similarity.

4.0pt

3.4.0.1 Recovering model families via hierarchical clustering.

We perform hierarchical clustering [41] of models using pairwise distances (Euclidean distance) between embeddings. The dendrogram and distance heatmap in 7 reveal coherent groupings that align with known model specializations (e.g., math- and code-finetuned families), suggesting that the embedding geometry also captures shared capability structure at the level of model families rather than just individual models.

3.5 Practical utilities enabled by model embeddings↩︎

We finally demonstrate that compact, stable capability embeddings enable useful operations beyond standard routing or correctness prediction tasks.

3.5.0.1 Nearest neighbors as proxies under model unavailability.

In practice, a preferred model for generating response to a query can be inaccessible due to capacity limits, hardware failure, policy constraints etc. In such cases, it is useful to identify proxy models with close matching capabilities. We therefore test whether nearest neighbors in embedding space can serve as reliable proxies. For each model \(m\), we rank its neighbors by embedding distance (e.g., cosine distance) and calculate correctness agreement on a common set of queries. 8 reports the average agreement of \(k^\text{th}\) neighbor across all models in the pool, with closest neighbor correctness agreement as high as \(\mathbf{79\%}\) and decreasing gradually with further neighbors (larger \(k\)). We extend the implications of this to resilient routing in 9.2.

3.5.0.2 Model portfolio selection from embeddings.

We evaluate whether the model embedding space can support practical deployment planning: selecting a portfolio of models such that routing among the selected models retains high accuracy on queries as compared to the full model pool (i.e., selecting a subset of models to eliminate redundancy in capabilities across models while maintaining performance).

Count-constrained portfolios. When limiting deployment to \(k\) models, we select diverse portfolios by maximizing coverage in embedding space using k-center (minimizes the maximum distance between all models’ embeddings to those in the chosen portfolio) [42] and k-medoids (minimize average distance of all models’ embeddings to models in chosen portfolio) [43]. We observe that a \(\mathbf{15}\)-model portfolio captures full routing accuracy of \(112\) models in the pool. Detailed algorithm is deferred to 9.4.2. We report routing accuracy restricted to models in the portfolio as \(k\) varies (9).

Parameter-budget portfolios. We solve a total-parameter-budget-constrained optimization problem (coverage-greedy): select a model subset which covers the model embedding space while respecting the overall parameter budget (9.4.3). 9 depicts strong routing performance with relatively small parameter budget (full pool of \(112\) models has a total of \(1930\)B parameters).

3.5.0.3 Additional utilities.

Beyond model capability proxies and portfolio selection, model embeddings enable other practical operations such as resilient routing, retrieving a real model that best matches a desired capability profile etc. We provide further discussion on additional use cases in 9.

4 Concluding Remarks↩︎

In this paper, we designed LOCUS, a novel attention-based framework that projects Large Language Models into a low-dimensional embedding space based on their inference capabilities. Unlike prior approaches, our method supports seamless, training-free onboarding of new models without disturbing existing embeddings, while demonstrating superior sample efficiency, requiring up to \(4.8 \times\) fewer query evaluations to generate robust representations. Beyond standard performance prediction, we demonstrated that the learned embedding geometry effectively captures behavioral similarity, enabling diverse downstream utilities such as optimal model portfolio selection and resilient fallback routing under model unavailability. As the ecosystem of proprietary and open-source models continues to expand, such capability-oriented embeddings offer a scalable, geometrically meaningful foundation for efficient model management and discovery. We leave extensions to multimodal models and adaptive query evaluations for generating embeddings as future directions.

Acknowledgments↩︎

This work was partially supported by NSF grants CCF 2045694, CNS-2112471, CPS-2111751, ONR grant N00014-23-1-2149, and an AI2C Seed grant. This work used Bridges-2 GPU at the Pittsburgh Supercomputing Center through allocation CIS250429 from the Advanced Cyberinfrastructure Coordination Ecosystem: Services & Support (ACCESS) program, which is supported by NSF grants #2138259, #2138286, #2138307, #2137603, and #2138296 [44].

Impact Statement↩︎

This paper presents work whose goal is to advance the field of Machine Learning. In particular, we design low-dimensional model embeddings that can be used to efficiently manage and utilize large model pools. There are many potential societal consequences of our work, none of which we feel must be specifically highlighted here.

5 Language Models and Query Datasets↩︎

We use the model-query evaluations presented in [22]. The queries are sampled from ten public benchmarks, representing a wide variety of real-world prompts and tasks: MathQA [45], LogiQA [46], MedMCQA [47], PIQA [48], TruthfulQA [49], MMLU [50], GSM8K[51], GPQA[52], ASDiv[53], SocialQA [54]. Semantic query encodings generated by sentence encoders are presented in 10 for four choices of sentence encoders, colored by dataset.

The 112 language models used are presented in 3 with their HuggingFace identifiers [55]. For new model onboarding analysis as presented in [subsec:exp_sample_efficiency] 3, we consider the following sixteen models as held-out from training (seen only during test time): FelixChao/llama2-13b-math1.2, MaziyarPanahi/WizardLM-Math-70B-v0.1, 01-ai/Yi-6B-200K, WizardLM/WizardLM-70B-V1.0, bigscience/bloom-7b1, sail/Sailor-7B, codellama/CodeLlama-13b-Instruct-hf, Writer/palmyra-med-20b, Qwen/Qwen1.5-0.5B-Chat, databricks/dolly-v2-12b, nomic-ai/gpt4all-13b-snoozy, stabilityai/stablelm-tuned-alpha-7b, AdaptLLM/medicine-chat, AdaptLLM/medicine-LLM, EleutherAI/pythia-12b, Q-bert/Optimus-7B.

6pt

6 Extended Numerical Comparison with Baselines↩︎

We present the extended numerical results on routing accuracy (4) and correctness prediction (5) of LOCUS, compared against three other baselines: EmbedLLM [22], EmbedLLM+ (more hidden layers in decoder with same architecture as original [22]), and IRT-Net [23]. We additionally present correctness prediction and routing accuracy of \(K\)-Means and \(k\)NN based approaches [20], [57]–[59]. \(K\)-Means routers perform clustering over query encodings and assign average model correctness on closest cluster for test queries as the predicted correctness probability, whereas \(k\)NN routers average model correctness over nearest \(k\) queries in encoding space as the prediction for a test query.

We observe that across all model embedding based approaches, LOCUS exhibits the highest overall routing accuracy (4) across all choices of number of training samples (number of queries each model is evaluated on). This is succinctly described in 1 as well. For correctness prediction accuracy (5) we find that all approaches perform similarly in terms of overall correctness prediction accuracy. Among nonparametric routing techniques, \(K\)-Means outperforms other approaches on certain choices of training data size, but we highlight that we cannot generate model embeddings and support their utilities via \(K\)-Means clustering based routers.

2.5pt

2.5pt

 

7 Ablation on choice of Sentence Encoder for Query Representation↩︎

2.5pt

2.5pt

In this section, we present correctness prediction accuracy and routing accuracy of LOCUS through different choices of sentence encoders for generating query representations. In our main analysis, we select a popular sentence encoder all_mpnet_base_v2 [37] which produces query representations \(\in\mathbb{R}^{768}\).

Here, we additionally present results for paraphrase_albert_small_v2 [29] (\(\in\mathbb{R}^{768}\)), all_minilm_l6_v2 [28] (\(\in\mathbb{R}^{384}\)), and all_distilroberta_v1 [56] (\(\in\mathbb{R}^{768}\)). Our model architecture is uniform across all sentence encoders (except for different input dimensions across choices of sentence encoders). We present routing accuracy in ¿tbl:tab:routing95acc95by95dataset95full95encoder95ablation? and correctness prediction accuracy in ¿tbl:tab:correctness95pred95by95dataset95full95encoder95ablation?.

We observe that overall routing accuracy and correctness prediction accuracy is comparable across sentence encoders, but individual dataset-wise performance varies. We attribute this to varying query representation capabilities across tasks for different sentence encoders.

8 Ablation on Embedding Dimension↩︎

Here, we present experimental results on our encoder-decoder approach \((F_\theta,G_\psi)\) with varying model embedding dimension (and corresponding internal dimensions as well). 6 denotes routing accuracy, and 7 denotes correctness prediction accuracy for embedding dimensions \(\in\{16,32,64,128,256,512\}\). We find that model embedding dimensions \(128/256\) show superior routing accuracy and correctness prediction accuracy, whereas smaller embedding dimensions are not informative enough and larger embedding dimensions increase training complexity, eventually affecting numerical performance.

2.5pt

2.5pt

9 Additional Utility of Model Embeddings↩︎

9.1 Retrieving models that match desired task profiles via hypothetical embeddings.↩︎

We next test whether embeddings preserve task-level performance structure rather than overfitting to query identity. For each model \(m\in \mathcal{M}\), we compute its empirical task-wise accuracy profile across a set of tasks. We then construct a hypothetical evaluation set by sampling queries and assigning synthetic correctness labels so that the resulting task-wise accuracies match (in expectation) the target profile of \(m\). Feeding this synthetic set through the encoder produces a hypothetical embedding \(\tilde{z}_{m}\), which we then compare against the library of real model embeddings \(\{z_m\}_{m\in\mathcal{M}}\) via nearest-neighbor search.

We report recall@k: the fraction of models for which the true model embedding \(z_m\) appears among the top-\(k\) nearest real embeddings to \(\tilde{z}_{m}\). Averaged over all models, recall@10 is almost \(97\%\) for evaluation sets of size \(8192\) queries, depicting that model embeddings capture high level task performance (11). These results indicate that the embedding generator maps evaluation sets to representations that reflect higher-level behavioral profiles: even when correctness labels are randomized at the query level subject to matching task-level rates, the resulting embedding is often close to the intended model.

9.2 Resilient fallback routing under model unavailability.↩︎

In deployed query routing systems, the language model selected by a router may become temporarily unreachable (e.g., due to overload, failures, or policy constraints). When routing decisions are produced by a black-box router and per-query correctness probabilities / suitability scores for models are not readily available, a natural fallback is to substitute the router selected model with a nearby model in embedding space. To study the validity of the fallback routing approach, for each test query \(x\) we take the routed model \(\widehat{m}(x)\) and simulate its unavailability, then generate a response using its \(k^{\text{th}}\) nearest neighbor by embedding distance (closest, second-closest, etc.). [fig:fallback95routing] reports routing accuracy versus neighbor rank: the closest-neighbor fallback retains up to \(\mathbf{85\%}\) of the original routing accuracy, with performance decreasing smoothly as \(k\) increases. This supports embedding proximity as a practical mechanism for resilient fallback under model outages.

9.3 Convergence relative to inter-model separation, and implications to security fingerprinting.↩︎

To quantify and contextualize the model embedding drift w.r.t. varying the evaluation set overlap and size, we plot the distance to the originally generated reference embedding \(\|z_m - z_m^{\mathrm{ref}}\|\) and compare them with embedding distances of reference embedding of chosen model to reference embeddings of other distinct models. 12 shows that the new model embeddings generated are close to the reference embeddings (as compared to distances to other reference model embeddings), and the distance/dissimilarity decays as the overlap fraction (or the number of evaluation queries) increases. A similar trend is also observed for subsampling a smaller number of queries in the evaluation set for generating model embeddings, as presented in 13.

Relative robustness of generated model embeddings to the choice of queries demonstrates possible applications in security fingerprinting of language models, where very close embeddings between two models can indicate the same language model hidden behind different APIs or endpoints.

9.4 Model portfolio selection using embedding geometry↩︎

9.4.1 Setup and portfolio routing accuracy↩︎

Let \(\mathcal{M}=\{1,\dots,M\}\) denote the full model pool and let \(z_m\in\mathbb{R}^d\) be the learned embedding for model \(m\). Given a query \(x\), the router assigns a correctness probability \(\widehat{p}^{(m)}(x)\) to each model \(m\). For a selected portfolio \(\mathcal{M}_\textrm{sel}\subseteq\mathcal{M}\), routing chooses \[\widehat m_\textrm{sel}(x) \;=\; \mathop{\mathrm{arg\,max}}_{m\in \mathcal{M}_\textrm{sel}} \widehat{p}^{(m)}(x).\] Let \(y^{(m)}(x)\in\{0,1\}\) indicate whether model \(m\) answers \(x\) correctly (used only for evaluation). We define portfolio routing accuracy as \[\mathrm{Acc}(\mathcal{M}_\textrm{sel}) \;=\; \mathbb{E}_{x}\big[y^{(\widehat m_\textrm{sel})}(x)\big],\] which we estimate on held-out data as routing accuracy after restricting to the models in \(\mathcal{M}_\textrm{sel}\). Our goal is to select \(\mathcal{M}_\textrm{sel}\) using only embedding geometry \(\{z_m\}\) under two different deployment constraints:

1. a count constraint \(|\mathcal{M}_\textrm{sel}|\leq k\), with \(k\in\{1,2,3,\cdots\}\)

2. a parameter budget constraint \(\sum_{m\in \mathcal{M}_\textrm{sel}} c_m \le \mathcal{C}\), where \(c_m>0\) denotes total number of parameters in model \(m\).

9.4.1.1 Similarity between model embeddings through Kernels.

Let \(K(z_i,z_j)\in[0,1]\) be a similarity kernel in the model embedding space (e.g., an RBF kernel), where larger values indicate close model embeddings and hence similar model behaviors. Concretely, we define \[K(z_i,z_j) \;=\; \exp\!\Big(\! -\tfrac{d_z(z_i,z_j)^2}{\sigma^2}\Big), \label{eq:rbf95kernel}\tag{6}\] with \(d_z(\cdot,\cdot)\) being a distance/dissimilarity measure (e.g., cosine distance) and bandwidth \(\sigma>0\) (which we set as the median pairwise distance between model embeddings).

9.4.2 Model Count Constrained Selection↩︎

To choose a subset of models such that collective abilities of the pool is maintained, we select model embeddings spread across the embedding space via the kernelized \(k\)-center objective [42]: \[\max_{\mathcal{M}_\textrm{sel} \subseteq \mathcal{M}:\;|\mathcal{M}_\textrm{sel}|\le k}\;\; \min_{i\in \mathcal{M}}\;\max_{s\in \mathcal{M}_\textrm{sel}} K(z_i,z_s), \label{eq:kcenter95obj95kernel}\tag{7}\] which seeks to choose \(k\) model embeddings such that every model in the full pool has high similarity to at least one chosen model in the subset (i.e., the worst-covered model is as well-covered as possible).

We additionally study the kernelized \(k\)-medoids objective [43] which maximizes the sum similarity of all model embeddings to the most similar among the chosen subset of models closest : \[\max_{\mathcal{M}_\textrm{sel} \subseteq \mathcal{M}:\;|\mathcal{M}_\textrm{sel}|\le k}\;\; \sum_{i\in \mathcal{M}} \max_{s\in \mathcal{M}_\textrm{sel}} K(z_i,z_s), \label{eq:kmedoids95obj95kernel}\tag{8}\] which selects \(k\) embeddings so that the total coverage (best similarity per model, summed over the pool) is maximized.

Both aforementioned approaches depend only on model embedding geometry and do not utilize per-query correctness.

9.4.2.1 k-center (farthest-first in similarity space) [42].

We apply a similarity-space analogue of farthest-first sampling. We initialize \(\mathcal{M}_\textrm{sel}=\{m_0\}\) using the global kernel medoid \[m_0 \;=\; \mathop{\mathrm{arg\,max}}_{m\in\mathcal{M}}\;\frac{1}{|\mathcal{M}|}\sum_{i\in\mathcal{M}} K(z_i,z_m),\] then iteratively add the least-covered model (the one whose best similarity to the current set is smallest): \[m^\star \;=\; \mathop{\mathrm{arg\,min}}_{m\in \mathcal{M}\setminus \mathcal{M}_\textrm{sel}}\;\; \max_{s\in \mathcal{M}_\textrm{sel}} K(z_m,z_s), \qquad \mathcal{M}_\textrm{sel} \leftarrow \mathcal{M}_\textrm{sel} \cup \{m^\star\}. \label{eq:farthest95first95kernel}\tag{9}\] Intuitively, this repeatedly adds a representative for the region of embedding space that is poorly covered at each step (lowest similarity to models in the selected set \(\mathcal{M}_\textrm{sel}\)).

9.4.2.2 k-medoids [43].

Initializing with the k-center solution, we utilize the PAM (Partitioning Around Medoids) refinement algorithm which swaps model embeddings under sum-similarity objective and replaces embeddings which result in improvement of the objective.

Define the sum-similarity objective as \[\label{eq:sum95similarity95objective} J(\mathcal{M}_\textrm{sel}) \;=\; \sum_{i\in \mathcal{M}} \max_{s\in \mathcal{M}_\textrm{sel}} K(z_i,z_s),\tag{10}\] so that maximizing \(J\) is equivalent to maximizing the total similarity coverage in 8 . PAM searches for swaps \(s\in \mathcal{M}_\textrm{sel}\) and \(h\notin \mathcal{M}_\textrm{sel}\) that improve the objective value: \[\mathcal{M}_\textrm{sel}' \;=\; (\mathcal{M}_\textrm{sel}\setminus \{s\}) \cup \{h\}, \qquad \text{accept if } J(\mathcal{M}_\textrm{sel}') > J(\mathcal{M}_\textrm{sel}),\] and stops when no improving swap exists (or after a fixed number of iterations).

9.4.2.3 Random baseline.

As a baseline, we sample a subset of models \(\mathcal{M}_\textrm{sel}\subseteq \mathcal{M}\) of size \(k\) uniformly at random, and report mean \(\pm\) standard deviation in routing accuracy across trials.

9.4.2.4 Evaluation.

For each \(k\), we select \(\mathcal{M}_\textrm{sel}\) using only \(\{z_m\}\) (via the kernel \(K(\cdot,\cdot)\)) and evaluate \(\mathrm{Acc}(\mathcal{M}_\textrm{sel})\) by restricting routing accuracy to the chosen subset of models.

9.4.2.5 Observations.

For the count constraint model portfolio selection problem, we find that simple approaches like k-medoids and k-center (9) achieve full model pool (comprising \(112\) models) routing accuracy with just \(15-20\) models in the selected subset, signifying that selecting a subset of models by maximizing coverage over the model embedding space leads to capturing the collective abilities of the entire model pool. Such efficient selection of capable model subsets/portfolios points towards an efficient management of a large pool of language models, and eliminating redundancy in model abilities for simplifying deployment and improving hardware utilization. We reiterate that the presented approaches do not require any query-correctness prediction over models on representative queries, and rather the model subset choice is made purely based on model embedding space.

9.4.3 Parameter Budget Constraint Selection↩︎

We utilize the sum-similarity objective as presented above in 10 and solve the parameter-budget-constrained maximization for respecting total parameter budget for model subset choice \[\max_{\mathcal{M}_\mathrm{sel}\subseteq\mathcal{M}}\;\; J(\mathcal{M}_\mathrm{sel}) \quad \text{s.t.}\quad \sum_{m\in \mathcal{M}_\mathrm{sel}} c_m \le \mathcal{C}. \label{eq:knapsack95objective}\tag{11}\]

9.4.3.1 Greedy selection (marginal gain per parameter).

For simplicity, we study the greedy solution to the formulation in 11 . Let \(b_i=\max_{s\in \mathcal{M}_\mathrm{sel}}K(z_i,z_s)\) denote most similar selected model in \(\mathcal{M}_\textrm{sel}\) for model \(i\). For a candidate \(m\notin \mathcal{M}_\mathrm{sel}\), the marginal gain is \[\Delta(m\mid \mathcal{M}_\mathrm{sel}) \;=\; J(\mathcal{M}_\mathrm{sel}\cup\{m\})-J(\mathcal{M}_\mathrm{sel}) \;=\; \sum_{i\in\mathcal{M}} \max\big\{0,\;K(z_i,z_m)-b_i\big\}. \label{eq:marginal95gain}\tag{12}\] Under the remaining budget, we add the model with the largest gain per parameter: \[m^\star \;=\; \mathop{\mathrm{arg\,max}}_{\substack{m\notin \mathcal{M}_\mathrm{sel},\\\; c_m \leq~ (\mathcal{C}-\sum_{s\in \mathcal{M}_\mathrm{sel}} c_s)}} \frac{\Delta(m\mid \mathcal{M}_\mathrm{sel})}{c_m}, \qquad \mathcal{M}_\mathrm{sel} \leftarrow \mathcal{M}_\mathrm{sel} \cup \{m^\star\}, \label{eq:greedy95ratio}\tag{13}\] and stop when no feasible model remains.

9.4.3.2 Evaluation.

For each parameter count constraint \(\mathcal{C}\), we select \(\mathcal{M}_\mathrm{sel}^\mathcal{C}\) using 13 and evaluate \(\mathrm{Acc}(\mathcal{M}_\mathrm{sel}^\mathcal{C})\) by restricting evaluation to models present in \(\mathcal{M}_\mathrm{sel}^\mathcal{C}\). We also compare to a random feasible baseline that samples subsets satisfying \(\sum_{m\in S}c_m\leq \mathcal{C}\) and reports mean \(\pm\) standard deviation across trials.

9.4.3.3 Observations.

For the parameter constrained choice of models (9), we note that the described greedy algorithm which maximizes coverage over the space of model embeddings consistently outperforms random choice of models, attaining very close to full model pool routing accuracy at \(\approx\)​150B parameters (full model pool has 112 models, with 1930B parameters in total). For practical deployment scenarios with limited hardware, clever choice of subset of models achieves almost full model pool’s routing accuracy by eliminating duplicated model capabilities (in total just \(8\%\) parameter budget across all models in the pool).

10 Regenerating Model Embeddings for Backpropagation-based Approaches↩︎

In this section, we describe the shortcomings of learnable embeddings based approaches. For the sake of brevity, we consider EmbedLLM [22] and examine the geometry and capability representation of model embeddings learned by original training vs retraining/regenerating embeddings with the same data. Concretely, consider \(\{e^{(m)}_\theta\}_{m\in\mathcal{M}}\) as the trainable embeddings for each model in the pool, and let \(G_\psi\) denote the correctness predictor which is of the form \(G_\psi (e^{(m)}_\theta, x) = \mathrm{MLP}_\psi(e^{(m)}_\theta \odot W_\psi x)\), where \(x\) is the query encoding generated through a sentence encoder, \(W_\psi\) is the trained query downprojection matrix, \(\mathrm{MLP}_\psi\) is a multilayer perceptron network, and \(\odot\) denotes elementwise product. We perform training over all model-query evaluation pairs, to minimize the binary cross entropy loss between predictions and observed correctness values from evaluations: \[\label{eq:embedllm95objective} \min_{\theta,\psi} \quad \mathbb{E}_m\;\mathbb{E}_x\;\; \Big[\mathrm{BCE}\big(G_\psi (e^{(m)}_\theta, x), \;y^{(m)}(x)\big)\Big]\tag{14}\]

This results in learned model embeddings \(\{e^{(m)}_\theta\}_{m\in\mathcal{M}}\), and correctness predictor \(G_\psi\). To evaluate the stability of the generated model embeddings, we freeze the correctness predictor parameters \(G_\psi\), and then retrain the model embeddings to obtain \(\{e^{(m)}_{\theta_\mathrm{re}}\}_{m\in\mathcal{M}}\) with the same model query evaluation data. Mathematically, we perform: \[\label{eq:relearn95embedllm95objective} \min_{\theta_\mathrm{re}} \quad \mathbb{E}_m\;\mathbb{E}_x\;\; \Big[\mathrm{BCE}\big(G_\psi (e^{(m)}_{\theta_\mathrm{re}}, x), \;y^{(m)}(x)\big)\Big]\tag{15}\]

This represents a practical scenario where unseen models are added to the pool, and embeddings need to be generated for them without affecting correctness predictor \(G_\psi\) and already learned embeddings \(\{e^{(m)}_\theta\}_{m\in\mathcal{M}}\) for models in the pool. For the sake of analysis, we consider new added models as clones of existing models with identical query evaluation data, and hence compare their newly generated embeddings \(\{e^{(m)}_{\theta_\mathrm{re}}\}_{m\in\mathcal{M}}\) with originally trained embeddings \(\{e^{(m)}_{\theta}\}_{m\in\mathcal{M}}\) for similarity.

Because the training objective in 15 differs from that in 14 , we anticipate small changes in the regenerated model embeddings, even though the embeddings are learned on the same query evaluation data. However, they should remain close to the original embeddings to preserve the practically desirable properties described in 1. If the regenerated embeddings \(\{e^{(m)}_{\theta_\mathrm{re}}\}_{m\in\mathcal{M}}\) are very different as compared to original model embeddings \(\{e^{(m)}_{\theta}\}_{m\in\mathcal{M}}\), we cannot utilize informative embedding geometry. In such scenarios, we cannot mutually compare model embeddings across models to identify similarities or differences in capabilities, making most of the analyses and practical applications presented in 3 infeasible.

We examine EmbedLLM [22] through a similar approach as described in 3.3, where we regenerate model embeddings through new query evaluation sets to test sensitivity to which queries are used for generating embeddings and how many queries are used to generate embeddings. 14 denotes geometric representation of original and regenerated embeddings (cf. 3.3), highlighting that regenerated embeddings are not in close proximity of the original embeddings. This is also presented numerically in 15, where average cosine distance between respective original and regenerated embeddings is high, even though regenerated embeddings are trained on the same evaluation data as original embeddings.

Discrepancy between original and regenerated embeddings also affects correctness prediction over queries. [fig:relearn95embedllm95heatmap95test95disagreement] denotes average correctness prediction disagreement (over a common set of test queries) by using original versus regenerated embeddings with the same correctness predictor \(G_\psi\). Again, regenerated embeddings are trained on the original evaluation set, but still display up to \(8.1\%\) disagreement rate on predicting correctness values for test queries.

11 Time Complexity of LOCUS for Generating Embeddings and Correctness Prediction↩︎

In this section, we examine the time complexity of LOCUS for generating embeddings using query evaluations, and for correctness prediction on unseen queries. We utilize a single NVIDIA Tesla V100 (32GB) GPU [60] for our experiments.

16 (a) denotes the wall clock time required to generate model embeddings (for \(|\mathcal{M}|\) models in parallel) using query evaluation sets of varying sizes. We observe that even for parallel embedding generation for the full model pool (112 models) on 4096 queries, required time is \(\approx\)​100ms. This indicates powerful yet efficient design of LOCUS for generating model embeddings from observed query evaluations. Additionally, 16 (b) denotes the correctness prediction time for varying number of unseen queries using the decoder \(G_\psi\). Calculating correctness probabilities over the entire model pool for 4096 queries requires \(\approx\)​20ms, which is orders of magnitude smaller than actual response generation time from language models (usually in seconds [61]). This further strengthens the practicality of LOCUS for applications like routing, model ranking etc.

12 Implementation Details and Reproducibility Checklist↩︎

We have provided training code and data processing scripts for LOCUS in the attached link https://github.com/patel-shivam/locus_code_release. Across experiments in this work, LOCUS utilizes \(r=64\) learnable vectors for Latent bottleneck Attention Blocks (3 ), and our default model embedding dimension is \(d=128\). We choose query, key and value dimensions same as \(d\) throughout all the attention layers. The feedforward network (\(\mathrm{FFN}\)) is a two-layer MLP with hidden dimension \(d_\mathrm{ff} = 2d\). The tokenizer module \(h_\omega\) is a single layer MLP with output dimension as \(d\). Correctness predictor \(G_\psi\) is a 2-layer MLP with hidden dimension 64.

For training encoder-decoder pair (\(F_\theta,G_\psi\)) in 2.3, we sample the query evaluations as input to the encoder \(S_m^\mathrm{enc}\subseteq S_m\) to generate model embeddings \(z_m\), and independently sample a subset \(S_m^\mathrm{dec}\subseteq S_m\) for training the correctness predictor (\(S_m^\mathrm{enc}\) and \(S_m^\mathrm{dec}\) can have overlap, but only during training). The number of queries chosen in \((S_m^\textrm{enc}, S_m^\textrm{dec})\) can be varied to adjust to training hardware, our choices are detailed in the supplementary code submission. For all of the presented numerical experiments on correctness prediction and routing accuracy, test queries are only used for evaluating the decoder performance \(G_\psi\), and never used for generating the embeddings \(\{z_m\}_{m\in\mathcal{M}}\). This ensures that there is no leakage of test data into generating the model embeddings. Additionally, experiments examining correctness agreement rate utilize test set of common queries over which model correctness agreement is evaluated (6 8), again ensuring no test queries used for generating embeddings and no training queries used for evaluation.

References↩︎

---

1. Presently, over \(300\text{K}\) text-to-text models are open-sourced on huggingface.co/models↩︎

2. trained on \(4096\) query evaluations for each language model, unless mentioned otherwise.↩︎