Orthros: Phase-Aware Heterogeneous Attention for Efficient Transformers

AI summary

Overview

Research area: Efficient attention mechanisms for large language models, specifically long-context inference and the pretraining-to-deployment architecture trade-offs of Transformers.

Technical level: Advanced. The paper assumes familiarity with self-attention, linear attention, sliding window attention, cross-attention, KV caching, FlashAttention, and asymptotic complexity analysis.

Scope: The paper proposes Orthros attention, a phase-aware heterogeneous attention module that uses sliding window attention during prefilling and full cross-attention at output-generation positions during decoding, and evaluates it from 160M to 6.8B parameters, through a 1.3B pretraining run on 100B tokens, conversions of Llama3.1-8B and Qwen3-8B, long-context retrieval tests, and system-level latency benchmarks.

What This Paper Is About

Efficient attention methods such as linear attention and sliding window attention lower Transformers from quadratic to linear complexity, but they lose in-context recall, so practitioners interleave full attention layers back into the network. Because the overall cost is then bounded by those full attention layers, hybrid Transformers only achieve a constant-factor speedup, and the quadratic bottleneck remains. Orthros attacks this by moving the efficiency trade-off out of the layer structure and into the phases of sequence processing: it applies sliding window attention across the whole sequence while adding full attention only at positions that predict the next token.

Key Contributions

1. A phase-aware heterogeneous attention module. Orthros attention serves as a single unified module that applies sliding window attention to all tokens to build a layer-wise KV cache, and additionally runs full cross-attention at output-generation positions against that same cache. It requires no KV eviction or compression.

2. A complexity argument for prefilling-dominated workloads. Because the full cross-attention term scales as O(T_out·T·d_model), the model's asymptotic cost is subquadratic when the generated output length T_out is much smaller than the total sequence length T, and reduces to strictly linear O(T·d_model²) when T_out is approximately constant.

3. Strict training–inference consistency without architecture change. Pretraining, supervised fine-tuning, prefilling, and decoding are all instantiated by the same formulation, differing only in the context length T_in and output length T_out, so existing pretrained full-attention models can be converted by swapping attention modules.

4. Broad empirical validation. Scaling and ablation studies from 160M to 6.8B parameters, a 1.3B model pretrained from scratch on 100B tokens, Transformer-to-Orthros conversion of Llama3.1-8B and Qwen3-8B, in-context recall benchmarks, a 128K needle-in-a-haystack test, and prefilling/end-to-end latency measurements.

Main Findings

• Scaling and ablation: Orthros is trained at 160M, 360M, 780M, 1.3B, 2.7B, and 6.8B parameters on 10B tokens with sequence length 2048 and a 0.25M-token batch size, and shows competitive scalability against Transformer baselines. The scale-offset (SO) design, which specializes the shared query projection for the different attention streams through lightweight element-wise transformations, gives a better trade-off between parameter efficiency and performance; removing it (Orthros w/o SO) is the ablated variant.

• Pretraining from scratch at 1.3B: Trained on 100B tokens of fineweb-edu with sequence length 2048 and a 2M-token batch size, Orthros reaches the best average on Common Sense Reasoning among the compared architectures: Orthros 53.98, Transformer 53.41, GDN 53.29, SWA 53.22, RetNet 51.96, GLA 51.90. The SWA baseline uses a window size of 512, while Orthros uses 128.

• Transformer-to-Orthros conversion works and even exceeds the originals: Converting Llama3.1-8B and Qwen3-8B and continual-training on an 80B-token Common Crawl corpus with a restricted math and code mixture, Orthros Conversion beats both the untrained Base and the same-architecture Base CPT on average. For Llama3.1 the average is 67.26 (Orthros Conversion) versus 67.00 (Base CPT) and 65.47 (Base); for Qwen3 it is 68.37 versus 67.55 (Base CPT) and 67.87 (Base).

• In-context recall is preserved: On recall-intensive tasks (DROP, NarrativeQA, SWDE, TriviaQA, FDA, SQuAD) at 1.3B scale, Orthros achieves the best average of 34.4, versus 33.9 for the full-attention Transformer, 28.8 for SWA, 28.4 for GDN, 27.1 for GLA, and 25.7 for RetNet.

• Long-context retrieval: An 8B Orthros model converted from Qwen3 and further trained for 20B tokens of long-context extension on ProLong data maintains near-perfect retrieval accuracy on the 128K needle-in-a-haystack benchmark.

• Prefilling speedup: Benchmarked with FLA-style models and high-performance kernels, 8B-Orthros shows consistently lower prefilling latency than Qwen3 on a single GPU from 8K to 128K input lengths and on 8 GPUs from 64K to 1M, reaching up to 21.9x speedup at one million tokens.

• End-to-end inference: On a single H100 GPU with the same SGLang inference engine, batch size 2, and the Triton attention backend, at 128K input and 128 output tokens Transformer reports 77.249 s time to first token and 84.911 s end-to-end, while Orthros reports 17.016 s and 24.647 s, Orthros-H reports 10.184 s and 12.771 s, and Linear Hybrid reports 25.566 s and 28.130 s. Decode-only throughput stays comparable to the corresponding baselines (Transformer 33.7 tokens/s versus Orthros 33.3; Linear Hybrid 98.8 versus Orthros-H 97.3), and the doubled FLOP count during decoding is masked by the I/O-bound nature of that phase. Gain shrinks as output length grows because decoding occupies more of the total time.

