Preface
前言
楊松琳專注于線性注意力機制、機器學習與大語言模型交叉領域，聚焦高效序列建模的硬件感知算法設計，圍繞線性注意力、循環(huán)神經網絡優(yōu)化開展研究，并取得多項研究成果，多篇論文被ICLR 2025、NeurIPS 2024等頂會收錄。
此外，她還開源了Flash-Linear-Attention項目，已在 GitHub 擁有 2k 星標。通過線性注意力與硬件優(yōu)化技術，將 Transformer 模型的注意力計算復雜度大幅降低，顯著提升長序列任務（如文本、圖像處理）的訓練推理速度與內存效率。
Background
背景介紹
為什么要做基礎模型新架構研究？
今天想和大家分享的主題是下一代基礎模型的架構展望。今天的這個分享不會涉及太多特別細節(jié)的內容，主要是一個比較high level和general的overview。
首先，從背景來說，序列建模本身涉及的任務非常廣，比如文本就是一個典型的 sequence modeling 任務，音頻、視頻也可以用序列建模的方法來處理，包括強化學習中的決策序列，DNA 和蛋白質的序列等等，甚至很多其他任務都可以看作是序列建模的問題。而 Transformer 其實就是一個非常典型的序列建模模型。
之前很多人的觀點是，基礎模型的架構其實沒必要換，大家更關注數據，覺得數據比架構重要得多。因為過去在堆計算資源、堆數據的情況下，Scaling Law 還沒有撞到那個墻，這種情況下，大公司就沒有動力去研究新的架構，而是直接通過增加算力和數據規(guī)模讓模型變得更強。所以，一直以來很多人都認為架構沒有那么重要，數據才是關鍵。
但現(xiàn)在的情況不太一樣了，比如大家可能比較熟悉的Ilya，他去年在一次 keynote 里提到，他覺得預訓練已經撞墻了。現(xiàn)在主要的 Scaling 方向其實是在 Test Time Scaling，比如像 DeepSeek-R1，它有很長的思維鏈，而在這種情況下，Transformer 在長文本上的效率問題就會被進一步放大。所以，在這個背景下，新架構的研究就顯得越來越重要。
研究新的架構不僅有可能改善 Test Time Scaling 的效率問題，也可能在 Training Time Scaling 方面帶來更多的想象空間。
比如，如果現(xiàn)在數據和計算資源的邊際效益遞減，其實可能是因為架構本身已經到了瓶頸。
如果我們去設計一些更expressive的架構，會不會就能進一步提升模型的表現(xiàn)？這其實是一個很值得探索的問題。
Component?????????
序列模型的核心組件
序列建模通常包含兩個核心部分:Token Mixing和ChannelMixing。
Token Mixing 主要是在序列維度上面做一個混合
Channel Mixing 是在特征維度上面進行一個混合
因為這個輸入就是序列長度乘特征長度，所以只能在這兩個維度上進行混合，這兩個部分可以說是非常本質的。
之前有一個叫 MetaFormer 的工作，它把整個 Transformer 的架構抽象成兩個模塊： 一個是 Token Mixer 一個是 Channel Mixer Token Mixer 可以換成各種在序列維度上做混合的 layer，而 Channel Mixer 也可以換成在特征維度上做混合的 layer。Token Mixing
我們先來看 Token Mixing，目前大家主要還是用全局的 Full Attention，不過這里面其實還是有兩個問題。
第一個問題是 RoPE，RoPE 是一種相對位置編碼，它在長度外推上有一些局限性，沒法很好地進行長度外推。而且 RoPE 和Transformer 本身的表達能力也不夠強，所以需要借助一些 Chain-of-Though 來增強它的表達能力。
另外一個方面是效率問題，Self-Attention 機制的時間復雜度是平方級的，這在長序列建模任務中有很大的問題。比如在視頻生成任務中，如果直接在幾百萬長度的序列上跑 Full Attention，計算成本會非常高。像視頻這種數據，隨便幾分鐘的視頻就能達到幾百萬的序列長度，在這種情況下，Attention 機制的平方復雜度即便有 Flash Attention 來優(yōu)化，依然是不容忽視的。另外還有一個線性的空間復雜度問題。在解碼過程中，KV Cache 的大小會隨著解碼長度的增加而不斷增長，很容易導致 out of memory 的問題。
RoPE后已有的改進
首先，我們來講 Position Encoding。也就是說 RoPE 還能換成什么別的東西？
之前有個方法叫 NoPE，它的做法是直接把 Position 空間去掉，但實際效果并不好。
另外，也有很多方法嘗試改進 RoPE，讓它可以做長度外推。但最后大家發(fā)現(xiàn)，直接增大 RoPE 的 base，然后在長文本上繼續(xù)訓練，整體上的超長文本能力會更好一些。
上下文相關的位置編碼——看好的方向
那接下來，Position Encoding 未來可能的發(fā)展方向是什么？
我個人比較看好上下文相關的 Position Encoding。RoPE 其實是上下文無關的位置編碼，它的值和前后詞的內容沒有關系。而有研究提出了一種Contextual Position Encoding，在一些任務上，比如 selective Copy、counting、Flip-Flop，比起 RoPE 這種相對位置編碼，效果會更好一些。
然后最近 ICLR 25 也有一些關于位置編碼的改進工作，比如我們這篇Scaling Stick-Breaking Attention，我也是作者之一。
論文鏈接：https://arxiv.org/abs/2412.06464 https://arxiv.org/abs/2412.06464
在這篇工作里，我們提出了一個完全不需要 RoPE 的編碼方式，核心思路是改進 Attention算法。之前的全局 Softmax Attention 是位置不敏感的，而我們把它換成從右到左的 Sigmoid Attention，它本身的累乘機制就自帶位置敏感性，所以不需要額外的位置編碼?？梢园堰@種自帶的位置敏感性，看作是一種上下文相關的位置編碼。然后下面有幾個公式，核心的思想就是從右到左的 Sigmoid 計算，即 (1 - sigmoid) 的累乘。
實驗結果表明，這種Stick-Breaking Attention 在性能上是穩(wěn)壓 RoPE 的。
另外的話，還有一篇 Mila 的工作，叫做Forgetting Transformer，它其實可以看成是 ALiBi 的一個動態(tài)版本。ALiBi 的話，大家應該也比較熟悉，它是一種位置編碼，但是它的衰減方式在長文本任務上容易出問題，所以現(xiàn)在大家基本都不用 ALiBi，而是更多選擇 RoPE。而Forgetting Transformer的核心想法，就是把 ALiBi 的衰減因子從一個固定的常數變成一個動態(tài)值。它通過類似 Forgetting Gate這種機制來計算這個衰減值。因為這個衰減是跟上下文相關的，所以它也可以看成是一種上下文相關的位置編碼。
高效注意力機制
然后，除了位置編碼之外，就是注意力機制的一些改進了。主要的話，還是有三種機制：
第一種是線性注意力，這部分的話，之前我其實也寫過一些內容，包括 tutorial video，還有 slides，大家可以到我主頁上去看，這里就不細講了，簡單提一下。像線性注意力，它基本上可以分成幾類，比如加衰減率的模型，像MSR 做的RetNet，然后是Lightning Attention，就是MiniMax-01里面用的注意力機制。然后像Mamba-2，它其實就是一個線性注意力，再加上一個動態(tài)的衰減。而 Gating Linear Attention 這個是我之前在 ICML 的一個工作，它也是一個動態(tài)衰減的線性注意力。另外最近比較流行的一種線性注意力機制叫做Test-time 在線學習的線性注意力。最近這類方法比較流行，比如DeltaNet、Test-time Training、Titans。像Titans這篇工作，大家最近應該能經常在微信公眾號、三大頂會上面看到。然后像RWKV-7、Gated DeltaNet這些都是表達能力更強的在線學習線性注意力的一些變體。如果大家對這些更感興趣的話，可以去看我主頁的 slides。
然后除了線性注意力，還有一種東西叫做稀疏注意力。稀疏注意力一般分為靜態(tài)稀疏和 動態(tài)稀疏。靜態(tài)稀疏指的是那個稀疏的 pattern 是事先給定的，它在計算過程中不會發(fā)生變化。像這種靜態(tài)稀疏的一些經典工作，比如 BigBird，這其實是谷歌之前做的一個工作。然后 StreamingLLM，它是用一個 sliding window，再取最前面的那個 attention sink。動態(tài)稀疏 就是模型會根據不同位置的信息，動態(tài)地改變那個稀疏模式。最近比較火的就是 DeepSeek 和 Kimi，他們兩個都同時出了兩個動態(tài)稀疏機制。第一個是 Native Sparse Attention，另一個是 MoBA。
在這兩種方法之外，我們還會碰到一種叫做混合注意力的一些模型，它主要可以分為兩類：
一種是在層間混合；
另一種是在層內混合。層間混合的意思是，在不同的層里有不同的注意力機制，比如有的層使用 Linear Attention，有的層使用 Softmax Attention。而 層內混合 則是同一層中會同時使用 Sliding Window Attention 和 Linear Attention 等不同的機制。
- 線性注意力
這里我展示了一些來自我主頁的 slides，基本上想表達這都是一些線性注意力在硬件上非常高效的RNN。
它通過一些高效的操作，比如外積，來更新狀態(tài)，而外積有很多良好的性質，這也讓其在硬件上訓練時非常高效。具體細節(jié)這里就不再展開了。另外，加Decay的方法是在 Hidden State 前加一個 Gamma（Gamma 介于 0 和 1 之間）。這樣做的效果是，逐步忘掉之前的信息，更加關注當前的狀態(tài)。
這種機制其實就是 Lightning Attention 在 MiniMax-01 中使用的那種方法，是一個相對簡單的機制。
另外，還有一種方法是 test-time 訓練模型。這種方法的目標函數其實是在線學習的目標函數，而隱狀態(tài)的更新是通過在線梯度下降來進行的。例如 DeltaNet，它通過在線優(yōu)化一個線性回歸損失來更新模型的參數，通過一次梯度下降，進而得到 DeltaNet 的更新公式。
以上這些方法都是典型的線性注意力方法，它們的基本思想在之前的工作中已經有了一些總結，具體細節(jié)可以參考我的主頁。
2. 稀疏注意力
稀疏注意力可以分為靜態(tài)稀疏和動態(tài)稀疏。
靜態(tài)稀疏的模式是事先設定好的，并且不會隨著輸入的變化而變化。通常這些方法因為給定了很強的先驗，所以它們的表達能力相對較弱?，F(xiàn)今，這類方法已經不太被使用了。
打開網易新聞 查看精彩圖片
動態(tài)稀疏是近年來比較流行的一種方式。它的特點是稀疏模式是由每個 token 動態(tài)選擇的。動態(tài)稀疏的主要挑戰(zhàn)在于，由于其稀疏模式非常靈活，沒有固定規(guī)則，因此可能會出現(xiàn)不連續(xù)的稀疏模式，而不連續(xù)的模式在硬件上會比較慢。為了應對這一問題，近期的研究通過將選擇分為 block 來進行優(yōu)化。因為一個 block 通常包含多個連續(xù)的 tokens，所以它的讀取是連續(xù)的，這樣有助于提高硬件效率。
接下來簡單介紹一下 NativeSparseAttention 的思想，一共有三個分支（compression/block selection/sliding window）。它將整個序列劃分成若干個 chunk，每個 chunk 首先通過 compression 將其壓縮成一個 token 表示。然后使用新的橙色的querytoken對這些 compressed tokens 進行注意力操作，得到一個 compressed attention的output。在這個過程中，注意力分數會用來選擇 top-k 的 attention blocks，然后在這些 blocks 上執(zhí)行 sparse attention。最終，再補充一個 sliding window attention，這就是 Native Sparse Attention 的主要做法。
Native Sparse Attention 之所以被稱為原生稀疏注意力，是因為它可以直接進行預訓練，并且有一些假設：每個 query group 下的每個 head 選擇相同的 KV block，這樣可以避免不同的 head 重復讀取不同的 KV block，從而減少 IO 開銷。
此外，Native Sparse Attention 還可以將每個位置的 attention heads 合并成一個矩陣乘法，這對硬件加速有很大幫助。
一個有趣的發(fā)現(xiàn)是Native Sparse Attention 甚至在一些任務上能超過全局注意力 的效果。傳統(tǒng)上大家認為稀疏注意力只是掉點，也就是效果只能保持，無法提高，但 Native Sparse Attention 通過學習自己的注意力分布，能夠達到不輸于全局注意力的效果。這表明，通過設計新的架構并進行預訓練，可能會得到比現(xiàn)有架構更好的效果。
MoBA 的思想與 NativeSparseAttention 非常相似。它的 chunk pooling 使用的是 mean pooling，不像 NSA 中那樣有多個分支，因此不會引入額外的參數，也不強制選擇相鄰的 blocks。它也比較兼容 Flash Attention 的 kernel，但在硬件訓練效率上不如 Native Sparse Attention。不過，MoBA 通過預訓練，能夠實現(xiàn)與 full attention 相比完全不遜色的性能。
3. 混合注意力
接下來就是混合注意力。
前面我們講到，它有很多種混合方法，基本上就是說不同的層就用不同的混合方式，也就是說在不同層之間我們用不同的注意力機制。大部分層都用的是高效的注意力機制，只有少數層用的是全局的 softmax attention。
舉個例子，通常我們會有 7/8 層用線性注意力，而 1/8 層用全局 softmax attention。
這樣做的好處就是，我們可以把 7/8 層 的 KV cache 直接省下來，因為線性注意力的 hidden state 可以表示為 RNN，它的 KV cache 基本上可以忽略不計。
這方面比較典型的例子，比如 Jamba 和 MiniMax-01。他們的做法就是每七層使用 Lightning attention。前面提到過，Lightning attention 其實就是最簡單的線性注意力，再加一個衰減項。這里的 m 通常取的是 7，所以形成了 7:1 的比例，這個比例在實際中是比較常見的。
前面我們也說過，這種方法需要存儲 1/8 層的 KV cache。如果層間混合做得更極致的話，可能只需要存儲一個 KV cache。也就是說，前半部分的層采用高效的注意力機制，然后它算出來的 KV cache 就可以用于后面所有的全局注意力層。這樣的話，首先能節(jié)省預填充的時間，因為在預填充時，我們只需要計算前面的高效層，得到 KV cache 就好了，不需要計算后面的全局注意力層，我們只需要存儲中間那一層的 KV cache，后面的層可以共享這個中間層的 KV cache。
這篇就叫 YOCO，也就是 You Only Cache Once，意思是我只存一層的KVcache，然后做層間混合。
相比之下，注意力的混合方法主要是層內混合，比如 Jamba 這種比較典型的層內混合方法，它就是在一層里同時用 sliding window attention 和 linear attention。具體來說，它先跑一個 sliding window attention，然后在此基礎上再加一個 linear attention，這樣就能在同一層里實現(xiàn)兩種不同的注意力機制。
Channel Mixing
再來看注意力機制的 token mixing，它主要是在位置編碼和注意力機制本身上做一些研究。而 channel mixing 一般就是 FFN 層，F(xiàn)FN 層可以被看成是一種key-value形式的associative memory。它的參數可以理解為存儲了一些全局知識，所以大家普遍認為 FFN 和知識存儲是相關的。那么很自然的一個思路就是，我想增大 FFN 的容量，但又不想大幅增加計算量，于是就有了 Mixture of Experts（MoE）模型。現(xiàn)在 MoE 結構基本上成了主流架構，很多支線模型都會帶一些 MoE 組件，因為它是目前非常高效的擴展整體參數規(guī)模的方式。
MoE 還是有很多可以做的東西。
比方說，UltraMem 這篇文章就是用張量分解進一步加速 MoE 的 inference。它提出了一個叫做 Memory Layers at Scale 的方法。
前面提到，F(xiàn)FN layer 可以理解成一個 memory layer，是用來存全局知識的。而這篇文章的做法就是用了一層非常非常大的 memory layer，來存儲更多的信息。因為這層很大，所以他們不想做全局更新，而是通過一些高效的檢索算法，找到激活的部分，只對這些部分進行查詢和更新。這樣的話，就能在擁有一個超大 memory layer 的同時，不會讓整體的計算復雜度變得不可控。
Thinking
未來方向預測
前面主要講了 token mixing 和 channel mixing 能做的一些事情，那架構層面還能有哪些新的思路呢？
目前的模型基本上都是基于 next word prediction，但能不能探索一些非 next word prediction 的方法？
比如，最近有很多研究在做 diffusion large language model，用擴散來做文本生成。還有像 Meta 去年年底推出的一個模型，它做的是 next concept prediction，而不是直接做 next word prediction。此外，還有一些研究在做 Multi-token Prediction，比如 DeepSeek-V3 也有相關的模塊。
另外，現(xiàn)在大家特別關注 test-time scaling，而 test-time scaling 一般是通過長的思維鏈來做的。這種方式可以看作是水平方向的 scaling，那能不能在垂直方向上也做 test-time scaling？
這里的垂直方向指的是layer，如果從這個維度來做，一種比較自然的思路就是 Universal Transformer。那么，在 test-time scaling 的背景下，這類方法能不能有更大的應用空間？這其實也可以歸為 latent reasoning，因為它不像 Chain of Thought 那樣顯式地依賴推理鏈，而是通過隱式的方式進行推理，因此被稱為 latent reasoning。
另外的話，就是關于 memory 機制的進一步探索。之前，F(xiàn)FN 的權重基本上都是靜態(tài)的，那能不能讓這些權重隨著輸入的變化而動態(tài)調整？這樣的話，其實就有點類似于 fast weight 這種比較經典的概念。然后還有一個方向，就是我們能不能加一些顯式的外部記憶？比如 Transformers with Explicit Memory 這篇文章，它就在模型外部外掛了一個很大的知識庫，來作為 KV memory。
最后的話，我們前面講的 token mixing 和 channel mixing 是完全分開的，那能不能把它們結合在一起，直接做一個統(tǒng)一的更新呢？這些都是可以繼續(xù)探索的方向。然后整體來看，在架構層面，其實還有很多優(yōu)化的空間。
熱門跟貼