Embedding Geometry Shapes LLM Routing: Tracing Failure Modes and Improving Generalization
Abstract
With the growing number of LLMs available, selecting which one to use for a prompt has become a key question. A large part of the routing literature relies on semantically embedding prompts into a vector space and then applying a lightweight classifier to the embedded representations. These methods perform robustly in-distribution but they struggle to generalize to unseen tasks. By studying the geometry of semantic embeddings, we trace a clear failure mode: features that determine the neighbors of unseen prompts are an unreliable proxy for LLM selection. We therefore replace the semantic embedding with the following map: we run a partial forward pass of Qwen2.5-7B-Instruct on the prompt and use the hidden state of its last token at a deep layer as the representation. When applying kNN and other simple classifiers to this new representation, we obtain a clear result in generalization: on average across cost penalties, their performance improves. Moreover, kNN and another simple classifier that we propose outperform every evaluated semantic-embedding baseline. Our two simple methods also beat a recent SoTA hidden-states based router on RouterBench.
Related
- Policy Regret for Embedding Model Routing: Contextual Bandits with Low-Rank Experts
- Generalization properties of the token perceptron: the role of embedding and vocabulary
- Overfitting of Spectral Gradient Descent: How Matrix Geometry shapes Generalization and Implicit Bias
- From Geometry to Generalization: Why Row Normalization Can Beat Adam and Muon
- When Does Sparse Routing Help? The Geometry of Expert Separation in Mixture of Experts
BibTeX
@misc{torti2026embedding,
title = {{Embedding Geometry Shapes LLM Routing: Tracing Failure Modes and Improving Generalization}},
author = {Leonardo Torti and Emanuele Rimoldi and Tomer Galanti and Tomaso Poggio and Pierfrancesco Beneventano},
howpublished = {arXiv},
year = {2026},
month = oct,
doi = {10.48550/arxiv.2610.19198},
eprint = {2610.19198},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2610.19198},
note = {CC BY 4.0},
}
RIS
TY - UNPB
TI - Embedding Geometry Shapes LLM Routing: Tracing Failure Modes and Improving Generalization
AU - Leonardo Torti
AU - Emanuele Rimoldi
AU - Tomer Galanti
AU - Tomaso Poggio
AU - Pierfrancesco Beneventano
PY - 2026
DA - 2026/10/13
PB - arXiv
DO - 10.48550/arxiv.2610.19198
UR - https://arxiv.org/abs/2610.19198
AB - With the growing number of LLMs available, selecting which one to use for a prompt has become a key question. A large part of the routing literature relies on semantically embedding prompts into a vector space and then applying a lightweight classifier to the embedded representations. These methods perform robustly in-distribution but they struggle to generalize to unseen tasks. By studying the geometry of semantic embeddings, we trace a clear failure mode: features that determine the neighbors of unseen prompts are an unreliable proxy for LLM selection. We therefore replace the semantic embedding with the following map: we run a partial forward pass of Qwen2.5-7B-Instruct on the prompt and use the hidden state of its last token at a deep layer as the representation. When applying kNN and other simple classifiers to this new representation, we obtain a clear result in generalization: on average across cost penalties, their performance improves. Moreover, kNN and another simple classifier that we propose outperform every evaluated semantic-embedding baseline. Our two simple methods also beat a recent SoTA hidden-states based router on RouterBench.
ER -