TriRoute: Unified Learned Routing for Joint Adaptive Attention, Experts, and KV-Cache Allocation

• URL: http://arxiv.org/abs/2607.06601v1
• Date: Tue, 07 Jul 2026 00:12:46 GMT
• Title: TriRoute: Unified Learned Routing for Joint Adaptive Attention, Experts, and KV-Cache Allocation
• Abstract summary: TriRoute is a single lightweight controller shared across all three axes.<n>It trains end-to-end via a heterogeneous relaxation.<n>It allocates full attention and high-precision cache to sentence-initial positions, rare subwords, and named entities.
• Score: 0.0
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: Conditional computation can decouple language model quality from per-token inference cost, yet leading techniques act on a single axis in isolation: Mixture-of-Experts (MoE) sparsifies the FFN, Mixture-of-Depths (MoD) skips whole transformer blocks, and KV-cache quantization compresses attention memory. We argue these three decisions (attention resolution, expert selection, and cache bit-width) are strongly coupled and should be made jointly: a token rare enough to warrant full attention may also need high-precision caching regardless of which expert processes it. We introduce TriRoute, a single lightweight controller shared across all three axes that, for every token at every layer, emits a coordinated policy: (i) an attention mode (skip/local/full), (ii) a sparse set of FFN experts (with a null expert recovering MoD), and (iii) a KV-cache bit-width. The controller trains end-to-end via a heterogeneous relaxation (Gumbel-Softmax with straight-through estimation for categorical decisions and load-balanced top-k gating for experts) under a Lagrangian budget constraint that turns the average compute and memory cost into a controllable knob. We identify a cross-axis routing-collapse cascade in naive joint training, where collapse on one axis propagates to the others, and address it with per-axis normalization and a coupling-aware balancing loss. On decoder-only models from 160M to 1.3B parameters at compute-optimal token counts, TriRoute Pareto-dominates the best independent MoD+MoE+KV-quantization combination at matched inference FLOPs and memory, while better preserving tail-case robustness on rare entities, code, and arithmetic that pure perplexity optimization erodes. Post-hoc analysis reveals interpretable structure: the controller allocates full attention and high-precision cache to sentence-initial positions, rare subwords, and named entities, while cheaply routing function words.

Related papers

• Empowering Compact LLMs with Fusion of Layer-wise Exits for Recommendation [59.216721553953]
  Large language model-based recommender systems (LLM-RSs) have demonstrated remarkable capabilities, but are computationally unsustainable for many real-world applications.<n>We propose a discriminative framework that enhances compact LLMs while retaining scalable full-corpus ranking.<n>Experiments on three real-world datasets with Qwen 3 1.7B and Llama 3.2 3B show that FLEXRec achieves state-of-the-art accuracy while remaining highly efficient.
  arXiv  Detail & Related papers  (2026-08-18T03:14:28Z)
• Omni2LoRA: Coherence-Preserving Parametric Memory for Efficient Omni Language Models [65.49950267695267]
  We introduce Omni2LoRA, a framework for efficient parametric memory compression via coherence-preserving context distillation.<n>By converting multimodal memory into a fixed-budget, reusable parameter state, our method drives answer-time multimodal-token load to zero.
  arXiv  Detail & Related papers  (2026-08-10T07:49:10Z)
• Controller-Augmented Hidden Markov Models: A Computational Framework for Constrained Sequential Inference [0.0]
  We present Controller-Augmented Hidden Markov Models (CHMMs)<n>CHMMs is a framework that compiles each constraint into a finite-state controller tracking the minimal sufficient history.<n>We evaluate CHMMs against 6 alternative decoders on 3 real-world sequence-labeling tasks.
  arXiv  Detail & Related papers  (2026-06-11T19:29:09Z)
• SHAPE: Coalition-Aware Expert Pruning for Sparse Mixture-of-Experts LLMs [4.035485260002141]
  Sparse Mixture-of-Experts (MoE) models achieve strong quality with low per-token compute, yet their deployment is often limited by the memory wall.<n>We propose textbfSHAPE, a task-driven pruning framework that explicitly models emphintra-layer expert cooperation.
  arXiv  Detail & Related papers  (2026-06-03T08:41:48Z)
