近日，由微軟研究院的孫宇濤、葉天竹、董力、夏玉清、陳健、高亦昭、曹世杰，以及清華大學的王建勇和香港大學的研究團隊共同發(fā)表了一篇名為《Rectified Sparse Attention》（矯正稀疏注意力機制）的研究論文。這項研究針對大型語言模型處理超長文本時的效率問題提出了一種全新的解決方案，團隊的研究成果可在https://aka.ms/ReSA-LM獲取。

隨著ChatGPT、Claude等大型語言模型的普及，人們越來越希望AI能夠處理更長的文本內(nèi)容。想象一下，你希望AI能閱讀并分析一本完整的書籍，或者處理成百上千頁的法律文件。這樣的長文本處理需求正在迅速增長，一些最新的模型甚至可以處理數(shù)百萬個詞的文本。但是，處理如此長的文本帶來了一個嚴峻的挑戰(zhàn)：效率問題。

為什么處理長文本會如此耗時呢？這里需要簡單了解一下大型語言模型的工作方式。當模型生成文本時，它需要"記住"之前的所有內(nèi)容，這些記憶被存儲在所謂的"KV緩存"（Key-Value緩存）中。想象這個緩存就像模型的短期記憶，存儲著它剛剛處理過的信息。在標準的處理方式中，模型每生成一個新詞，都需要查看所有之前的內(nèi)容，這就像是一個人在寫作時，每寫一個詞都要回頭完整閱讀之前寫的所有內(nèi)容。當文本變得很長時，這個過程會變得異常緩慢。

有研究者提出了"稀疏注意力"的解決方法，就像是讓模型只關(guān)注最相關(guān)的部分內(nèi)容，而不是所有內(nèi)容。這就好比寫作時只回顧最重要的幾段，而不是整篇文章。這種方法確實能提高速度，但隨著生成文本的增長，質(zhì)量會逐漸下降。為什么會這樣呢？因為在稀疏處理過程中，微小的計算誤差會不斷積累，最終導(dǎo)致模型的"記憶"與訓練時的狀態(tài)越來越不一致，就像是復(fù)述一個故事時，每次復(fù)述都會有小的偏差，最終導(dǎo)致故事變得面目全非。

微軟研究團隊提出的"矯正稀疏注意力"（ReSA）方法巧妙地解決了這個問題。這個方法就像是定期"校準"模型的記憶，確保它不會偏離軌道太遠。具體來說，ReSA結(jié)合了兩個關(guān)鍵步驟：首先，它使用"分組塊稀疏注意力"讓模型只關(guān)注最相關(guān)的內(nèi)容，加快處理速度；其次，它會定期進行"密集矯正"，即完整地檢查一遍最近生成的內(nèi)容，修正可能積累的誤差。這就像是一個人在講述長故事時，大多數(shù)時候只回憶關(guān)鍵部分，但每隔一段時間會完整地回顧一下最近的內(nèi)容，確保故事不會偏離原意。

研究團隊通過多項實驗證明，這種方法不僅保持了生成內(nèi)容的高質(zhì)量，還顯著提高了處理速度。在數(shù)學推理、語言建模和信息檢索等任務(wù)上，ReSA的表現(xiàn)幾乎與完整的密集注意力機制相當，但速度快了許多。特別是在處理25.6萬詞的超長文本時，整體速度提升了2.42倍，這是一個相當顯著的改進。

一、分組塊稀疏注意力：智能篩選重要信息

為了理解ReSA的工作原理，我們需要先了解其核心組件之一：分組塊稀疏注意力。想象你正在閱讀一本千頁的百科全書，當你閱讀到第800頁時，你不可能記住之前所有頁面的每一個細節(jié)。相反，你的大腦會選擇性地記住一些重要的信息塊，比如關(guān)鍵概念、重要事實或與當前閱讀內(nèi)容相關(guān)的段落。

分組塊稀疏注意力機制模擬了這種人類的選擇性注意力。它將之前的內(nèi)容劃分為多個"記憶塊"，每個塊包含一組連續(xù)的詞。當模型需要生成新內(nèi)容時，它不會檢查所有之前的內(nèi)容，而是智能地選擇最相關(guān)的幾個記憶塊進行參考。

那么，模型如何知道哪些塊是重要的呢？每當模型需要生成新內(nèi)容時，它會進行一個"相關(guān)性評分"的過程。對于每個記憶塊，模型會計算它與當前生成需求的相關(guān)程度。這就像是你在寫作時，會下意識地想到與當前段落最相關(guān)的之前內(nèi)容。

