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.