Meta-Attention: Bayesian Per-Token Routing for Efficient Transformer Inference

• URL: http://arxiv.org/abs/2605.28384v1
• Date: Wed, 27 May 2026 12:21:28 GMT
• Title: Meta-Attention: Bayesian Per-Token Routing for Efficient Transformer Inference
• Abstract summary: We propose Meta-Attention, a framework that dynamically routes each token to the most appropriate attention strategy.<n>Unlike prior routing approaches that use deterministic or prior-free learned routing, the Meta-Controller treats per-token mechanism selection as posterior inference.<n>This design produces principled routing uncertainty estimates that govern the soft-to-hard routing transition, mitigates routing collapse without ad hoc load-balancing losses, and yields better compute-performance trade-offs.
• Score: 0.0
• License: http://creativecommons.org/licenses/by-sa/4.0/
• Abstract: Standard transformer architectures apply a single attention mechanism uniformly across all tokens and sequence positions, irrespective of local context or computational budget. We propose Meta-Attention, a framework that dynamically routes each token to the most appropriate attention strategy -- full softmax attention, linear (kernel) attention, or sliding-window local attention -- via a Bayesian Meta-Controller. Unlike prior routing approaches that use deterministic or prior-free learned routing, the Meta-Controller treats per-token mechanism selection as posterior inference under a compute-aware Dirichlet prior: routing weights are the output of an amortised variational posterior q(alpha | x_t; phi) trained with an Evidence Lower Bound (ELBO) objective that jointly encodes task performance and attention-mechanism cost. This design produces principled routing uncertainty estimates that govern the soft-to-hard routing transition, mitigates routing collapse without ad hoc load-balancing losses, and yields better compute-performance trade-offs than deterministic or prior-free learned routing at negligible overhead. Phase 1 empirical results on a Tiny LM benchmark confirm core predictions: the Bayesian controller's learned routing distribution implies a projected normalised FLOP cost of 25.1% under hard routing, vs. 59.3% for the prior-free baseline (-34.2 pp), and reduces routing entropy from 55.8% to 43.3% (-12.5 pp), demonstrating that the Dirichlet prior prevents routing collapse while the non-Bayesian model defaults to full attention. We present the Bayesian architecture, ELBO training objective, and a Phase 1 PyTorch prototype validating forward-pass correctness, posterior diversity, and a controlled ablation against a prior-free baseline. Code available at: https://github.com/KFEAL/meta-attention

Related papers

• CRISP: Cliff-awaRe Input-adaptive Sparse Prefilling with Structural-Mass-Motivated Routing [90.63350360476294]
  We present CRISP (Cliff-awaRe Input-adaptive Sparse Prefilling), which identifies and addresses two structural challenges in this dynamic routing paradigm.<n>We show that the routing decision can be read directly off the structure of the proxy attention map.<n>We formalize the post-softmax mass cliff and demonstrate theoretically that strictly cumulative coverage thresholds accumulate O(n) background noise at long contexts.
  arXiv  Detail & Related papers  (2026-09-01T22:49:46Z)
• Architecture-Dependent Causal Transfer of Activation States Across Large Language Models [51.56484100374058]
  Internal language activation states can be transferred causally between different large language models.<n>Injecting projected activations into a target model during generation produces a statistically significant, preregistered causal effect.<n>We conclude that end-to-end activation-state transfer between LLMs, as currently implemented, is architecture-dependent rather than universal.
  arXiv  Detail & Related papers  (2026-08-17T09:53:27Z)
• Adaptive Error Budget Allocation for Fault-Tolerant Quantum Resource Estimation: A Metaheuristic Approach [0.14323566945483496]
  System-level resource estimation is a key component of fault-tolerant quantum computing toolchains.<n>We propose a training-free optimization framework that performs derivative-free search directly on the Azure Quantum Resource Estor.<n>We evaluate our framework across 433 circuits spanning 2 to 91 qubits from 31 families in the MQT Bench.
  arXiv  Detail & Related papers  (2026-08-16T07:14:29Z)
• EntroRouter: Learning Efficient Model Routing via Entropy Regulation [45.88530258803968]
  We propose a single-round routing framework that treats entropy regulation as a core objective.<n>We demonstrate that Entro retains 98.3% of the strongest expert's accuracy while reducing computational costs by 48.25%.
  arXiv  Detail & Related papers  (2026-06-28T14:39:06Z)
• When Does Routing Become Interpretable? Causal Probes on Block Attention Residuals [0.0]
  Block Attention Residuals (Block AttnRes) replace fixed additive residuals with a learned softmax over earlier depth-source representations.<n>We ask whether such exposure suffices for mechanistic interpretation.
  arXiv  Detail & Related papers  (2026-06-11T10:37:04Z)
