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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

為了驗(yàn)證ReSA的有效性，研究團(tuán)隊(duì)在多種任務(wù)上進(jìn)行了廣泛的實(shí)驗(yàn)評(píng)估，包括數(shù)學(xué)推理、語(yǔ)言建模和信息檢索。

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

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

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

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

四、效率分析：ReSA帶來(lái)的速度提升

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

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

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

延遲分析還顯示，稀疏估計(jì)和注意力計(jì)算占用了相當(dāng)?shù)臅r(shí)間，這是因?yàn)樗鼈兊膬?nèi)存訪問(wèn)模式類似（一個(gè)與緩存大小/塊大小成比例，另一個(gè)與緩存大小×稀疏率成比例）。而矯正過(guò)程的額外開(kāi)銷相對(duì)較小，在25.6萬(wàn)詞長(zhǎng)度下僅占總延遲的32.7%。

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

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

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

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

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

值得一提的是，ReSA的設(shè)計(jì)是完全無(wú)需訓(xùn)練的，它可以直接應(yīng)用于現(xiàn)有的預(yù)訓(xùn)練模型，無(wú)需任何微調(diào)或重新訓(xùn)練。這使得它特別適合在實(shí)際部署中使用，因?yàn)橛脩艨梢粤⒓聪硎艿叫侍嵘?，而不需要耗費(fèi)大量資源進(jìn)行模型再訓(xùn)練。

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

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

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