Show HN: Proxy Block-CAGE，一種新的稀疏塊注意力機制
本文介紹Proxy Block-CAGE，一種用於長上下文預填充的稀疏注意力原型。通過將Q/K/V分塊，使用池化摘要選擇關鍵塊，僅對選定塊計算精確注意力。在小型字符級Transformer上，16K上下文下比密集註意力快2.61倍，僅使用0.39%的密集註意力容量，且保持100%最終token一致。文章詳細闡述了從遺傳算法到數學優化的演變過程。
在人工智能研究領域，處理長序列的Transformer模型面臨着注意力機制計算量隨序列長度平方增長的問題。近日，一位名為Iqbal Addou的博士研究生在GitHub上發佈了一個名為Proxy Block-CAGE的注意力原型，旨在通過稀疏化方法降低長上下文預填充階段的計算開銷。該項目不僅是一項技術探索，更是一次將ChatGPT作為研究夥伴的實踐。
Proxy Block-CAGE的核心思想是將查詢（Q）、鍵（K）和值（V）分割成固定大小的塊，然後使用每個塊的平均池化摘要來計算塊之間的代理分數。根據這些分數，每個查詢塊會選擇得分最高的幾個鍵塊，隨後只在這些選擇的塊內執行標準的精確注意力計算。這種方法在保持注意力質量的同時，大幅減少了需要計算的查詢-鍵對數量。
作者在配備NVIDIA RTX 2080 Max-Q GPU的筆記本電腦上，對一個兩層字符級Transformer進行了測試。結果顯示，在16,384個token的上下文長度下，Proxy Block-CAGE實現了2.61倍的加速，注意力密度僅為密集註意力的0.39%，而最終token的top-1和top-5準確率保持100%。雖然測試規模較小，但足以證明該方法的潛力。
文章詳細記錄了項目的演變過程。最初，作者嘗試使用遺傳算法（GA）為每個查詢token選擇稀疏的注意力掩碼（稱為“CAGE”），但token級別的搜索耗時巨大。隨後，他將搜索單元從token升級為塊，大幅減少了搜索次數，但遺傳算法的計算開銷仍然過高。最終，作者發現對於僅基於log-sum-exp的簡單目標函數，最優解等價於簡單的top-k選擇，遺傳算法由此被更高效的有序選擇取代。這一發現促成了Proxy Block-CAGE的最終形式，即通過池化摘要的代理分數進行top-k塊選擇。
從數學上看，Proxy Block-CAGE的計算複雜度約為O(BHd(n²/(b_q b_k) + n m b_k))，其中b_q和b_k是塊大小，m是每個查詢塊選擇的鍵塊數。在典型配置（b_q=b_k=32, m=2）下，16K上下文時的計算位置數僅為密集註意力的約1/100。儘管由於工程優化程度不同，實際加速比未達到理論值，但已足夠表明稀疏注意力在長上下文場景中的優勢。
作者強調，這只是一個可重複的研究原型，而非生產級解決方案。他希望通過公開此項目，獲得來自社區的反饋，並鼓勵更多研究者探索稀疏注意力中的數學與算法創新。