為什麼脈絡無法就這樣拉長
叫一個在 4k 詞元上訓練的 RoPE 模型去讀 16k,品質會直接跳崖。原因可回溯到第一篇:慢的 RoPE 頻率如今旋轉到模型訓練時從未見過的角度,於是遙遠的位置看起來像它沒有表徵可對應的亂碼。延展脈絡長度(context length)因此不是一個設定旗標——而是要把位置訊號保持在模型能詮釋的範圍內的問題。
這裡有三個不同的子問題,把它們搞混會浪費好幾週。延展(extension)是教模型處理比訓練時更大的脈絡視窗(context window)。串流(streaming)是無限期生成而不讓快取或品質爆掉。分散(distribution)則是把單一龐大序列硬塞進多張 GPU。接下來三節逐一處理。
YaRN:重縮頻率
YaRN(Yet another RoPE extensioN)藉由選擇性地重縮 RoPE 的旋轉頻率來延展脈絡。其洞見是:並非所有頻率都該一視同仁——慢的、長程的頻率才是快超出範圍的那些,所以把它們內插到模型已經理解的角度,同時讓快的、局部的頻率——負責解析鄰近字序的那些——幾乎原封不動。這種逐頻率、考量波長的處理,正是 YaRN 區別於樸素線性內插之處。
RoPE 每个维度的旋转角与波长——YaRN 按因子 s 重新缩放位置,对高频维度的缩放强于低频维度。
回報是效率:YaRN 只需在少量長資料上做短暫微調,就能把脈絡推到遠超訓練長度——比重新訓練便宜好幾個數量級。它是許多「32k」與「128k」開源變體背後的標準配方,那些變體通常就是一個基礎模型加上 YaRN 式的重縮與一段短暫的適應訓練。
注意力匯點:無止盡串流的祕密
假設你想要一個永不停止的聊天機器人,於是在 KV 快取裡只留最近詞元的滑動視窗、逐出最舊的。出人意料的是,品質會在最開頭那幾個詞元被逐出的瞬間崩潰——儘管它們遙遠且看似無關。這就是注意力匯點(attention sink)現象:模型學會把多餘的注意力機率傾倒到開頭幾個位置,當成一種無作用操作,因為 softmax 強迫權重總和為一,總得有東西吸收這份餘額。
交互式自注意力:点击一个词会高亮它所关注的词元。
解法便宜得令人愉快:永遠把開頭那幾個「匯點」詞元連同滑動視窗一起留在快取裡。保住匯點後,一個在幾千詞元上訓練的模型,就能串流數百萬詞元而不退化、也不需任何微調。它與滑動視窗注意力天然成對——視窗處理局部內容,匯點穩定 softmax。
環形注意力:一條序列,多個裝置
即使快取很小,一條百萬詞元序列的啟動值也可能超過任何單張 GPU 的記憶體。環形注意力(ring attention)的解法是把序列切成多個區塊、每個裝置一塊,排成一個邏輯環。每個裝置為自己那塊查詢計算注意力,同時鍵/值區塊沿環在裝置間傳遞。由於下一塊的通訊與當前塊的計算重疊,網路傳輸大致被有用的工作藏了起來。
注意力是对各键值块求和,因此当块沿环旋转时,每个设备都能累加自己那块的贡献——精确的 softmax,无需任何单个 GPU 容纳整条序列。
美妙的性質是:最大脈絡隨裝置數量縮放,而非任一張的記憶體——往環裡加節點,可定址的序列就線性增長,原則上沒有上限。再搭配每個裝置上 FlashAttention 式的局部核心,環形注意力就是各團隊在訓練與長文件推論中觸及數十萬到數百萬詞元脈絡的方法。