本文經(jīng)計(jì)算機(jī)視覺研究院公眾號(hào)授權(quán)轉(zhuǎn)載，轉(zhuǎn)載請(qǐng)聯(lián)系出處。

論文地址：https://arxiv.org/pdf/2312.10794.pdf

2017 年，Vaswani 等人發(fā)表的 《Attention is all you need》成為神經(jīng)網(wǎng)絡(luò)架構(gòu)發(fā)展的一個(gè)重要里程碑。這篇論文的核心貢獻(xiàn)是自注意機(jī)制，這是 Transformers 區(qū)別于傳統(tǒng)架構(gòu)的創(chuàng)新之處，在其卓越的實(shí)用性能中發(fā)揮了重要作用。

事實(shí)上，這一創(chuàng)新已成為計(jì)算機(jī)視覺和自然語言處理等領(lǐng)域人工智能進(jìn)步的關(guān)鍵催化劑，同時(shí)在大語言模型的出現(xiàn)中也起到了關(guān)鍵作用。因此，了解 Transformers，尤其是自注意處理數(shù)據(jù)的機(jī)制，是一個(gè)至關(guān)重要但在很大程度上尚未充分研究的領(lǐng)域。

深度神經(jīng)網(wǎng)絡(luò)（DNNs）有一個(gè)共同特征：輸入數(shù)據(jù)按照順序，被逐層處理，形成一個(gè)時(shí)間離散的動(dòng)態(tài)系統(tǒng)（具體內(nèi)容可以參考 MIT 出版的《深度學(xué)習(xí)》，國(guó)內(nèi)也被稱為「花書」）。這種觀點(diǎn)已被成功地用于將殘差網(wǎng)絡(luò)建模到時(shí)間連續(xù)的動(dòng)態(tài)系統(tǒng)上，后者被稱為神經(jīng)常微分方程（neural ODEs）。在神經(jīng)常微分方程中，輸入圖像 在時(shí)間間隔 （0，T） 上會(huì)按照給定的時(shí)變速度場(chǎng) 進(jìn)行演化。因此，DNN 可以看作是從一個(gè)  到另一個(gè)的流映射（Flow Map）。即使在經(jīng)典 DNN 架構(gòu)限制下的速度場(chǎng)中，流映射之間也具有很強(qiáng)的相似性。

研究者們發(fā)現(xiàn)，Transformers 實(shí)際上是在上的流映射，即 d 維概率測(cè)度空間（the space of probability measures）間的映射。為了實(shí)現(xiàn)這種在度量空間間進(jìn)行轉(zhuǎn)換的流映射，Transformers 需要建立了一個(gè)平均場(chǎng)相互作用的粒子系統(tǒng)（mean-field interacting particle system.）。

具體來說，每個(gè)粒子（在深度學(xué)習(xí)語境下可以理解為 token）都遵循向量場(chǎng)的流動(dòng)，流動(dòng)取決于所有粒子的經(jīng)驗(yàn)測(cè)度（empirical measure）。反過來，方程決定了粒子經(jīng)驗(yàn)測(cè)量的演變進(jìn)程，這個(gè)過程可能會(huì)持續(xù)很長(zhǎng)時(shí)間，需要進(jìn)行持續(xù)關(guān)注。

對(duì)此，研究者的主要觀察結(jié)果是，粒子們往往最終會(huì)聚集到一起。這種現(xiàn)象在諸如單向推導(dǎo)（即預(yù)測(cè)序列中的下一個(gè)詞）的學(xué)習(xí)任務(wù)中會(huì)尤為明顯。輸出度量對(duì)下一個(gè) token 的概率分布進(jìn)行編碼，根據(jù)聚類結(jié)果就可以篩選出少量可能的結(jié)果。

