“大模型設計迎來伽利略時刻。正如比薩斜塔實驗推動了現(xiàn)代物理學的發(fā)展，我們開發(fā)的受控合成預訓練平臺揭示了大模型架構(gòu)的真正極限。這可能是一個轉(zhuǎn)折點，將大模型的研究劃分為‘之前’和‘之后’?！?清華本科校友、美國麻省理工學院博士畢業(yè)生、Meta 研究員朱澤園在 X 上寫道。
這則推文介紹了一個名為“語言模型物理學”的長期項目。在這則推文發(fā)布之前，當?shù)貢r間 5 月 2 日關(guān)于這一長期項目的一篇論文上線 [1]。
在這篇論文中，朱澤園和同事提出一種名為“Canon 層”的輕量級結(jié)構(gòu)。
“Canon”原本是一個音樂術(shù)語，它指的是一種復調(diào)音樂的作曲技法，其特點是讓多個聲部以相同的旋律和不同時間進入，從而形成層疊交錯的效果。
而本次提出的“Canon 層”和上述音樂技法有著異曲同工之妙，它能促進相鄰 tokens 之間的水平信息流動，提升推理深度、推理廣度、知識容量和操控能力，并能在跨多種架構(gòu)的層次結(jié)構(gòu)學習上實現(xiàn)突破。
通過開展這些合成任務，研究團隊實現(xiàn)了架構(gòu)性能的嚴格分離，讓有著 13 億參數(shù)和 1000 億 tokens 數(shù)據(jù)量的學術(shù)級預訓練實驗中的噪聲與隨機干擾因素，得到了有效消除。
研究人員在論文中表示，Canon 層能夠顯著提升傳統(tǒng)弱勢架構(gòu)的性能，例如它能將無位置編碼（NoPE，No Positional Encoding）的 Transformer 提升至旋轉(zhuǎn)位置編碼（RoPE，Rotary Positional Encoding）水平，使線性注意力（GLA，Generalized Linear Attention）達到 Mamba 級表現(xiàn)，同時大幅降低對旋轉(zhuǎn)位置編碼的依賴，從而有效改善長程泛化能力。（注：Mamba，是此前由其他研究團隊提出的一種高效序列建模架構(gòu)，下文的 Mamba2 是 Mamba 的升級版。）
據(jù)介紹，Canon 層能夠計算附近 tokens 表示的加權(quán)組合，并能無縫集成到 Transformer、線性注意力、狀態(tài)空間架構(gòu)或任何通用序列模型中。
本次研究還表明，線性模型推理深度不足的原因在于 tokens 的壓縮過程和檢索過程效率低下。同時，本次研究中的合成任務也驗證了受控可預測基準在架構(gòu)評估中的核心價值。
與此前經(jīng)典的學術(shù)級別實驗場景相比，本次研究中的合成實驗平臺能夠更清晰地揭示模型的真實能力和擴展行為。
憑借無限高質(zhì)量的數(shù)據(jù)，研究團隊希望該框架能夠預測架構(gòu)如何隨著訓練流程的改進而演變。例如，預測該如何通過更好的數(shù)據(jù)整理或基于強化學習的后訓練，解鎖更深層次的推理能力和層次化推理能力。
研究團隊在論文中表示，他們更建議將智能分解為原子級組件，比如分為推理深度和推理廣度，并通過構(gòu)建高度可控的合成預訓練任務，來獨立地分離和評估這些能力。
這樣一來，在理想化的實驗環(huán)境之下，就能清晰地界定不同架構(gòu)的性能優(yōu)勢與擴展?jié)摿Α?/p>
據(jù)研究人員介紹，本次方法通過實現(xiàn)單一技能的評估，消除了真實世界預訓練數(shù)據(jù)中的混雜因素。例如，本次方法既可以針對架構(gòu) A 在推理深度上是否優(yōu)于架構(gòu) B 開展嚴謹對比，還能確保修改不會降低其他能力。
通過隔離固有的架構(gòu)偏差，他們在合成預訓練任務中揭示了那些在傳統(tǒng)設置中往往會被噪聲掩蓋的特性。
研究中該團隊打造了一個合成訓練環(huán)境，并引入了五項合成預訓練任務，以便系統(tǒng)地評估語言模型中的關(guān)鍵認知能力。
通過這一實驗平臺研究人員揭示了在小規(guī)模場景下的架構(gòu)趨勢：其一，線性注意力模型在各類任務中表現(xiàn)持續(xù)滯后；其二，Mamba2 這樣的遞歸架構(gòu)在依賴記憶的任務中優(yōu)勢顯著；其三，標準 Transformer 在復雜推理任務中占據(jù)主導型的優(yōu)秀表現(xiàn)。
對于分析不同架構(gòu)下的固有偏差和可擴展性來說，這些受控基準測試能為其提供一個較為嚴謹?shù)目蚣?。研究人員指出，Transformer 在層內(nèi)缺乏水平信息流，即使在關(guān)聯(lián)回憶這樣的簡單任務上也會出現(xiàn)效率低下的情況。
另據(jù)研究人員在論文中表示，Canon 層可以靈活地嵌入網(wǎng)絡的不同位置，其中：
- Canon-A：插入注意力機制之前；
- Canon-B：嵌入注意力機制內(nèi)部；
- Canon-C：置于多層感知機（MLP，Multilayer Perceptron）之前；
- Canon-D：整合至多層感知機內(nèi)部。
雖然 Canon 層的實現(xiàn)方式有很多種，即便使用簡單的隨機平均操作也能取得顯著效果。但是，本次研究聚焦于采用可訓練的一維線性卷積核（核大小為 4），這種方法只需使用輕量計算、而且易于集成，只需少量代碼即可適配任意架構(gòu)。
在研究團隊的實驗平臺中，Canon 層分別能將推理深度提升 200%-400%、將推理廣度提升 30%、將知識操作長度提升 30%，并且只需要極小的計算開銷能夠?qū)崿F(xiàn)這些提升，而這主要依賴于更優(yōu)的層次化學習動態(tài)。
將 Canon 層集成到模型中，能夠提高無位置編碼模型的性能，其性能與“旋轉(zhuǎn)位置編碼模型+Canon 層”的組合相當，甚至更勝一籌。在性能上，Canon 層優(yōu)于像線性偏置注意力（ALiBi，Attention with Linear Biases）或混合線性偏置注意力（H - Alibi，Hybrid Attention with Linear Biases）這樣的位置修正方法，并且能夠降低或消除旋轉(zhuǎn)位置編碼的使用，從而顯著提升模型的長程泛化能力。
研究人員在論文中表示，Canon 層通過不同子層位置的累積式作用機制來實現(xiàn)性能提升，其效果獨立于注意力或多層感知機組件。殘差連接提高了訓練效率，在不影響穩(wěn)定性的前提下，幾乎不需要進行參數(shù)調(diào)整。
通過提高訓練效率和穩(wěn)定性，Canon 層還能恢復在門控多層感知機或混合專家架構(gòu)中損失的部分知識承載能力。
研究中，該團隊還將線性注意力與 Canon 加以結(jié)合。通過此，Canon 層能夠顯著提升門控線性注意力的推理能力。其中，推理深度能從單跳（1-hop）被擴展至四跳（4-hop），推理廣度與知識操作長度能夠?qū)崿F(xiàn)雙倍增長。最終可以使門控線性注意力達到與狀態(tài)空間模型架構(gòu) Mamba2 相當?shù)男阅埽⒛茉?Brevo 等任務中實現(xiàn)反超。
據(jù)介紹，Mamba2 的性能得益于其內(nèi)置的一維卷積（conv1d）機制，該機制相當于一個應用于選定坐標的非線性 CanonB 層。移除一維卷積機制之后，其性能下降到與門控線性注意力模型相當?shù)乃健６褂猛暾?Canon 層替代它則能進一步提升效果，這凸顯了在結(jié)構(gòu)化狀態(tài)空間模型設計中水平信息流的重要性。
通過消融研究，該團隊發(fā)現(xiàn) Canon 層在集成點和殘差鏈接上的選擇會影響到 Mamba2 的性能。作為一種使用學習線索來初始化 Transformer 自注意力權(quán)重的方法，模仿初始化（Mimetic initialization）雖然能在長度泛化上實現(xiàn)優(yōu)化，但卻會給短上下文任務造成損害，而這恰恰凸顯多樣化預訓練環(huán)境的必要性。
研究中，該團隊還開展了架構(gòu)比較。他們分別在旋轉(zhuǎn)位置編碼、無位置編碼、Mamba2 和門控線性注意力中同時使用完整的 Canon 層。結(jié)果表明，在層次推理任務中，完整 Transformer 的表現(xiàn)優(yōu)于線性模型，其推理深度達到了線性模型的兩倍。
研究人員指出，線性模型在推理深度方面存在一定的局限性，這種局限性源于壓縮與檢索過程中累積的誤差，而非記憶容量的不足。而將滑動窗口 Transformer 與線性模型相結(jié)合的混合架構(gòu)，則能為深度推理任務提供可擴展的解決方案。
與此同時，他們還開展了學術(shù)級別的真實預訓練。具體來說，他們在 1000 億個 tokens 上訓練具有 4096 個上下文長度的 13 億參數(shù)模型，結(jié)果顯示噪聲較大且分辨率有限，這使得大多數(shù)架構(gòu)差異在統(tǒng)計上并不顯著。
盡管如此，課題組也收獲了一些發(fā)現(xiàn)：
首先，在檢索任務繁重的任務上，線性架構(gòu)即使配備了 Canon 層，也始終落后于完整的 Transformer 架構(gòu)。
其次，Canon 層能夠顯著提升門控線性注意力和無位置編碼，使門控線性注意力能夠達到與 Mamba2 相當?shù)乃?，以及使無位置編碼能夠達到與旋轉(zhuǎn)位置編碼相當?shù)乃?。然而，移?conv1d 會使 Mamba2 性能下降到與門控線性注意力相當。
再次，即使在只有 100 個 tokens 的短文本中，所有模型在處理兩跳推理任務時都面臨一定的困難，這凸顯了學術(shù)級別預訓練的局限性。
最后，通過減少或去除旋轉(zhuǎn)位置編碼這一做法，尤其是在添加了 Canon 層的情況下使用這一做法，能在不影響整體性能的前提下，提升長上下文的泛化能力。
總之，Canon 層從根本上改善了跨不同架構(gòu)的水平信息流，從而實現(xiàn)了更深入的推理和高效的可擴展性。未來，該團隊希望借助無限高質(zhì)量的數(shù)據(jù)，通過本次成果來預測未來架構(gòu)的演進方向，例如通過更完善的數(shù)據(jù)策展或基于強化學習的訓練后優(yōu)化，從而釋放更深層次的推理與分層推斷能力，進而推動大模型訓練流程的持續(xù)進步。
參考資料：
1.論文地址：https://papers.ssrn.com/sol3/papers.cfm?abstract_id=5240330
https://www.linkedin.com/in/zeyuan-allen-zhu/details/education/
https://x.com/ZeyuanAllenZhu
http://zeyuan.allen-zhu.com/index.php
運營/排版：何晨龍
熱門跟貼