Methodology in Plain English

Orthros splits a sequence into two roles. Every token goes through a sliding window attention stream with window size 128, producing key and value vectors that form a shared, layer-wise KV cache. Separately, for the positions that must predict an output token, the model duplicates those input embeddings into extra "output-specific" tokens and runs full cross-attention from those tokens to the entire KV cache. The two streams use distinct learned scale-offset transformations of the query projection so the module can specialize each role while still sharing parameters. The outputs of both streams are concatenated and passed through the usual output projection, and the non-attention parts of the Transformer (feed-forward networks, normalization) process the extra tokens unchanged because they act token-wise.

The trick is that the same code path covers all four stages of use by changing two lengths. During pretraining, every token is a prediction target, so the whole sequence is duplicated (T_in = 0, T_out = T). During supervised fine-tuning, only response tokens are duplicated. During prefilling, only the final prompt token is duplicated (T_in = T−1, T_out = 1), so the expensive full attention runs essentially once. During decoding, each step processes a decoupled pair: one token updates the KV cache via sliding window attention, the other performs full cross-attention over the history. The extra per-step arithmetic is hidden behind the memory-I/O cost of loading weights, which is why decoding throughput stays close to the baseline.

Evaluation is done on Common Sense Reasoning benchmarks (ARC-C, ARC-E, HellaSwag, OBQA, PIQA, SIQA, WinoGrande) with lm-evaluation-harness, plus MMLU for the conversion study, recall-intensive datasets, needle-in-a-haystack, prefilling latency measurements, and end-to-end SGLang inference. The Orthros-H variant adds a layer-wise mixture of full-attention and linear-attention layers at a 1:3 ratio, with GDN used as the cross-attention part and the query bypassing its convolution module.

Why This Matters

For research, the paper reframes the efficiency question: instead of asking where to place cheap layers inside a stack, it asks when computation is cheap relative to hardware limits. It argues that conventional scaling analyses, which judge architectures by training FLOPs versus loss, can undervalue designs that spend extra training compute to buy sustained inference savings, and it calls for evaluation frameworks that account for phase-specific inference cost and workload shape.

Real-world applications named or implied by the paper:

• High-resolution multimodal understanding, which is prefilling-heavy.
• Repository-level code analysis over long contexts.
• Long-horizon agentic workflows where interaction sequences span hundreds of thousands to millions of tokens, and where structured tool invocations keep generated output short.
• Deployments that need long-context retrieval, such as question answering over long documents, validated by the 128K needle-in-a-haystack result.

For industry, the plug-and-play conversion path matters most: existing well-trained models such as Llama3.1-8B and Qwen3-8B can be converted with a moderate continual-training budget and then run with substantially lower prefilling latency, which is the dominant cost in interactive and agentic serving.

Future Directions

• Reducing training overhead and KV cache cost. Orthros currently computes a doubled sequence during training, and the memory efficiency of its KV cache is explicitly left unaddressed.

• Mixture-of-Experts compatibility. MoE settings need specialized designs so that the decoupled token pairs are not routed to different experts during decoding, which would worsen the I/O bound.

• Agentic reinforcement learning. The paper states that the agentic RL paradigm for Orthros still needs thorough investigation to demonstrate efficiency in multi-turn, long-context agentic applications.

• Inference framework adaptation. Optimizing the decoding mechanism across diverse scenarios, such as single-sample generation and batched decoding, is named as future work, along with the open question of how to fairly compare Orthros's redistribution of computation against the uniform computation of standard Transformers.

Target Audience

Researchers and engineers working on efficient attention, long-context language models, and LLM serving systems; practitioners who need to convert existing pretrained models rather than pretrain new ones; and architecture researchers interested in how training-time and inference-time costs trade off in phase-structured designs. Readers without a background in attention variants and complexity analysis will find the formal sections difficult, since the paper is written at an advanced level.

Authors’ abstract

Efficient attention mechanisms, such as linear and sliding window attention, aim to reduce the computational complexity of Transformers from quadratic to linear. However, their degraded in-context recall often necessitates interleaving full attention layers, leaving the quadratic bottleneck unresolved in long-context processing. To address this, we introduce Orthros attention, which shifts the efficiency paradigm from structural allocation to phase-aware decoupling. As a unified module, it dynamically switches its computation pattern across different phases of sequence processing, enabling linear complexity prefilling and full attention decoding. Our design grants Orthros attention a lower asymptotic complexity compared to vanilla Transformers in prevalent long-context scenarios characterized by heavy prefilling and light decoding. Beyond theoretical formulation and complexity analysis, extensive empirical evaluations demonstrate Orthros's superiority across multiple dimensions. First, we validate its fundamental language modeling capabilities, demonstrating scalability across varying model sizes and strong competitiveness in comprehensive architectural comparisons. Second, Orthros exhibits exceptional conversion feasibility via a data-efficient Transformer-to-Orthros conversion, yielding performance that even surpasses the original Transformer. Finally, Orthros achieves a $21.9\times$ prefilling speedup over Transformers with end-to-end inference gains, while preserving long-context capabilities, as demonstrated by near-perfect 128K needle-in-a-haystack results and Transformer-level performance on real-world in-context recall tasks. Overall, these advantages in high performance, flexibility, and efficiency position Orthros as a promising candidate to succeed the Transformer in industrial applications.