English

Sub-Token Routing in LoRA for Adaptation and Query-Aware KV Compression

Machine Learning 2026-05-07 v2 Computation and Language

Abstract

Sub-token routing provides a finer compression axis for transformer efficiency than the coarse units used in most prior work, such as tokens, pages, heads, or layers. In this paper, we study routing within a token representation itself in LoRA-adapted transformers. We consider two settings. In the query-independent setting, we combine routed subspace LoRA with value-group routing on the KV path for compression-aware language modeling. In the query-aware setting, we use a predictor-based selector to allocate a global retention budget over context-token/value-group pairs using query-conditioned relevance. Experiments show that the query-independent design improves language-model quality under reduced KV budgets, while the query-aware design preserves downstream behavior well under KV compression. We further show that sub-token routing is most effective as a complementary compression axis to token-level query-aware selection: token-level methods decide which tokens survive globally, while sub-token routing determines how the surviving tokens are compressed internally. Their combination enables deeper KV compression at nearly unchanged task accuracy.

Keywords

Cite

@article{arxiv.2604.21335,
  title  = {Sub-Token Routing in LoRA for Adaptation and Query-Aware KV Compression},
  author = {Wei Jiang and Wei Wang},
  journal= {arXiv preprint arXiv:2604.21335},
  year   = {2026}
}

Comments

17 pages, 13 tables, 2 figures

R2 v1 2026-07-01T12:31:57.306Z