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

The Meta-Attention framework is proposed, which dynamically routes each token to the most suitable attention strategy via a Bayesian Meta-Controller. It achieves a 34.2 percentage point reduction in FLOP cost on the Tiny LM benchmark, providing a new approach to solving the routing collapse problem.

• Original Author/Maintainer: KFEAL Research Team
• Source Platform: arXiv
• Original Title: Meta-Attention: Bayesian Per-Token Routing for Efficient Transformer Inference
• Original Link: http://arxiv.org/abs/2605.28384v1
• Source Publication/Update Time: 2026-05-27

Standard Transformer architectures apply a single attention mechanism uniformly to all tokens and sequence positions, regardless of local context or computational budget. This one-size-fits-all design means that even if some tokens only need simple local attention, the model still computes full global attention for them, resulting in significant computational waste.

As sequence length increases, this efficiency issue becomes more severe. In scenarios like long document processing and code generation, attention computation often becomes a bottleneck for inference speed. How to allocate appropriate computational resources to different tokens while maintaining model performance has become a key challenge in improving Transformer efficiency.

The core idea of the Meta-Attention framework is to dynamically select the most suitable attention strategy for each token. The framework supports three attention mechanisms:

1. Full Softmax Attention: Provides the strongest global context understanding capability
2. Linear (Kernel) Attention: More computationally efficient, suitable for long sequences
3. Sliding Window Local Attention: Balances efficiency and local context capture

The key point is that this routing decision is not static but dynamically made based on the local context of each token. Some tokens may need global attention to understand long-distance dependencies, while others may only require local attention.

Unlike previous methods that use deterministic or prior-free learning for routing, Meta-Attention adopts a Bayesian framework to handle routing decisions. Specifically:

Meta-Controller treats the mechanism selection for each token as posterior inference under a computation-aware Dirichlet prior. The routing weights are outputs of the variational posterior q(alpha | x_t; phi), which is trained via an Evidence Lower Bound (ELBO) objective, encoding both task performance and attention mechanism cost.

This design has several notable advantages:

1. Principled Uncertainty Estimation: The Bayesian framework naturally provides uncertainty quantification for routing decisions
2. Soft-to-Hard Routing Transition: Uncertainty estimation guides the transition from soft routing (probabilistic mixing) to hard routing (discrete selection)
3. Prevention of Routing Collapse: The Dirichlet prior prevents the collapse phenomenon where all tokens are routed to a single mechanism
4. No Additional Load Balancing Loss: The Bayesian prior itself achieves load balancing without the need for ad hoc loss functions

Phase 1 experiments on the Tiny LM benchmark validate Meta-Attention's core predictions:

FLOP Cost大幅降低: The learned routing distribution of the Bayesian controller means that the normalized FLOP cost projected under hard routing is 25.1%, compared to 59.3% for the prior-free baseline—a reduction of 34.2 percentage points. This implies that Meta-Attention can achieve similar performance with less than half the attention computation.

Routing Entropy Reduction: Routing entropy decreased from 55.8% to 43.3% (a 12.5 percentage point reduction), indicating that the Dirichlet prior indeed prevents routing collapse. In contrast, non-Bayesian models tend to default to full attention.

Negligible Additional Overhead: The additional computational overhead from these gains is minimal, making Meta-Attention attractive for practical deployment.

The technical architecture of Meta-Attention includes several key components:

Variational Posterior Network: Outputs distribution parameters for the three attention mechanisms for each token. This is a lightweight network that usually adds only a small number of parameters.

Dirichlet Prior Design: The prior design considers computational cost, favoring more efficient attention mechanisms (e.g., linear attention) unless task performance requires full attention.

ELBO Training Objective: The training objective balances task performance and routing efficiency; this trade-off can be controlled by adjusting hyperparameters.

Soft-to-Hard Routing Scheduling: Soft routing (probabilistic weighting) is used in the early stages of training to ensure gradient flow, and gradually transitions to hard routing (discrete selection) in later stages to maximize efficiency gains.

Meta-Attention provides new insights for efficient Transformer inference:

First, token-level dynamic routing is more effective than layer-level static mixing. Tokens at different positions have very different attention needs; uniform processing inevitably leads to waste.

Second, the Bayesian framework provides a theoretical foundation for routing decisions. Uncertainty estimation not only helps prevent collapse but can also be used for adaptive inference—when the model is uncertain about a routing decision, it can conservatively choose a stronger attention mechanism.

Finally, computation-aware prior design is key to achieving efficient routing. The prior should encode our knowledge of the efficiency of different attention mechanisms, guiding the model to make informed trade-offs between performance and efficiency.