本文的研究結(jié)果表明，極限分布實(shí)際上是一個(gè)點(diǎn)質(zhì)量，不存在多樣性或隨機(jī)性，但這與實(shí)際觀測(cè)結(jié)果不符。這一明顯的悖論因粒子存在長(zhǎng)時(shí)間的可變狀態(tài)得到解決。從圖 2 和圖 4 中可以看出，Transformers 具有兩種不同的時(shí)間尺度：在第一階段，所有 token 迅速形成幾個(gè)簇，而在第二階段（較第一階段速度慢得多），通過簇的成對(duì)合并過程，所有 token 最終坍縮為一個(gè)點(diǎn)。

本文的目標(biāo)有兩個(gè)方面。一方面，本文旨在提供一個(gè)從數(shù)學(xué)角度研究 Transformers 通用且易于理解的框架。特別是，通過這些相互作用粒子系統(tǒng)的結(jié)構(gòu)，研究者可以將其與數(shù)學(xué)中的既定主題建立具體聯(lián)系，包括非線性傳輸方程、Wasserstein 梯度流、集體行為模型和球面上點(diǎn)的最優(yōu)化配置等。另一方面，本文描述了幾個(gè)有前景的研究方向，并特別關(guān)注長(zhǎng)時(shí)間跨度下的聚類現(xiàn)象。研究者提出的主要結(jié)果指標(biāo)都是新的，并且還在整篇論文中提出了他們認(rèn)為有趣的開放性問題。

本文的主要貢獻(xiàn)分為三個(gè)部分。

第 1 部分：建模。本文定義了 Transformer 架構(gòu)的理想模型，該模型將層數(shù)視為連續(xù)時(shí)間變量。這種抽象方法并不新穎，與 ResNets 等經(jīng)典架構(gòu)所采用的方法類似。本文的模型只關(guān)注 Transformer 架構(gòu)的兩個(gè)關(guān)鍵組成部分：自注意力機(jī)制和層歸一化。層歸一化有效地將粒子限制在單位球 的空間內(nèi)部，而自注意力機(jī)制則是通過經(jīng)驗(yàn)度量實(shí)現(xiàn)粒子之間的非線性耦合。反過來，經(jīng)驗(yàn)度量根據(jù)連續(xù)性偏微分方程進(jìn)行演化。本文還為自注意引入了一個(gè)更簡(jiǎn)單好用的替代模型，一個(gè)能量函數(shù)的 Wasserstein 梯度流，而能量函數(shù)在球面上點(diǎn)的最優(yōu)配置已經(jīng)有成熟的研究方法。

第二部分：聚類。在這一部分，研究者提出了在較長(zhǎng)時(shí)間跨度下，token 聚類的新的數(shù)學(xué)結(jié)果。如定理 4.1 表明，在高維空間中，一組隨機(jī)初始化在單位球上的 n 個(gè)粒子會(huì)在時(shí)聚成一個(gè)點(diǎn)。研究者對(duì)粒子集群收縮率的精確描述對(duì)這一結(jié)果進(jìn)行了補(bǔ)充說明。具體來說，研究者繪制了所有粒子間距離的直方圖，以及所有粒子快要完成聚類的時(shí)間點(diǎn)（見原文第 4 節(jié)）。研究者還在不假設(shè)維數(shù) d 較大的情況下就得到了聚類結(jié)果（見原文第 5 節(jié)）。

第 3 部分：未來展望。本文主要以開放式問題的形式提出問題，并通過數(shù)字觀測(cè)加以證實(shí)，以此提出了未來研究的潛在路線。研究者首先關(guān)注維數(shù) d = 2 的情況（見原文第 6 節(jié)），并引出與 Kuramoto 振蕩器的聯(lián)系。然后簡(jiǎn)要展示了如何通過對(duì)模型進(jìn)行簡(jiǎn)單而自然的修改，解決球面最優(yōu)化相關(guān)的難題（見原文第 7 節(jié)）。接下來的章節(jié)探討了相互作用的粒子系統(tǒng)，這些粒子系統(tǒng)使得對(duì) Transformer 架構(gòu)中的參數(shù)進(jìn)行調(diào)整成為可能，日后可能會(huì)進(jìn)一步產(chǎn)生實(shí)際應(yīng)用。