帶有 Log-Linear 複雜度的區塊稀疏注意力
為何重要
將 Self-Attention 的複雜度從 O(N²) 降至 O(N log N),為縮放語言模型至長上下文解決了核心算力瓶頸,特別有利於 RAG 與長文檢索場景。對開發者而言,這代表複雜的高效能注意力運算元具備落地可行性,可降低長文本模型部署的計算成本。
為突破語言模型長上下文訓練受自注意力二次方開銷的受限,本研究提出 PISA,一種運用 Pyramid Top-K 策略的區塊稀疏注意力機制。
- PISA 透過多層遞進式篩選,將總計算複雜度從二次方最佳化為 O(N log N)。
- 研究團隊開發了硬體感知的 Triton kernels,在訓練與推理階段直接融合層級路由與 LogSumExp 評分。
- 評估顯示,模型在常識推理基準上與基線相仿,但在檢索任務上表現優於基線。