Research map / Large language models / Language, safety and interpretability
EmbedLLM: Learning Compact Representations of Large Language Models
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Learns compact vector embeddings of 112 open LLMs via a matrix-factorization-style encoder-decoder trained on binary correctness labels, then uses them for correctness forecasting, model routing, and benchmark accuracy prediction.
Problem
Evaluating and routing among the growing number of LLMs requires repeatedly re-learning task-specific model representations, which is time- and compute-expensive.
Approach
Treat each LLM as a learnable embedding indexed by model ID and reconstruct a model-by-question correctness matrix with a Hadamard-product decoder and linear classifier (framed as matrix factorization), so the resulting compact model embeddings can be reused across downstream tasks with only a light linear head.
Key results
- On the held-out question split, matrix factorization reaches 0.7409 correctness-forecasting accuracy with the full training set versus 0.7152 for KNN.
- The MF router achieves accuracy near the single-best model on each benchmark and beats the random weighted router overall, routing 3,000 questions in about 3.8 seconds on one A100.
- Benchmark accuracy prediction from embeddings trained without the target benchmark is statistically significant (Kendall's Tau, 5% level) on 7 of 10 benchmarks.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use Matrix Factorization instead of KNN | Correctness forecasting · Unseen questions across training set sizes from 1K to Full (29K) | accuracy: 0.6372 -> 0.6443; 0.7078 -> 0.7331; 0.7107 -> 0.7362; 0.7128 -> 0.7378; 0.7143 -> 0.7390; 0.7146 -> 0.7394; 0.7152 -> 0. | Helped |
| Increase the correctness-training dataset size | Matrix Factorization · Correctness forecasting on unseen questions | accuracy: 0.6443 -> 0.7409 | Helped |
| Use the Matrix Factorization router instead of single-best or weighted-random routing | Model routing · Overall test set | router accuracy | Helped |
| Use the Matrix Factorization router rather than a causal LLM router | Model routing · Model selection throughput | model selections per second: 15x faster | Helped |
| Use the Matrix Factorization router rather than fine-tuning Llama-3-8B as a router | Model routing training · GPU memory usage | memory cost: 60x cheaper | Helped |
| Add MMLU to the embedding training benchmarks (ablation) | Matrix Factorization embeddings · Accuracy prediction on other benchmarks | MSE | Helped |
| Add other benchmarks to the embedding training data (ablation) | Matrix Factorization embeddings · GPQA accuracy prediction | predictive power | Hurt |
| Add GSM8K to the embedding training data (ablation) | Matrix Factorization embeddings · MathQA accuracy prediction | MSE: 0.271 | Helped |
| Add MathQA to the embedding training data (ablation) | Matrix Factorization embeddings · ASDiv accuracy prediction | MSE: 0.190 | Helped |
| Add ASDiv to the embedding training data (ablation) | Matrix Factorization embeddings · MathQA accuracy prediction | MSE: 0.103 | Helped |
Surprising findings
- Incorporating other benchmarks into embedding training harms predictive power on GPQA.
- Adding MMLU to embedding training noticeably improves prediction on other benchmarks.
Limitations
- Dataset is small (112 models, correctness only), making embeddings moderately sparse.
- Adding new models requires retraining the matrix factorization model.
- Restricted to binary correctness data; other signals such as output text embeddings are unexplored.
Experimental setup
Datasets: MMLU TruthfulQA SocialQA PIQA MedMCQA MathQA LogiQA GSM8K GPQA ASDiv
Benchmarks: MMLU TruthfulQA SocialQA PIQA MedMCQA MathQA LogiQA GSM8K GPQA ASDiv
Metrics: prediction accuracy router accuracy MSE Kendall's Tau significance L2 distance latency
Baselines: KNN single-best model router random weighted router causal LLM router
Builds on: matrix factorization KNN classifier all-mpnet-base-v2 lm-evaluation-harness
Closest papers by meaning
- Scaling Embedding Layers in Language Models (NeurIPS 2025 · similarity 0.90)
- Repurposing Language Models into Embedding Models: Finding the Compute-Optimal Recipe (NeurIPS 2024 · similarity 0.89)
- NV-Embed: Improved Techniques for Training LLMs as Generalist Embedding Models (ICLR 2025 · similarity 0.89)
- SparseLLM: Towards Global Pruning of Pre-trained Language Models (NeurIPS 2024 · similarity 0.89)
- MaskLLM: Learnable Semi-Structured Sparsity for Large Language Models (NeurIPS 2024 · similarity 0.89)
- DEPT: Decoupled Embeddings for Pre-training Language Models (ICLR 2025 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.