Technical Summary: AMOR (Adaptive Metacognitive Output Router)

Problem Statement

Recurrent-attention hybrid models aim to combine the inference efficiency of linear recurrent state-space models (SSMs) with the expressive power of attention mechanisms. However, existing hybrid architectures typically rely on fixed-schedule strategies, where attention is applied uniformly across all tokens or at predetermined layers (e.g., every $N$-th layer). This approach incurs the full quadratic cost of attention ($O(N^2)$) even when the recurrent state alone is sufficient to predict the next token with high confidence. Consequently, these models fail to fully capitalize on the efficiency gains of SSMs, particularly in long-context scenarios where the computational overhead of unnecessary attention blocks becomes a bottleneck. Furthermore, fixed hybrids often degrade under distribution shifts (e.g., length extrapolation) because they force attention mechanisms that may not be calibrated for the specific input difficulty.

Methodology: AMOR Architecture

The authors propose AMOR (Adaptive Metacognitive Output Router), a post-hoc hybrid architecture that dynamically allocates attention based on predictive uncertainty rather than fixed positions.

Core Architecture

AMOR is built upon a standard recurrent backbone (tested with Mamba2 and Gated DeltaNet) and appends $K$ (typically 3) entropy-gated attention blocks after the recurrent layers but before the final LM head.

• Backbone: The recurrent mixer processes the entire sequence, maintaining a fixed-size state.
• Post-Hoc Blocks: Unlike serial hybrids that replace recurrent layers, AMOR blocks are additive. They operate on the normalized output of the backbone.
• Asymmetry: The first AMOR block (Block 0) operates directly on the backbone's final normalized residual (post-norm), while subsequent blocks (1 and 2) use a standard pre-norm formulation. This design preserves the backbone as a standalone, valid language model, treating AMOR blocks as conditional refinements.

The Entropy Gate Mechanism

The central innovation is a hard, gradient-free gating mechanism inspired by the System 1 / System 2 distinction in cognitive science (Kahneman, 2011).

1. Uncertainty Estimation: For each token, the shared LM head computes logits. The normalized output entropy ($H_t$) is calculated from the softmax distribution.
2. Adaptive Threshold: A dynamic threshold $\tau$ is computed as an exponential moving average (EMA) of the batch median entropy ($\mu$) plus a scaled standard deviation ($\kappa\sigma$):
   $$\tau = \mu + \kappa\sigma$$
   This allows the threshold to track shifts in the model's uncertainty distribution during training without requiring learned router parameters.
3. Conditional Execution:
   • If $H_t > \tau$, the gate fires ($g_t=1$): The token undergoes causal attention (SDPA) with RoPE, and the output is added to the residual with a learned per-channel scaling factor $\alpha$.
   • If $H_t \le \tau$, the gate remains inactive ($g_t=0$): The attention computation is skipped, and the token passes through unchanged.
4. Gradient Flow: The gate is non-differentiable. Gradients flow through the attention weights ($W_Q, W_K, W_V, W_O$) only at firing positions. Crucially, $W_K$ and $W_V$ receive gradients from future firing positions to ensure the KV cache remains valid for retrieval, even if the current token skipped attention.

Training and Inference

• Training: AMOR is trained in a "dense with mask" mode using FlashAttention-2 for efficiency, mathematically equivalent to true sparse execution. The threshold $\tau$ is frozen at inference.
• Inference: The model performs a conditional skip. When the gate is off, the expensive $Q$ projection, attention calculation, and $O$ projection are bypassed, saving compute. However, $K$ and $V$ projections and cache updates always occur to maintain context continuity.

Key Contributions

1. Uncertainty-Driven Routing: AMOR introduces a simple, post-hoc mechanism that selectively invokes attention only when the recurrent backbone exhibits high predictive uncertainty (high entropy), rather than relying on fixed layer schedules.
2. Parameter-Efficient Hybridization: The method adds minimal parameter overhead (approx. <5% for the gating and attention blocks) while achieving performance comparable to or exceeding fixed-schedule hybrids.
3. Robustness to Distribution Shift: By relying on the recurrent backbone for low-uncertainty tokens, AMOR mitigates the performance degradation seen in fixed hybrids during length extrapolation, where attention mechanisms often suffer from positional drift.
4. Gradient-Free Gating: The approach avoids the complexity of learned routers or reinforcement learning, using a statistical threshold derived from the model's own output distribution.

Experimental Results

The authors evaluated AMOR across three scales (180M, 440M, 1.5B parameters) using Mamba2 and Gated DeltaNet backbones, comparing against pure recurrent models, pure Transformers, and fixed-schedule serial/fused hybrids.

• Common-Sense Reasoning: AMOR consistently matched or outperformed both pure recurrent backbones and fixed-schedule hybrids across all scales and benchmarks (LAMBADA, HellaSwag, PIQA, ARC, etc.).
• In-Context Retrieval: While pure recurrent models struggle with exact-copy retrieval, AMOR substantially improved over its backbone, achieving competitive results with fixed hybrids while activating attention on only ~22% of tokens.
• Long-Context Performance: On the LongBench benchmark, AMOR maintained stable performance as context length increased. In contrast, fixed-schedule hybrids and pure Transformers degraded significantly due to RoPE distribution shifts and memory constraints. AMOR's ability to skip attention on predictable tokens allowed it to scale more gracefully.
• Inference Efficiency: AMOR demonstrated better decode throughput scaling in long-context regimes. Theoretical analysis and empirical benchmarks (up to 448K context) showed that the savings from skipping attention outweigh the overhead of entropy calculation once the context length exceeds a certain crossover point (approx. 256K for 440M models).

Significance and Claims

The paper argues that when attention is applied is as critical as how much is applied. The authors claim that:

1. Efficiency and Expressivity are not a trade-off: By treating the allocation of compute as a routing problem based on input difficulty, models can achieve high expressivity without incurring the full cost of attention.
2. Simplicity is effective: Complex learned routers or auxiliary objectives are not necessary; a simple, output-driven entropy gate can effectively coordinate hybrid architectures.
3. Adaptive Computation: AMOR points toward a class of models that dynamically match computation to the difficulty of the input, preserving the strengths of lightweight state-space backbones while selectively deploying heavy attention mechanisms only when the model is uncertain.

The results suggest that selective allocation of attention based on predictive uncertainty offers a robust alternative to uniform or fixed routing strategies, particularly for long-context generation and retrieval tasks.