• LimiX-2M: Mitigating Low-Rank Collapse and Attention Bottlenecks in Tabular Foundation Models [56.999481798138625]
  LimiX-2M is a 2M- parameter model that outperforms larger TabPFN-v2 and TabICL baselines on widely used benchmarks.<n>We present a unified tokenize-and-route framework for strong Tabular foundation models (TFMs)<n>Results highlight value-aware tokenization and readout-aligned routing as key levers for improving the accuracy--efficiency trade-off in TFMs.
  arXiv  Detail & Related papers  (2026-06-03T06:07:33Z)
• When Model Merging Breaks Routing: Training-Free Calibration for MoE [67.27839627141972]
  Hessian-Aware Router (HARC) is a training-free framework that leverages second-order curvature information to realign the merged router.<n> Experiments on mathematical reasoning and code generation tasks show that HARC effectively mitigates routing breakdown across diverse MoE merging baselines.
  arXiv  Detail & Related papers  (2026-06-02T09:33:33Z)
• OScaR: The Occam's Razor for Extreme KV Cache Quantization in LLMs and Beyond [50.440302567029654]
  Multi-modal intelligence has made the Key-Value cache a dominant memory bottleneck for efficient deployment.<n>In this work, we revisit the inherent limitations of the per-channel quantization paradigm.<n>We propose OScaR, an accurate and lightweight KV cache compression framework for X-LLMs.
  arXiv  Detail & Related papers  (2026-05-19T10:53:03Z)
• Q-LocalAdam: Memory-Efficient Client-Side Adaptive Optimization for Edge Federated Learning [1.4610038284393168]
  Federated learning on edge devices must cope with non-IID client data and tight memory budgets.<n>We propose textbfQ-LocalAdam, which applies distribution-aware 8-bit quantization block-wise linear encoding for momentum and log-space encoding for variance.
  arXiv  Detail & Related papers  (2026-05-17T17:23:23Z)
• TriAxialKV: Toward Extreme Low-Precision KV-Cache Quantization for Agentic Inference Tasks [20.031258206130154]
  We introduce TriAxialKV, a novel mixed-precision KV-cache quantization scheme for agentic workloads.<n>We implement TriAxialKV as an end-to-end serving system, comprising calibration, mixed-precision quantization and memory management.
  arXiv  Detail & Related papers  (2026-05-16T21:58:28Z)
• AdaFuse: Accelerating Dynamic Adapter Inference via Token-Level Pre-Gating and Fused Kernel Optimization [84.25316984309725]
  Integration of dynamic, sparse structures with parameter-efficient adapters (e.g., LoRA) is a powerful technique for enhancing Large Language Models (LLMs)<n>Despite minimal increases in computational load, the computation latency often skyrockets, leading to decoding speeds slowing by over 2.5 times.<n>AdaFuse is a framework built on a tight co-design between the algorithm and the underlying hardware system to enable efficient dynamic adapter execution.
  arXiv  Detail & Related papers  (2026-03-12T12:46:42Z)
• Towards Zero-Shot Point Cloud Registration Across Diverse Scales, Scenes, and Sensor Setups [41.033896572767496]
  We present BUFFER-X, a training-free registration framework that achieves zero-shot generalization.<n>Our approach employs hierarchical multi-scale matching to extract correspondences across local, middle, and global receptive fields.<n>For efficiency-critical applications, we introduce BUFFER-X-Lite, which reduces total time by 43%.
  arXiv  Detail & Related papers  (2026-01-06T06:51:24Z)
• OmniSparse: Training-Aware Fine-Grained Sparse Attention for Long-Video MLLMs [43.78743496579736]
  We introduce OmniSparse, a training-aware fine-grained sparse attention framework for long-video MLLMs.<n>Experiment results show that OmniSparse matches the performance of full attention while achieving up to 2.7x speedup during prefill and 2.4x memory reduction during decoding.
  arXiv  Detail & Related papers  (2025-11-15T13:14:17Z)
• Unsupervised Conformal Inference: Bootstrapping and Alignment to Control LLM Uncertainty [49.19257648205146]
  We propose an unsupervised conformal inference framework for generation.<n>Our gates achieve close-to-nominal coverage and provide tighter, more stable thresholds than split UCP.<n>The result is a label-free, API-compatible gate for test-time filtering.
  arXiv  Detail & Related papers  (2025-09-26T23:40:47Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.