EmbedLLM: Learning Compact Representations of Large Language Models

We thank the reviewer for their comments and for taking the time to review our paper. We also appreciate their comments about finding our work innovative. We respond to some of the reviewer’s comments below:

W1: “With only 30,000 data for training, can the resulting encoding truly capture all the characteristics of the models?”

We would like to clarify that our current embedder certainly does not capture ALL characteristics of the models. Intuitively, if the model behavior data has zero information about how models perform in complex reasoning tasks, then we argue that no method could effectively predict model behaviors on them. However, we would like to emphasize that the key contribution of our work is that our framework could effectively capture and compress whatever information that is available in our dataset in a compact embedding form, and our experiment in correctness forecasting as well as model routing suggests that the model embeddings are powerful summarizers of model strengths and weaknesses within the range of model behaviors that is tested in our selected benchmarks. We believe that our framework could also extend on complex reasoning tasks if we're able to at least reconstruct some of the relevant tasks.

W2: “Why are two values (p(m,q)_0 and p(m,q)_1) output here? Is it simply to add a nonlinear operation?”

The decoder can be seen as a binary classifier so the two outputs are the two logits for classifying to label 0 or 1, and we combine them into one logit by taking their difference.

W3.1: “The experiments do not appear to categorize models by scale.”

In fact in figure 6 we performed some analysis on scale that indicates our embedding intrinsically captures the scale of the model (which is likely to be represented by the overall ability across all benchmarks). Specifically, we demonstrated that for all models of size 7B, 13B, or 70B, their embeddings’ intra-community distance is smaller than inter-community distance, indicating a presence of clusters, even though the data doesn’t explicitly contain the size information of the models.

W3.2: “Intuitively, the larger the model, the smaller the performance differences between models. What is the authors' view on this issue?”

We respectfully disagree with the statement that the performance gap between larger models is smaller. While many large models share overlapping training data, recent developments focus on specialization through domain-specific fine-tuning, or utilizing various feedback like program compiler for code generation or theorem verifiers for math. This led to remarkable divergences of performance even among similarly scaled models. We believe that larger models are exhibiting more and more distinctive strengths and weaknesses specific to certain tasks rather than converging in performance, and so our method could be a sustainable solution in this future we envision.

W4: “The paper does not seem to provide a clear and intuitive illustration of the overall architecture, including training, input, and output processes.”

We thank the reviewer for this suggestion. We will refine the illustrations and add a figure that explains our architecture explicitly.

Finally, we would like to thank the reviewer once more for their comments which we find to be really constructive. If the reviewer finds our responses above to answer their questions, we gently request the reviewer to take them into account before finalizing their reviews and revise their scores as they find appropriate.

Q4: “Some qualitative discussion with examples on what happens when the resulting data points are widely different from already seen data points will be useful.”

Some of our probing experiments serve for that purpose. For instance, in figure 5, a linear regressor put on top of our embedding, which is trained from all 10 benchmarks except SocialQA, does not exhibit statistical significance in predicting SocialQA performance.

We would also like to clarify that our current embedder certainly does not capture ALL characteristics of the models. Intuitively, if the model behavior data has zero information about how models perform in complex reasoning tasks, then we argue that no method could effectively predict model behaviors on them. However, we would like to emphasize that the key contribution of our work is that our framework could effectively capture and compress whatever information that is available in our dataset in a compact embedding form, and our experiment in correctness forecasting as well as model routing suggests that the model embeddings are powerful summarizers of model strengths and weaknesses within the range of model behaviors that is tested in our selected benchmarks.

Q5: “Does the prompt embedding have an impact on the results, i.e., using better recent LLM-based embeddings.”

Specifically on our dataset, the choice of the embedder that produces the prompt embedding does not affect the result to a large extent. We provide some experimental detail below:

For the experiments in the paper, we chose “sentence-transformers/all-mpnet-base-v2” as our embedder and on the full training set training from embeddings from this embedder gives 74% test accuracy. We further explored 4 other embedders from the MTEB leaderboard on HuggingFace: “jxm/cde-small-v1”, “nomic-ai/nomic-embed-text-v1-ablated”, “sentence-transformers/msmarco-bert-co-condensor”, and “sentence-transformers/LaBSE”. Using the same set of hyperparameters, we found that all of the training using these embedders gives a test accuracy ranging from 72% to 74%.