為了提高效率，研究團隊采用了一種巧妙的表示方法。對于每個記憶塊，模型不會存儲其中的所有詳細信息，而是存儲兩個關(guān)鍵"摘要"：塊中最小值和最大值。這就像是給每個內(nèi)容塊做了一個"范圍標簽"，說明這個塊中包含的內(nèi)容大致在什么范圍內(nèi)。這種簡化的表示方法大大減少了計算量，同時保持了足夠的信息來判斷相關(guān)性。

此外，研究團隊還引入了"共享分組"技術(shù)，允許多個注意力頭共享同一個稀疏模式。這就像是團隊成員共享同一份閱讀筆記，而不是每個人都單獨做筆記，進一步提高了效率。

在實際運行中，模型會動態(tài)確定需要關(guān)注的記憶塊數(shù)量。它總是會保留最近的幾個塊（稱為"局部窗口"），確保短期連貫性；同時，它還會根據(jù)內(nèi)容長度，選擇一定比例的歷史塊進行關(guān)注。例如，當處理非常長的文本時，模型可能只關(guān)注10%的歷史內(nèi)容，但這些內(nèi)容是經(jīng)過精心篩選的最相關(guān)部分。

二、密集矯正：定期校準確保質(zhì)量不下滑

雖然稀疏注意力機制大大提高了處理速度，但正如前面提到的，它會導(dǎo)致誤差累積問題。這就好比是導(dǎo)航系統(tǒng)中的"漂移"現(xiàn)象：如果你只依靠估算行走距離，而不參考GPS信號進行校正，那么誤差會隨著時間不斷累積，最終導(dǎo)致你的位置估計與實際位置相差甚遠。

ReSA通過"密集矯正"巧妙地解決了這個問題。密集矯正的核心思想非常直觀：每生成一定數(shù)量的新內(nèi)容后（比如每32個詞），模型會暫停稀疏處理，對最近生成的這批內(nèi)容進行一次完整的密集注意力計算。這就像是導(dǎo)航系統(tǒng)定期接收GPS信號來校準位置，確保累積誤差不會太大。

這個過程在技術(shù)上是如何實現(xiàn)的呢？當達到預(yù)設(shè)的矯正頻率（默認是32個詞）時，模型會將最近生成的這批詞一次性輸入，使用完整的密集注意力機制重新計算它們的表示，并用這些精確的表示更新KV緩存。這樣，即使之前的稀疏計算引入了一些誤差，這些誤差也會被定期"重置"，不會無限制地累積。

值得注意的是，這種矯正是批量進行的，而不是逐詞進行。這意味著雖然矯正過程使用了計算成本較高的密集注意力，但由于是一次性處理多個詞，平均到每個詞的額外成本其實很小。就像是你可能每周花一個小時整理一次筆記，而不是每寫一段就停下來整理，這樣總體效率會更高。

研究者們發(fā)現(xiàn)，通過適當選擇矯正頻率，可以在保持生成質(zhì)量和提高計算效率之間取得良好的平衡。當矯正頻率設(shè)為32（即每生成32個詞進行一次矯正）時，模型的表現(xiàn)接近于完全使用密集注意力的上限，同時仍然保持了顯著的速度優(yōu)勢。

三、全面評估：ReSA在各種任務(wù)中的表現(xiàn)

為了驗證ReSA的有效性，研究團隊在多種任務(wù)上進行了廣泛的實驗評估，包括數(shù)學推理、語言建模和信息檢索。

在數(shù)學推理任務(wù)上，研究者使用了多個具有挑戰(zhàn)性的數(shù)據(jù)集，包括Minerva Math、高考2023英文數(shù)學題、奧林匹克數(shù)學題集、AIME24和AMC23。這些任務(wù)特別適合測試長文本處理能力，因為解決復(fù)雜數(shù)學問題通常需要模型維持長期的推理鏈。結(jié)果顯示，使用純稀疏注意力的模型表現(xiàn)明顯不如密集注意力，特別是在生成較長答案時。而使用ReSA的模型幾乎能達到與密集注意力相同的準確率，同時保持了更高的效率。

在語言建模任務(wù)中，研究者設(shè)計了一個巧妙的實驗，將輸入序列分為兩部分：前綴使用密集注意力處理，后綴使用稀疏注意力或ReSA處理。這樣可以模擬不同長度生成過程中的性能變化。結(jié)果再次證明，隨著生成長度的增加，純稀疏注意力的性能明顯下降，而ReSA能夠保持穩(wěn)定的高性能。特別是當矯正頻率設(shè)為32時，ReSA的表現(xiàn)幾乎達到了理論上限。

研究者還嘗試了不同的稀疏率設(shè)置。當模型只關(guān)注10%的歷史內(nèi)容時（即稀疏率為0.9），ReSA仍能保持與密集注意力相當?shù)男阅?，這證明了該方法在高稀疏率下的有效性。

