LLM 的成功，某種程度上要?dú)w功于 Transformer 架構(gòu)在自然語言處理任務(wù)上的突破。該架構(gòu)最初是為了克服循環(huán)模型的 sequential training 問題而提出的。這些年來，Transformer 已經(jīng)成為 LLM 普遍采用的架構(gòu)。

然而，Transformer 的訓(xùn)練并行性是以低效推理為代價(jià)的：每一步的復(fù)雜度為 O (N) 且鍵值緩存受內(nèi)存限制，讓 Transformer 不適合部署。不斷增長(zhǎng)的序列長(zhǎng)度會(huì)增加 GPU 內(nèi)存消耗和延遲，并降低推理速度。

研究者們一直在努力開發(fā)下一代架構(gòu)，希望保留訓(xùn)練并行性和 Transformer 的性能，同時(shí)實(shí)現(xiàn)高效的 O (1) 推理。針對(duì)這個(gè)問題，此前的方法都沒能同時(shí)實(shí)現(xiàn)這幾點(diǎn)，至少與 Transformer 相比沒有顯示出絕對(duì)的優(yōu)勢(shì)。

現(xiàn)在，微軟研究院和清華大學(xué)的研究者已經(jīng)在這個(gè)問題上取得了重大突破。

論文鏈接：https://arxiv.org/pdf/2307.08621.pdf

在這項(xiàng)工作中，研究者提出了 retentive 網(wǎng)絡(luò)（RetNet），同時(shí)實(shí)現(xiàn)了低成本推理、高效長(zhǎng)序列建模、媲美 Transformer 的性能和并行模型訓(xùn)練，打破了「不可能三角」。

具體來說，RetNet 引入了一種多尺度 retention 機(jī)制來替代多頭注意力，它有三種計(jì)算范式：并行、循環(huán)和分塊循環(huán)表征。

首先，并行表征使訓(xùn)練并行化，以充分利用 GPU 設(shè)備。其次，循環(huán)表征法在內(nèi)存和計(jì)算方面實(shí)現(xiàn)了高效的 O (1) 推理。部署成本和延遲可以顯著降低，同時(shí)無需鍵值緩存技巧，大大簡(jiǎn)化了實(shí)現(xiàn)過程。此外，分塊循環(huán)表征法能夠執(zhí)行高效的長(zhǎng)序列建模。研究者對(duì)每個(gè)局部塊進(jìn)行并行編碼以提高計(jì)算速度，同時(shí)對(duì)全局塊進(jìn)行循環(huán)編碼以節(jié)省 GPU 內(nèi)存。

論文進(jìn)行了大量實(shí)驗(yàn)來對(duì)比 RetNet 和 Transformer 及其變體。實(shí)驗(yàn)結(jié)果表明，RetNet 在 scaling 曲線和上下文學(xué)習(xí)方面始終具有競(jìng)爭(zhēng)力。此外，RetNet 的推理成本與長(zhǎng)度無關(guān)。對(duì)于 7B 模型和 8k 序列長(zhǎng)度，RetNet 的解碼速度是帶鍵值緩存的 Transformers 的 8.4 倍，內(nèi)存節(jié)省 70%。

在訓(xùn)練過程中，RetNet 也能夠比標(biāo)準(zhǔn) Transformer 節(jié)省 25-50% 的內(nèi)存，實(shí)現(xiàn) 7 倍的加速，并在高度優(yōu)化的 FlashAttention 方面具有優(yōu)勢(shì)。此外，RetNet 的推理延遲對(duì)批大小不敏感，從而實(shí)現(xiàn)了巨大的吞吐量。

這些令人驚艷的特質(zhì)讓不少研究者驚呼「好得不可思議」，甚至有人將其比作當(dāng)初「M1 芯片」登場(chǎng)所帶來的變革意義。看來，RetNet 有望成為 Transformer 的有力繼承者。

不過，也有研究者提出疑問：這么優(yōu)秀的表現(xiàn)是否意味著 RetNet 要在某些方面有所權(quán)衡？它能擴(kuò)展到視覺領(lǐng)域嗎？

接下來，讓我們深入了解 RetNet 方法的細(xì)節(jié)。