Intuitively, since our questions are more like knowledge tests, we expect the embedder to learn the differences between question embeddings in the form like “which field is the question belonging to”. Nevertheless, we agree that a better embedder could result in better performance when the model behavior data become more intricate or when the label of the data becomes related to the actual model response.

Q6: “I would also like to see the TSNE plots of the embeddings.”

As mentioned in our limitation section, because of the limited number of model (112), the embeddings are often quite sparse and therefore dimension reduction techniques did not apply very well on the embeddings. However, we would like to reassure the reviewer that we conducted extensive probing experiments (Section 6) to ensure that the embeddings capture meaningful information. Specifically, the performance of our linear prober placed on top of the embeddings are highly affected by changes in the training set and their changes align with the qualitative natures of the benchmark we selected. We hope that these probing experiments could serve as an alternative to tSNE plots.

On the Minor issues: We appreciate the reviewer for pointing out these details and we will certainly refine our writing in an updated version.

Finally, we would like to thank the reviewer once more for their comments which we find to be really constructive. If the reviewer finds our responses above to answer their questions, we gently request the reviewer to take them into account before finalizing their reviews and revise their scores as they find appropriate.

We thank the reviewer for their comments and for taking the time to review our paper. We also appreciate their comments about finding our work intuitive.

We respond to some of the reviewer’s comments below

Q1: “Why does the method assume that answer have to be a binary label? Is this a limitation of this method?”

Our method does not need the label to be binary only. Our emphasis is to convey that we can employ such encoder-decoder architecture to learn compact representations of various LLMs that facilitate downstream tasks, given some data of model behaviors. We choose a dataset of binary correctness labels because it is the most accessible feature that quantifies model behavior. Nonetheless, our framework is certainly generalizable to learning from other formats of data, for instance we could learn to reconstruct exact model responses to questions (of course in this setting the encoder and decoder needs to be more carefully designed and the scale of the training is larger).

Q2: “The explanation of kNN-classifer baseline is unclear. How is final label chosen? -- based on the majority voting of k-nearest neighbours?”

Yes, the KNN classifier aggregates the label by getting the majority vote of the K nearest neighbors where we performed hyperparameter tuning on the optimal number of neighbors.

Q3: “What is mxp and nxp represent?”

The correctness matrix has shape (n,p) where n is the number of models and p is the number of prompts(questions). As a result, m is the dimension of our model embeddings which is itself a hyperparameter and thus can be flexible. We thank the reviewer for pointing this out and we will update our writing with an additional explanation on what these symbols mean.

Thank you for your response! I still feel some important questions are unanswered.

Q6: I would also like to see the TSNE plots of the embeddings.

I am puzzled by the response. Why is the resulting model embedding sparse? In Q3, you indicate "m" is quite flexible, and if so, is it not possible to build a dense embedding. Can you point me to the discussion in the paper where you describe what values of 'm' are used. But TSNE plots are based on distances, so I don't see sparsity as a problem, or is it? As I pointed in weakness 5, Figure 6 isn't that useful to give a wholistic view of model similarities.

Q5: “Does the prompt embedding have an impact on the results, i.e., using better recent LLM-based embeddings.”

Here I meant LLM-based embeddings not MLM based embeddings. For example, LLM based embeddings like LLM2Vec or GritLM style. In your explanation, you imply that question category (domain) is more important than the difficulty of the question. That is intriguing because I thought the main usefulness of the work is at estimating whether a model can solve a given question but not how good it is in a domain? Isn't it?

Could you also comment on weakness 3 and 4?

Thanks for the further discussion! Here are responses to some of your questions:

1. About TSNE plots, upon further review, we agree that sparsity may not inherently preclude TSNE from being used effectively. In our case, TSNE plots of our embeddings were not sparse; however, the resulting plots do not reveal clear trends or clusters. The points appear more scattered. We acknowledge this as a current limitation. To improve this aspect, we plan to explore approaches such as training with a subset of more specialized models (e.g., only include models specialized in math and in medicine) to potentially create more distinguishable clusters. Additionally, adding more models to the dataset could address sparsity concerns and allow for clearer separations. We appreciate this suggestion and will address it in the final version of the paper to enhance the interpretability of our embeddings.

