近日,arxiv 上發布了一篇論文,對 Transformer 的數學原理進行全新解讀,內容很長,知識很多,十二分建議閱讀原文。 2017 年,Vaswani 等人發表的 《Attention is all you need》成為神經網路架構發展的一個重要里程碑。這篇論文的核心貢獻是自註意機制,這是 Transformers 區別於傳統架構的創新之處,在其卓越的實用性能中發揮了重要作用。 事實上,這項創新已成為電腦視覺和自然語言處理等領域人工智慧進步的關鍵催化劑,同時在大語言模型的出現中也起到了關鍵作用。因此,了解 Transformers,尤其是自註意處理資料的機制,是一個至關重要但在很大程度上尚未充分研究的領域。 論文網址:https://arxiv.org/pdf/2312.10794.pdf深度神經網路( DNNs)有一個共同特徵:輸入資料依照順序,被逐層處理,形成一個時間離散的動態系統(具體內容可以參考MIT 出版的《深度學習》,國內也被稱為「花書」)。這種觀點已被成功地用於將殘差網路建模到時間連續的動態系統上,後者被稱為神經常微分方程(neural ODEs)。在神經常微分方程中,輸入影像 在時間間隔 (0,T) 上會依照給定的時變速度場 進行演化。因此,DNN 可以看成是從一個 到另一個的流映射(Flow Map)。即使在經典 DNN 架構限制下的速度場中,流映射之間也具有強烈的相似性。 研究者發現,Transformers 實際上是在上的流映射,即 d 維概率測度空間(the space of probability measures)間的映射。為了實現這種在度量空間間進行轉換的流映射,Transformers 需要建立了一個平均場相互作用的粒子系統(mean-field interacting particle system.)。 具體來說,每個粒子(在深度學習情境下可以理解為token)都遵循向量場的流動,流動取決於所有粒子的經驗測度( empirical measure)。反過來,方程式決定了粒子經驗測量的演變過程,這個過程可能會持續很長時間,需要持續關注。 對此,研究者的主要觀察結果是,粒子們往往最終會聚集在一起。這種現像在諸如單向推導(即預測序列中的下一個詞)的學習任務中會特別明顯。輸出量測對下一個 token 的機率分佈進行編碼,根據聚類結果就可以篩選出少量可能的結果。 本文的研究結果表明,極限分佈實際上是一個點質量,不存在多樣性或隨機性,但這與實際觀測結果不符。這明顯的悖論因粒子存在長時間的可變狀態而解決。從圖2 和圖4 可以看出,Transformers 具有兩種不同的時間尺度:在第一階段,所有token 迅速形成幾個簇,而在第二階段(較第一階段速度慢得多),透過簇的成對合併過程,所有token 最終坍縮為一個點。 本文的目標有兩個面向。一方面,本文旨在提供一個從數學角度研究 Transformers 通用且易於理解的框架。特別是,透過這些相互作用粒子系統的結構,研究者可以將其與數學中的既定主題建立具體聯繫,包括非線性傳輸方程式、Wasserstein 梯度流、集體行為模型和球面上點的最佳化配置等。另一方面,本文描述了幾個有前景的研究方向,並特別關注長時間跨度下的聚集現象。研究者提出的主要結果指標都是新的,並且在整篇論文中也提出了他們認為有趣的開放性問題。 本文的主要貢獻分為三個部分。 第 1 部分:建模。本文定義了 Transformer 架構的理想模型,該模型將層數視為連續時間變數。這種抽象方法並不新穎,與 ResNets 等經典架構所採用的方法類似。本文的模型只關注 Transformer 架構的兩個關鍵組成部分:自註意力機制和層歸一化。層歸一化有效地將粒子限制在單位球 的空間內部,而自註意力機制則是透過經驗度量實現粒子之間的非線性耦合。反過來,經驗測量根據連續性偏微分方程進行演化。本文也為自註意引入了一個更簡單好用的替代模型,一個能量函數的 Wasserstein 梯度流,而能量函數在球面上點的最適配置已經有成熟的研究方法。 第二部分:聚類。在這一部分,研究者提出了在較長時間跨度下,token 聚類的新的數學結果。如定理 4.1 表明,在高維空間中,一組隨機初始化在單位球上的 n 個粒子會在時聚成一個點。研究者對粒子集群收縮率的精確描述對此結果進行了補充說明。具體來說,研究者繪製了所有粒子間距離的直方圖,以及所有粒子快要完成聚類的時間點(見原文第 4 節)。研究者還在不假設維數 d 較大的情況下就得到了聚類結果(見原文第 5 節)。 第 3 部分:未來展望。本文主要以開放式問題的形式提出問題,並透過數位觀測加以證實,以此提出了未來研究的潛在路線。研究者首先關注維數 d = 2 的情況(見原文第 6 節),並引出與 Kuramoto 振盪器的連結。接著簡單展示如何透過對模型進行簡單而自然的修改,解決球面最優化相關的難題(見原文第 7 節)。接下來的章節探討了相互作用的粒子系統,這些粒子系統使得對 Transformer 架構中的參數進行調整成為可能,日後可能會進一步產生實際應用。