你沒有看過的全新版本，Transformer數學原理揭秘
本文經計算機視覺研究院公眾號授權轉載，轉載請聯系出處。
論文地址：https://arxiv.org/pdf/2312.10794.pdf
2017 年，Vaswani 等人發表的 《Attention is all you need》成為神經網絡架構發展的一個重要里程碑。這篇論文的核心貢獻是自注意機制，這是 Transformers 區別于傳統架構的創新之處，在其卓越的實用性能中發揮了重要作用。
事實上，這一創新已成為計算機視覺和自然語言處理等領域人工智能進步的關鍵催化劑，同時在大語言模型的出現中也起到了關鍵作用。因此，了解 Transformers，尤其是自注意處理數據的機制，是一個至關重要但在很大程度上尚未充分研究的領域。
深度神經網絡（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 架構中的參數進行調整成為可能，日后可能會進一步產生實際應用。
