1000000000！微軟改進(jìn)Transformer一次能記住這么多token了

豐色 發(fā)自 凹非寺

量子位 | 公眾號 QbitAI

微軟亞研院最新研究，有點震撼：

他們開發(fā)了一個Transformer變體，居然將其token拓展到了10億。

什么概念？

目前最強(qiáng)的GPT-4也才最大支持一次處理32k?token，相當(dāng)于50頁文字。

而能夠只用1分鐘看完一本數(shù)萬字小說的Claude，其token數(shù)也不過“才”100k（10萬）。

一次性擴(kuò)展到10億，并且這個數(shù)字理論上其實還是無限的，這不就意味著：

不久的將來，整個語料庫甚至互聯(lián)網(wǎng)都能視為一個序列？

若真如此，簡直不敢想象。

如何做到？

大預(yù)言模型時代，擴(kuò)展序列長度已成為一大關(guān)鍵需求。

然而，現(xiàn)有的方法要么難以解決計算復(fù)雜性，要么搞不定模型表達(dá)力，導(dǎo)致長度很受限。

在此，作者提出一個Transformer變體：LongNet，它應(yīng)用了一種叫做“膨脹注意力（dilated attention）”的機(jī)制，可以隨著距離的增長，讓注意力場（模型感知范圍）呈指數(shù)級擴(kuò)展。

具體而言，dilated attention替代了普通Transformer中的注意力機(jī)制的，其一般的設(shè)計原則是：

讓注意力的分配隨著token之間距離的增長，呈指數(shù)級下降。

如下圖構(gòu)建塊所示，dilated attention由一系列用于建模短期和長期依賴關(guān)系的注意力pattern組成。pattern的數(shù)量可以根據(jù)序列長度進(jìn)行擴(kuò)展。

實驗證明，dilated attention能夠產(chǎn)生線性計算復(fù)雜度和token之間的對數(shù)依賴性，從而解決了注意力資源有限，但每一個token都可訪問的矛盾。

除此之外，作者還引入了多頭dilated attention。如下圖所示，不同頭之間的注意力pattern會因位置的連續(xù)變化而不同。

在最終實現(xiàn)中，LongNet可轉(zhuǎn)換為密集的Transformer，從而無縫集成Transformer現(xiàn)有的優(yōu)化方法（比如內(nèi)核融合、量化和分布式訓(xùn)練）。

而利用線性復(fù)雜度，LongNet還可以跨節(jié)點并行訓(xùn)練，利用分布式算法打破計算和內(nèi)存的限制。

如下圖所示，LongNet通過劃分序列維度在兩個GPU上進(jìn)行并行化訓(xùn)練。隨著設(shè)備數(shù)量的增長，計算和通信成本幾乎仍然是不變的。

因此，這使得LongNet可以在幾乎恒定的運行時間內(nèi)有效地將序列長度擴(kuò)展到10億token（同時無需犧牲較短序列的性能）。

相比之下，普通Transformer還在受困于二次復(fù)雜度帶來的影響。

實驗評估

作者將LongNet與普通Transformer和稀疏Transformer進(jìn)行了比較。

三者之間的體系結(jié)構(gòu)差異只體現(xiàn)在注意力層，模型的序列長度都從2K擴(kuò)展到32K，同時減少batch size以保持每個batch的token數(shù)量不變。

對于LongNet，作者使用w＝{2048,4096,8192,16384,32768}的分段長度，同時擴(kuò)張比率為r＝{1,2,4,6,12}，并實現(xiàn)了稀疏注意力的固定模式，再將稀疏比進(jìn)行調(diào)整為匹配LongNet的計算flops，從而使實驗公平。

（由于計算限制，只能將LongNet擴(kuò)展到32K序列長度）。

下表為以上三個模型在Stack數(shù)據(jù)集上的結(jié)果，衡量指標(biāo)是困惑度，測試了每個模型不同的序列長度。

當(dāng)輸入長度超過模型支持的最大長度時，作者就使用一種最先進(jìn)的語言模型推理外推方法：分塊因果注意（BCA）。

此外，還去除了絕對位置編碼。

結(jié)果表明：

首先，在訓(xùn)練過程中增加序列長度一般會得到效果更好的語言模型。

其次，在長度遠(yuǎn)大于模型支持的情況下，序列長度外推法（推理時）并不適用。

最后，可以看到，LongNet的每一項數(shù)據(jù)都優(yōu)于其他兩個基線，證明了它在語言建模方面的有效性。

除此之外，作者還證明，在擴(kuò)展語言模型上下文長度時，LongNet也是一種更有效的方法。

網(wǎng)友：Memery is all you need

不少人看完這項成果總結(jié)道，說來說去，還是“memery is all you need”（手動狗頭）

有點厲害！非常希望它是真的且可復(fù)現(xiàn)。

不過，也不乏爭議。

比如有人吐槽實驗只將LongNet擴(kuò)展到了32k長度，離10億還差的遠(yuǎn)呢。

還有不少人稱，方法不算很創(chuàng)新，其實就跟稀疏注意力模式差不多。

作者介紹

本論文作者7位，全部來自微軟亞研院。

兩位共同一作：

Jiayu Ding，微軟亞研院實習(xí)生，本碩畢業(yè)于北京郵電大學(xué)計算機(jī)科學(xué)專業(yè)；

馬樹銘，微軟亞研院自然語言計算組研究員，2019年加入，本碩畢業(yè)于北京大學(xué)。

通訊作者為Furu Wei，微軟亞研院自然語言計算組全球研究合伙人，博士畢業(yè)于武漢大學(xué)，曾任IBM中國研究院研究員、香港理工大學(xué)電子計算系研究助理。

論文地址：
https://arxiv.org/abs/2307.02486

參考鏈接：
[1]https://twitter.com/arankomatsuzaki/status/1676765133362675712
[2]https://twitter.com/AlphaSignalAI/status/1677009901334192143?s=20
[3]https://www.reddit.com/r/LocalLLaMA/comments/14rym30/longnet_scaling_transformers_to_1000000000_tokens/