2. About LLM-based embedding, we appreciate this clarification and we will conduct more experiments on the effect of the embedder. We believe that more informative embedding could definitely enhance performances on downstream tasks and we will consider this enhancement for future work.

About question category vs. difficulty, we believe that the behavior of our routing mechanism suggests that the embeddings are capturing characteristics aligned with question domains rather than question difficulty. Specifically, our MF router utilizes a diverse set of models for different questions, which indicates that the embeddings capture nuances about which models are better suited for specific domains. This behavior aligns with our classifier design—a linear classifier applied to the Hadamard product of the question and model embeddings—which encourages an alignment between model and question domains. If the embeddings were purely reflective of difficulty, we would expect the router to favor a smaller subset of strong models across all questions, which is not observed in our experiments. We will conduct more controlled experiments to validate this claim.

3. For weakness 3, we will incorporate references to relevant prior work, including Xia et al. (2020), in the final version of the paper. One thing worth mentioning is that while related works have focused on predicting performance on new tasks based on existing tasks, our primary focus is on compactly vectorizing model characteristics. Performance prediction serves as an auxiliary task within our broader framework, which is designed to adapt to any measurable model behavior data.

4. About weakness 4, we thank the reviewers for pointing out this possible confusion. In the specific setting of our binary correctness dataset, the decoder is indeed a binary classifier, and the reason why we call it a decoder is because it is used to reconstruct every entry of our correctness matrix. We will place more emphasis in explaining this in the updated writing to make it less ambiguous. Nonetheless, we re-emphasize that our framework is generalizable to learning from other formats of data, for instance we could learn to reconstruct exact model responses to questions and in that setting the decoder would be reconstructing text as the reviewer suggested. In this work, the main focus is to convey that we can employ such encoder-decoder architecture to learn compact representations of various LLMs that facilitate downstream tasks, given some data of model behaviors.

Overall, we thank the reviewer again for their constructive and detailed comments. These insights have been invaluable in identifying areas for improvement, and we are committed to addressing them comprehensively in the final version.

We thank the reviewer for their comments and for taking the time to review our paper. We also appreciate their comments about finding our work novel and effective. We respond to some of the reviewer’s comments below:

W1: “The term "decoder" in this paper is a bit misleading.”

We thank the reviewers for pointing out this possible confusion. In the specific setting of our binary correctness dataset, the decoder is indeed a binary classifier, and the reason why we call it a decoder is because it is used to reconstruct every entry of our correctness matrix. We will place more emphasis in explaining this in the updated writing to make it less ambiguous.

Nonetheless, we re-emphasize that our framework is generalizable to learning from other formats of data, for instance we could learn to reconstruct exact model responses to questions and in that setting the decoder would be reconstructing text as the reviewer suggested. In this work, the main focus is to convey that we can employ such encoder-decoder architecture to learn compact representations of various LLMs that facilitate downstream tasks, given some data of model behaviors.

W2: “I'd like to see details cost metrics for re-train embedder for new model vs traditional benchmarking approaches.”

About the retraining cost, we would like to reassure the reviewer that training our embedder takes a trivial amount of computation compared to traditional benchmarking approaches. Below is some experimental details:

Using fvcore package to perform flop analysis, we performed an example on a training set leaving out all questions from MMLU, which results in a correctness matrix of around 20,000 questions on 112 models. Training our embedder on this dataset for 50 epochs using a batch size of 2048 results in peak memory usage of less than 1GB and costs 107.71 TFlops in total. This is approximately equivalent to the cost of running 60 inferences of length 128 on a 7B model (calculation referring https://arxiv.org/pdf/2001.08361). That is, benchmarking all 112 models on MMLU test set (about 14,000 questions) will be at least (as we have larger models) 112 * 14000 / 60 = 26000 times more expensive. Hence, we believe that our framework could effectively leverage the extensive demand on performing traditional benchmarking.

Finally, we would like to thank the reviewer once more for their comments which we find to be really constructive. If the reviewer finds our responses above to answer their questions, we gently request the reviewer to take them into account before finalizing their reviews and revise their scores as they find appropriate.