• CR^2: Cost-Aware Risk-Controlled Routing for Wireless Device-Edge LLM Inference [52.849509991178884]
  Large language models (LLMs) move from centralized clouds to mobile edge environments.<n> query-level routing between lightweight on-device models and stronger edge models provides a flexible mechanism to navigate this trade-off.<n>Existing routers are designed for centralized cloud settings and optimize token-level costs, failing to capture the dynamic latency and energy overheads in wireless edge deployments.
  arXiv  Detail & Related papers  (2026-05-12T11:50:15Z)
• Same Signal, Opposite Meaning: Direction-Informed Adaptive Learning for LLM Agents [24.71914991820339]
  LLM agents aim to invoke extra computation only when it improves performance.<n>Existing methods typically use confidence-, uncertainty-, or difficulty-based gates, assuming a fixed direction from the gating signal through compute need to the value of computation.<n>We show that this alignment is unstable, with reversals across environments and backbones even when the task is fixed.<n>We propose DIAL (Direction-Informed Adaptive Learning), a sparse gate trained from signal-agnostic counterfactual exploration to learn the utility direction of state features per.
  arXiv  Detail & Related papers  (2026-05-07T20:15:51Z)
• Posterior Augmented Flow Matching [64.1559809786948]
  Posterior-Augmented Flow Matching (PAFM) is a generalization of flow matching (FM)<n>We show that PAFM improves over FM by up to 3.4 FID50K across different model scales.
  arXiv  Detail & Related papers  (2026-05-01T17:59:59Z)
• Directional Routing in Transformers [0.0]
  directional routing is a lightweight mechanism that gives each transformer attention head learned suppression directions controlled by a shared router.<n>We trace the resulting circuits through mechanistic interpretability.<n> routing reduces perplexity 31-56% relative to the baseline, though downstream multiple-choice benchmarks do not yet reflect these gains.
  arXiv  Detail & Related papers  (2026-03-16T07:28:22Z)
• Proactive Routing to Interpretable Surrogates with Distribution-Free Safety Guarantees [0.5518378568494161]
  We study proactive (input-based) routing, in which a lightweight gate selects the model before either run.<n>We derive a feasibility condition linking safe routing to the base safe rate and risk budget, along with sufficient AUC thresholds ensuring that feasible routing exists.
  arXiv  Detail & Related papers  (2026-03-15T21:38:43Z)
• Efficient and Interpretable Multi-Agent LLM Routing via Ant Colony Optimization [58.59491516762626]
  We propose AMRO-S, an efficient and interpretable routing framework for Multi-Agent Systems (MAS)<n>AMRO-S models MAS routing as a semantic-conditioned path selection problem, enhancing routing performance through three key mechanisms.<n>Extensive experiments on five public benchmarks and high-concurrency stress tests demonstrate that AMRO-S consistently improves the quality--cost trade-off over strong routing baselines.
  arXiv  Detail & Related papers  (2026-03-13T12:26:05Z)
• Think When Needed: Model-Aware Reasoning Routing for LLM-based Ranking [25.69863022367215]
  Reasoning prompting can enhance ranking utility, but its benefits are inconsistent and come at a substantial computational cost.<n>We propose a reasoning routing framework that employs a lightweight, plug-and-play router head to decide whether to use direct inference (Non-Think) or reasoning (Think) for each instance before generation.
  arXiv  Detail & Related papers  (2026-01-26T05:09:07Z)
• You Only Need Your Transformer 25% of the Time: Meaning-First Execution for Eliminating Unnecessary Inference [0.0]
  We introduce Meaning-First Execution (MFEE), a control-plane architecture implementing this framework.<n>MFEE achieves 78.1% execution reduction while maintaining 100% exact-match equivalence for invoked executions.
  arXiv  Detail & Related papers  (2025-12-29T08:03:52Z)
• Dr.LLM: Dynamic Layer Routing in LLMs [55.11953638340419]
  Dr.LLM is a retrofittable framework that equips pretrained models with lightweight per-layer routers deciding to skip, execute, or repeat a block.<n>On ARC (logic) and DART (math), Dr.LLM improves accuracy by up to +3.4%p while saving 5 layers per example on average.
  arXiv  Detail & Related papers  (2025-10-14T17:51:26Z)
• What Layers When: Learning to Skip Compute in LLMs with Residual Gates [66.23658560048241]
  GateSkip is a residual-stream gating mechanism that enables token-wise layer skipping in decoder-only LMs.<n>Each Attention/MLP branch is equipped with a sigmoid-linear gate that condenses the branch's output before it re-enters the residual stream.
  arXiv  Detail & Related papers  (2025-10-13T16:31:50Z)
• Beyond 'Aha!': Toward Systematic Meta-Abilities Alignment in Large Reasoning Models [86.88657425848547]
  Large reasoning models (LRMs) already possess a latent capacity for long chain-of-thought reasoning.<n>We explicitly align models with three meta-abilities: deduction, induction, and abduction, using automatically generated, self-verifiable tasks.<n>Our three stage-pipeline individual alignment, parameter-space merging, and domain-specific reinforcement learning, boosts performance by over 10% relative to instruction-tuned baselines.
  arXiv  Detail & Related papers  (2025-05-15T17:58:33Z)

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.