Organizations: State Key Laboratory for Multimedia Information Processing, School of Computer Science, PKU-Anker LLM Lab, Peking University, Beijing, China · Contribution during internship at Huawei Technologies Co., Ltd. · Huawei Technologies Co., Ltd.
The attention mechanism has been the core component in modern transformer architectures. However, the computation of standard full attention scales quadratically with the sequence length, serving as a major bottleneck in long-context language modeling. Sliding window attention restricts the context length for better efficiency at the cost of narrower receptive fields. While existing efforts attempt to take the benefits from both sides by building hybrid models, they often resort to static, heuristically designed alternating patterns that limit efficient allocation of computation in various scenarios. In this paper, we propose Switch Attention (SwiAttn), a novel hybrid transformer that enables dynamic and fine-grained routing between full attention and sliding window attention. For each token at each transformer layer, SwiAttn dynamically routes the computation to either a full-attention branch for global information aggregation or a sliding-window branch for efficient local pattern matching. An adaptive regularization objective is designed to encourage the model towards efficiency. Moreover, we adopt continual pretraining to optimize the model, transferring the full attention architecture to the hybrid one. Extensive experiments are conducted on twenty-three benchmark datasets across both regular (4K) and long (32K) context lengths, demonstrating the effectiveness of the proposed method.
Figures & tables
Figure 1: Conventional hybrid transformers (a) often adopt manually designed, static hybrid schemes. In contrast, the proposed SwiAttn (b) dynamically selects the type of attention for each input token at each transformer layer, enabling fine-grained allocation of computation.
Figure 2: During the continual pretraining stage (a), we compute both the full attention branch and the SWA branch separately. Subsequently, the router is applied to select the attention output from different branches for each token’s hidden representation. The selection result is then processed by a feed-forward network (FFN Layer). During the decoding stage (b), the router is used to decide the branch for attention computation. The full attention branch and the SWA branch share a unified KV cache for better efficiency.
Datasets
LMB.
Wiki.
AVG
BoolQ
PIQA
SIQA
Hella.
Wino.
ARC-e
ARC-c
OBQA
AVG
Metrics
ppl ↓
ppl ↓
ppl ↓
acc ↑
acc ↑
acc ↑
acc_n ↑
acc ↑
acc ↑
acc_n ↑
acc_n ↑
acc ↑
SWA-ZS
7.16
16.6
11.9
62.2
74.0
41.4
60.1
57.1
62.9
34.3
35.4
50.5
SWA-CPT
7.02
15.6
11.3
62.1
74.5
42.0
60.1
57.9
62.7
35.5
35.2
52.6
SH (3:1)
7.32
15.2
11.3
62.3
73.7
42.0
60.2
57.9
62.3
33.6
36.7
53.6
FullAttn
7.16
15.0
11.1
62.2
74.0
41.4
60.1
57.1
62.9
34.3
35.4
53.4
SwiAttn
7.28
15.0
11.1
62.0
74.0
42.3
60.2
57.8
63.0
34.0
36.8
53.7
Table 1: Performance on common sense reasoning datasets. Best results are in bold and second-best underline .
Datasets
SWDE
FDA
TQA
NQ
DROP
SQD
AVG
SWA-ZS
52.0
22.5
64.6
30.9
31.1
49.2
41.7
SWA-CPT
49.1
24.0
65.7
28.4
32.2
49.7
41.5
SH (3:1)
56.5
67.9
65.1
32.7
31.3
50.1
50.6
FullAttn
60.8
76.4
64.6
35.1
31.1
49.2
52.9
SwiAttn
60.2
75.1
65.2
35.2
32.4
49.4
52.9
Table 2: Performance on in-context retrieval datasets. Best results are in bold and second-best underline .
Datasets
GvR
MNs
MFQA
Qas
TQA
SSM
TRec
AVG
SWA-ZS
19.2
18.2
18.5
10.5
51.1
27.4
18.7
23.4
SWA-CPT
19.2
18.7
24.2
13.1
57.8
30.7
53.0
31.0
DuoAttention
21.6
19.9
30.9
14.5
62.4
29.4
58.7
33.9
SH (6:1)
21.0
19.0
26.2
20.1
60.0
31.2
55.7
33.3
SH (3:1)
22.6
20.0
29.6
14.5
68.1
32.3
59.3
35.2
FullAttn
24.9
20.9
33.9
18.6
69.3
30.2
63.7
37.4
Table 3: Performance on long-context understanding. Best results are in bold and second-best underline .
Variants
GvR
MNs
SSM
AVG
V1: w/o Adaptive Reg.
24.1
20.7
30.0
24.9
V2: w/o the NLL term
24.3
20.9
30.7
25.3
V3: α=10
24.1
20.9
30.5
25.2
V4: α=1
23.8
20.9
29.9
24.9
V5: α=0
23.9
20.6
29.5
24.7
V6: ε=0.2
24.7
20.9
31.7
25.8
Table 4: Ablation studies of SwiAttn .
Figure 3: Needle-in-a-haystack retrieval accuracy in different context positions (depth percentage) across 32 K context lengths. The proposed SwiAttn achieves perfect retrieval accuracy with its dynamic and fine-grained hybrid attention mechanism.
Figure 4: Efficiency comparison under different token positions of SwiAttn (ours) and FullAttn baseline. We report GFLOPs for the prefill stage and memory access in terms of averaged tokens in the decode stage.
Figure 5: The ratio of using full attention branches across different transformer layers. The averaged full attention ratio is around 0.13 .
Figure 6: The ratio of full attention branches when generating the needle in the needle-in-a-haystack retrieval task. When the needles fall within the sliding window, the full attention ratio is low, while a higher ratio is observed when the needle is outside the sliding window.
Appendix figures & tables9 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 8: Language modeling loss curves of (a) pretraining of the full attention transformer on 4 K-length sequences, (b) continual pretraining of SwiAttn on 4 K-length sequences, (c) pretraining of the full attention transformer on 32 K-length sequences, and (d) continual pretraining of SwiAttn on 32 K-length sequences.
Figure 9: The changes of full attention ratio during the continual pretraining stage. The ratio is averaged within each transformer layer.
Figure 10: The changes of full attention ratio during the continual pretraining stage. The averaged ratio over all transformer layers is reported.
Datasets
LMB.
Wiki.
AVG
BoolQ
PIQA
SIQA
Hella.
Wino.
ARC-e
ARC-c
OBQA
AVG
Metrics
ppl ↓
ppl ↓
ppl ↓
acc ↑
acc ↑
acc ↑
acc_n ↑
acc ↑
acc ↑
acc_n ↑
acc_n ↑
acc ↑
FullAttn
6.14
13.8
9.99
63.3
74.8
43.9
63.5
61.1
66.3
35.4
39.0
55.9
SwiAttn
6.19
13.8
9.99
65.6
74.8
43.4
63.5
61.7
65.0
34.6
38.4
55.9
Appendix
Table 5: Performance of 3B models on common sense reasoning datasets.
In-Context Retrieval
LongBench-E
Methods
SWDE
FDA
TQA
NQ
DROP
SQD
AVG
GvR
MNs
MFQA
Qas
TQA
AVG
FullAttn
65.4
76.5
65.0
38.0
34.4
53.3
55.4
27.3
21.5
39.6
27.6
70.0
37.2
SwiAttn
66.3
77.7
65.6
38.1
33.5
53.7
55.8
26.8
20.2
39.7
26.6
69.5
36.6
Appendix
Table 6: Performance of 3B models on in-context retrieval and LongBench-E datasets.
Method
NIAH Accuracy
FullAttn
1.00
SwiAttn
1.00
SWA
0.20
SH (3:1)
0.95
Appendix
Table 7: Aggregate NIAH retrieval accuracy under the same evaluation setting.
Methods
Forward time (s)
Backward time (s)
FullAttn
0.0595
0.1459
SwiAttn
0.0645
0.1574
Overhead
8.4%
7.9%
Appendix
Table 8: Comparison of runtime per layer during continual pretraining (CPT) with a sequence length of 32K.
Standard Setting
Simple Setting
Methods
Memory
GFLOPs
Memory
GFLOPs
LongBench-E
SH (3:1)
8.0K
1.38
8.0K
1.38
35.2
SwiAttn
7.3K
1.26
2.6K
0.44
37.4
Appendix
Table 9: Efficiency comparison between SwiAttn and the SH (3:1) baseline.
Methods
Commonsense
Retrieval
LongBench
SH (3:1)
53.6 ± 0.32
50.6 ± 0.09
35.2 ± 0.44
FullAttn
–
–
37.4 ± 0.18
SwiAttn
53.7 ± 0.30
52.9 ± 0.11
37.4 ± 0.35
p -value (vs. SH)
0.713
<0.001
<0.01
p -value (vs. FullAttn)
–
–
0.941
Appendix
Table 10: Performance comparison over three runs and statistical significance ( p -value) across commonsense reasoning, in-context retrieval, and LongBench-E.
Hybrid attention models improve long-context efficiency by retaining only a subset of full-attention layers and replacing the remaining layers with linear attention. However, the effectiveness of Transformer-to-hybrid conversion critically depends on which layers preserve full attention. Existing hybrid layer selection methods typically rely on heuristic strategies such as fixed placement patterns or layerwise scoring, implicitly treating layer importance as isolated and overlooking the interdependent layer effect under a global hybrid configuration. In this work, we formulate hybrid layer selection as a budget-constrained subset optimization problem. We further propose FlashMorph (Fast LAyer Selection for Hybrid MORPHing), an effective, efficient and scalable layer selection method for Transformer-to-hybrid conversion. FlashMorph first constructs a morphable model by equipping each full-attention layer with a converted linear-attention branch. It then freezes all model weights and jointly optimizes layerwise gates on synthetic long-context retrieval data, with a linearization regularization that encourages the model to rely on linear attention for efficiency. The learned gates are discretized under a preset full-attention budget to instantiate the hybrid architecture, followed by standard logits distillation and long-context finetuning. Extensive experiments show that FlashMorph discovers more effective hybrid configurations, preserves strong long-context recall and general benchmark performance while substantially reducing layer selection cost compared with existing layer selection methods, demonstrating its effectiveness, efficiency, and scalability.
Disen Lan, Jianbin Zheng, Yuxi Ren +5
1Fudan University · 2ByteDance Seed · Work done at ByteDance Seed +1
The quadratic complexity of attention poses a critical bottleneck for long-context processing, spurring interest in hybrid attention designs. Most open-source hybrid models adopt a layer-wise strategy. Yet, prior work has noted the inherent difficulty of integrating Linear Attention (LA) with Full Attention (FA), suggesting that the design space of attention hybridization remains underexplored. To probe this space, we conduct interpretability analysis and observe that layers exhibit block-wise functional similarity, while individual heads within the same layer display distinct functional specialization despite sharing input features. This head-level heterogeneity suggests that the head dimension provides a natural and principled granularity for fusing heterogeneous attention signals. Building on this insight, we introduce HydraHead, a novel architecture that hybridizes FA and LA along the head axis. HydraHead features two key innovations: (1) an interpretability-driven selection strategy that identifies retrieval-critical heads and preserves FA only for them, and (2) a scale-normalized fusion module that reconciles the distributional gap between FA and LA head outputs. By leveraging a three-stage transfer pipeline with parameter reuse and distillation, we achieve high-performance hybrid models with minimal training overhead. Under a unified training setup, HydraHead outperforms other hybrid designs in long-context tasks while maintaining strong general reasoning. With interpretability-driven head selection, it matches a 3:1 layer-wise hybrid's long-context performance at a 7:1 LA-to-FA ratio. Crucially, trained on only 15B tokens, HydraHead achieves over 69% improvement over the baseline at 512K context length, approaching Qwen3.5, a leading model of comparable size with a native context length of 256K. This highlights the significant scaling potential of head-level hybridization.
The architectural design of Large Language Models (LLMs) is shifting from traditional full-attention-only models to hybrid models, which combine different attention modules to improve long-context efficiency and performance in length extrapolation and context extension. To explain why hybrid models work and how to design them better, we propose Mechanics of Long-Context Hybrid Models. As Part 1.1 of this series, we begin with hybrids of full attention and either sliding-window attention (SWA) or gated variants of linear attention (LA), represented by GLA and GDN. We first observe a Seesaw Effect in Context Extension: LA hybrids benefit more from long-context continual pretraining, whereas SWA hybrids perform better under length extrapolation. We attribute this behavior to differences in the positional inductive biases induced by these attention mechanisms. We find that SWA hybrids suffer from a Short-Context Learning Trap, Short-Window Weariness, and Long-Window Laziness, and require extended windows to enhance performance in continual long-context pretraining. For LA hybrids, we summarize the Matthew Effect of Hybrid Position Extrapolation and propose Sliding-Window Linear Attention, achieving 16× training-free length extrapolation while maintaining 100% accuracy on NIAH-SK1 in 64k context length.
Xiaoran Liu, Ziwei He, Xipeng Qiu
Fudan University · Shanghai Innovation Institute · OpenMOSS Team