在RULER信息檢索基準測試中，由于輸出序列相對較短，矯正的作用不太明顯，但結(jié)果仍然表明適當?shù)南∈杪试O(shè)置（p=0.9）能夠在保持性能的同時提高效率。

四、效率分析：ReSA帶來的速度提升

除了保持生成質(zhì)量外，ReSA的另一個主要優(yōu)勢是顯著提高了處理效率。研究團隊在標準GPU硬件（NVIDIA A100-80G）上進行了詳細的效率測試。

從理論上分析，ReSA的內(nèi)存訪問成本可以表示為：mem(KV緩存)×(1/b + p + 1/f)，其中b是塊大小，p是稀疏率，f是矯正頻率。相比于密集注意力需要訪問整個KV緩存，ReSA只需要訪問其中的一小部分，從而大大減少了內(nèi)存帶寬需求。

在實際測試中，研究者比較了在不同序列長度下的延遲情況。結(jié)果顯示，當序列長度增加時，密集注意力的延遲急劇增加，而ReSA的延遲增長則平緩得多。特別是在超長序列（25.6萬詞）下，ReSA的速度優(yōu)勢最為明顯。

延遲分析還顯示，稀疏估計和注意力計算占用了相當?shù)臅r間，這是因為它們的內(nèi)存訪問模式類似（一個與緩存大小/塊大小成比例，另一個與緩存大小×稀疏率成比例）。而矯正過程的額外開銷相對較小，在25.6萬詞長度下僅占總延遲的32.7%。

在端到端的吞吐量測試中，ReSA在FP16精度下提供了最高2.28倍的加速，在INT4量化精度下提供了最高2.44倍的加速。這種加速效果在序列越長時越明顯，充分證明了ReSA在實際應(yīng)用中的價值。

五、參數(shù)選擇與權(quán)衡

為了幫助用戶選擇最適合自己需求的參數(shù)設(shè)置，研究團隊還進行了詳細的消融實驗，分析了不同矯正頻率和稀疏率對性能的影響。

實驗結(jié)果顯示，在大多數(shù)任務(wù)上，矯正頻率f=32或f=64能夠在保持高質(zhì)量和提高效率之間取得最佳平衡。雖然更頻繁的矯正（f=16）可能帶來輕微的性能提升，但矯正開銷也會相應(yīng)增加，降低了整體效率。另一方面，即使較低頻率的矯正（f=128）也能保持顯著的性能提升，這表明矯正機制對不頻繁更新具有很強的魯棒性。

對于稀疏率，研究者發(fā)現(xiàn)p=0.9（即關(guān)注10%的歷史內(nèi)容）通常能提供最佳的效率與性能平衡。雖然更低的稀疏率（如p=0.8，關(guān)注20%的內(nèi)容）可能在某些情況下提供略好的性能，但考慮到較低稀疏率會導(dǎo)致更多的內(nèi)存訪問和計算成本，p=0.9成為了推薦的默認設(shè)置。

值得一提的是，ReSA的設(shè)計是完全無需訓練的，它可以直接應(yīng)用于現(xiàn)有的預(yù)訓練模型，無需任何微調(diào)或重新訓練。這使得它特別適合在實際部署中使用，因為用戶可以立即享受到效率提升，而不需要耗費大量資源進行模型再訓練。

這項研究的成功之處在于，它巧妙地結(jié)合了稀疏注意力的高效率和密集注意力的高質(zhì)量，為長文本生成找到了一條實用的中間道路。正如研究者所展示的，通過適當?shù)膮?shù)設(shè)置，ReSA可以在幾乎不損失性能的情況下，將超長文本處理的速度提高2倍以上，這對于實際應(yīng)用具有重要意義。

當人工智能技術(shù)日益融入我們的生活時，處理長文本的能力變得越來越重要。無論是分析法律文件、醫(yī)學文獻，還是與用戶進行深入的多輪對話，都需要模型能夠高效地處理長文本。微軟研究團隊提出的ReSA方法為這一挑戰(zhàn)提供了一個簡單而有效的解決方案，它不需要復(fù)雜的模型修改或重新訓練，就能顯著提高處理效率。

這項研究的代碼已在https://aka.ms/ReSA-LM公開，感興趣的讀者可以直接嘗試將這一技術(shù)應(yīng)用到自己的項目中。隨著這類高效處理技術(shù)的發(fā)展，我們可以期待人工智能在處理復(fù)雜長文本任務(wù)方面的能力將繼續(xù)提升，為用戶提供更加流暢、自然的交互體驗。