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

作者：量子位 2023-07-07 19:23:08

人工智能

目前最強(qiáng)的GPT-4也才最大支持一次處理32k token，相當(dāng)于50頁文字。而能夠只用1分鐘看完一本數(shù)萬字小說的Claude，其token數(shù)也不過“才”100k（10萬）。

微軟亞研院最新研究，有點(diǎn)震撼：

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

圖片

什么概念？

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

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

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

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

圖片

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

如何做到？

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

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

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

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

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

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

圖片

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

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

圖片

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

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

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

圖片

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

圖片

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

實(shí)驗(yàn)評估

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

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

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

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

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

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

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

圖片