當(dāng)年調(diào)試兩層 Transformer 時我踩過一個非常典型的坑把輸入句子里幾個詞的順序打亂之后模型輸出竟然幾乎沒變。查了半天才意識到問題不在實現(xiàn)而在位置信息。沒有位置編碼的注意力機制本質(zhì)上是對稱的它不關(guān)心詞與詞之間隔了多遠只關(guān)心內(nèi)容像不像。短句上這種毛病不明顯序列一長就徹底露餡。后來我開始逐個翻主流開源語言模型的訓(xùn)練和推理代碼發(fā)現(xiàn)這幾年位置編碼幾乎都從“加法偏置”換成了旋轉(zhuǎn)位置嵌入Rotary Position EmbeddingRoPE。RePE 不再把位置當(dāng)成一個加在詞向量上的噪聲而是當(dāng)成一個作用在 Query 和 Key 上的旋轉(zhuǎn)變換。這個設(shè)計有個非常優(yōu)雅的結(jié)果注意力分?jǐn)?shù)天然只依賴兩個 token 之間的相對距離而不依賴它們各自的絕對位置。這篇文章不談空泛結(jié)論只講幾件實際用得上的事為什么 Transformer 必須有位置編碼、RoPE 的旋轉(zhuǎn)矩陣到底做了什么、怎么從零手寫一遍能跑的 PyTorch 實現(xiàn)以及在實際訓(xùn)練和推理里 RoPE 最常見的那幾個坑該怎么繞。1. 先講動機沒有位置編碼的 Transformer 是什么樣子的1.1 自注意力的排列不變性Transformer 的核心是自注意力。給定一組輸入向量每一層的注意力輸出是“對整條序列做加權(quán)求和”權(quán)重由兩兩之間的點積決定Attention(Q, K, V) softmax(QK^T / sqrt(d)) · V這里面有一個容易被忽略的性質(zhì)如果把輸入序列的打亂順序再喂進去輸出也會跟著打亂但每個位置的輸出值本身不變。用術(shù)語說就是 permutation equivariant排列等變。也就是說模型本身根本不知道“誰在誰后面”。把這個問題類比成語義理解就很好懂了。同樣一句話“張三打了李四”和“李四打了張三”用的詞一模一樣唯一的差別在詞序。如果模型沒有位置信息它看到的就是一個詞袋只能知道出現(xiàn)了張三、李四、打卻分不清誰打了誰。對短句來說預(yù)訓(xùn)練數(shù)據(jù)里的大量語序先驗還能補救但一旦句子變長、結(jié)構(gòu)變得復(fù)雜位置信息的缺失就是致命的。1.2 為什么早期方案不徹底初代 Transformer 用的是正弦加法式位置編碼直接把一個隨位置變化的正弦波向量加到詞向量上。這種方式在當(dāng)時的機器翻譯任務(wù)上夠用但它有幾個明顯的先天問題第一位置信息是被“加”進去的和詞向量混在一起。模型必須自己從混合向量里把“內(nèi)容”和“位置”拆開這增加了學(xué)習(xí)負擔(dān)。第二注意力分?jǐn)?shù)計算時兩個 token 的點積里雖然包含了它們的絕對位置信號卻很難直接化簡成只依賴相對位置的形式。換句話說模型要學(xué)出“距離越近關(guān)系越強”這種歸納偏置需要靠大量數(shù)據(jù)硬學(xué)而不是由結(jié)構(gòu)天然保證。后來也有人用可學(xué)習(xí)位置編碼、相對位置偏置、T5 式相對 bias、ALiBi 線性距離懲罰等方案。這些方案各有進展但要么仍然是在加法層面打轉(zhuǎn)要么雖然顯式建模了相對位置卻沒有和 Q、K 的旋轉(zhuǎn)幾何結(jié)合起來無法同時滿足“形式簡單”“可外推”“不增加推理狀態(tài)”這三個要求。這才輪到 RoPE 出場。1.3 我們到底想要怎樣一種位置編碼站在工程角度理想位置編碼應(yīng)該滿足四個條件。一是相對位置敏感兩個 token 的注意力分?jǐn)?shù)最好只由它們的相對位移決定這和語言直覺一致。二是形式輕量不能給 KV Cache 或推理流程增加額外的復(fù)雜度。三是可外推訓(xùn)練時最長 2048推理時最好能體面地跑到 4096 甚至更長而不是直接斷崖式崩潰。四是向量范數(shù)穩(wěn)定對向量做變換時不要把向量的模長也跟著改掉否則注意力在長序列上的尺度會失控。RoPE 恰恰同時滿足這四點。它的核心思路是給每個 Query 和 Key 單獨施加一個只和它自身位置相關(guān)的旋轉(zhuǎn)變換再計算兩者的點積。三角函數(shù)和旋轉(zhuǎn)矩陣恰好滿足“先旋轉(zhuǎn)再點積 旋轉(zhuǎn)一點的差后點積”于是相對位置就自動浮現(xiàn)了。2. 旋轉(zhuǎn)矩陣如何把“位置”藏進夾角2.1 從二維旋轉(zhuǎn)說起回想二維平面上的旋轉(zhuǎn)矩陣。把一個向量逆時針旋轉(zhuǎn)角度m·θ可以得到R(mθ) [[cos(mθ), -sin(mθ)], [sin(mθ), cos(mθ)]]這里m是 token 在序列里的位置θ是一個基礎(chǔ)頻率。比如位置 1 旋轉(zhuǎn)θ位置 5 旋轉(zhuǎn)5θ位置差為 4 的兩個向量它們之間的相對角度就是4θ。如果把注意力計算里的二維 Query 和二維 Key 都各自旋轉(zhuǎn)再來做點積就會得到?R(mθ) q, R(nθ) k? q^T R(mθ)^T R(nθ) k q^T R((n-m)θ) k最終等式里只剩n-m跟 m 和 n 的絕對值無關(guān)。這就是 RoPE 的全部魔法只要 Q 和 K 用同一套旋轉(zhuǎn)體系點積天然只看相對位置。2.2 從二維推廣到高維語言模型的 head_dim 通常是偶數(shù)比如 64 或 128。RoPE 的做法是把高維向量拆成d/2個二維子向量每個子向量用各自獨立的頻率θ_i旋轉(zhuǎn)。這些頻率設(shè)計成從大到小的指數(shù)分布θ_i 10000^(-2i / d), i 0, 1, ..., d/2 - 1當(dāng)i0時頻率最高旋轉(zhuǎn)角度變化最快當(dāng)i接近d/2-1時頻率最低旋轉(zhuǎn)角度變化非常緩慢。高頻率通道負責(zé)捕捉相鄰詞的局部關(guān)系低頻率通道負責(zé)傳遞跨很長的位置信息。這和小波分析里的多尺度分解思路類似不同通道各司其職。完整的高維變換寫成公式大概是f(q, m) [ [cos(mθ_0), -sin(mθ_0)], [sin(mθ_0), cos(mθ_0)], ... [cos(mθ_{d/2-1}), -sin(mθ_{d/2-1})], [sin(mθ_{d/2-1}), cos(mθ_{d/2-1})] ] · q這塊對角旋轉(zhuǎn)矩陣雖然龐大但實際運算時完全不需要顯式構(gòu)建矩陣。我們可以借用復(fù)數(shù)觀點把相鄰的兩個維度看成一個復(fù)數(shù)旋轉(zhuǎn)就是乘上e^(i·m·θ_i)。這正是原始論文和很多開源實現(xiàn)采用復(fù)數(shù)形式的原因。2.3 三個容易被忽略的幾何性質(zhì)第一個性質(zhì)是模長不變。旋轉(zhuǎn)不改變向量的長度所以|R(m)q| |q|。這一點很重要它保證了注意力分?jǐn)?shù)不會因為位置編碼的引入而出現(xiàn)尺度漂移。第二個性質(zhì)是相對位置分解。這一點前面已經(jīng)推導(dǎo)過是 RoPE 名稱里“旋轉(zhuǎn)”的底氣。它讓模型在結(jié)構(gòu)上就具備相對位置建模能力而不是靠數(shù)據(jù)硬學(xué)。第三個性質(zhì)是遠程衰減。由于用了多層不同頻率的旋轉(zhuǎn)當(dāng)位置差變大時高頻率分量之間的點積會快速震蕩抵消整體上會出現(xiàn)一個隨距離增加而衰減的包絡(luò)。這個衰減不是某個規(guī)則寫死的而是由頻率分布自然形成的類似于原始正弦編碼里的衰減行為但尺度更平滑。實際調(diào)試時我經(jīng)?？吹接腥藛枴盀槭裁?RoPE 不直接把位置加在 embedding 上”。答案是加法會改向量長度旋轉(zhuǎn)不會加法無法化簡成n-m旋轉(zhuǎn)可以。RoPE 是乘法式位置編碼里目前最優(yōu)雅的一種。3. 從公式到 PyTorch一份可直接運行的實現(xiàn)3.1 先寫旋轉(zhuǎn)函數(shù)高維旋轉(zhuǎn)的落地寫法有兩種。論文原汁原味的寫法是用復(fù)數(shù)每兩個相鄰維度組成一個復(fù)數(shù)和預(yù)計算的復(fù)數(shù)頻率相乘。主流框架里更常見的是rotate_half寫法它利用向量前后兩半的錯位拼接省去復(fù)數(shù)操作。先看rotate_halfimport torch def rotate_half(x): 把向量的后半段取負接到前半段后面。 x1 x[..., : x.shape[-1] // 2] x2 x[..., x.shape[-1] // 2:] return torch.cat((-x2, x1), dim-1)再看應(yīng)用函數(shù)def apply_rotary_pos_emb(q, k, cos, sin): # 假設(shè) q/k 形狀都是 [batch, num_heads, seq_len, head_dim] # cos/sin 形狀為 [seq_len, head_dim] cos cos[: q.size(-2), :].unsqueeze(0).unsqueeze(0) sin sin[: q.size(-2), :].unsqueeze(0).unsqueeze(0) q_embed (q * cos) (rotate_half(q) * sin) k_embed (k * cos) (rotate_half(k) * sin) return q_embed, k_embed這里的關(guān)鍵是把向量x拆成前后兩半然后用(q * cos) rotate_half(q) * sin組合。它實際上把不同維度對按一定分組做了旋轉(zhuǎn)和復(fù)數(shù)乘法是同一族旋轉(zhuǎn)嵌入只是分組順序不同。在工程上這種寫法更直觀也完全不用處理復(fù)數(shù)類型。3.2 預(yù)計算 cos 和 sin 表位置編碼是確定的不需要學(xué)習(xí)所以可以預(yù)先算出一張表運行時按位置取用。頻率公式里的10000是 base后面會講它怎么影響外推。def precompute_rope_cache(dim, max_seq_len, base10000.0): # 生成 [θ_0, θ_1, ..., θ_(dim/2-1)] inv_freq 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim)) # 生成每個位置的角度 m * θ_i t torch.arange(max_seq_len).float().unsqueeze(1) angles t * inv_freq.unsqueeze(0) # [max_seq_len, dim // 2] angles torch.cat((angles, angles), dim-1) # [max_seq_len, dim] cos_cached angles.cos() sin_cached angles.sin() return cos_cached, sin_cached這里torch.cat((angles, angles), dim-1)是工程里很常見的加速技巧先算出dim // 2個角再復(fù)制一份拼到后面讓每個旋轉(zhuǎn)“配對”有兩個維度共享同一個頻率省去逐維度查表。3.3 復(fù)數(shù)寫法更貼近論文原理解讀如果你從論文或某些復(fù)現(xiàn)代碼里見到復(fù)數(shù)實現(xiàn)也不奇怪。它的核心是把相鄰兩個維度讀成一個復(fù)數(shù)def apply_rotary_emb(xq, xk, freqs_cis): # xq: [batch, num_heads, seq_len, head_dim] xq_ xq.float().view(*xq.shape[:-1], -1, 2) xk_ xk.float().view(*xk.shape[:-1], -1, 2) xq_ torch.view_as_complex(xq_) # [batch, heads, seq, dim/2] xk_ torch.view_as_complex(xk_) freqs_cis freqs_cis[: xq_.size(-2)] # [seq, dim/2] freqs_cis freqs_cis.unsqueeze(0).unsqueeze(0) xq_out torch.view_as_real(xq_ * freqs_cis).flatten(-2) xk_out torch.view_as_real(xk_ * freqs_cis).flatten(-2) return xq_out.type_as(xq), xk_out.type_as(xk)復(fù)數(shù)寫法對數(shù)學(xué)推導(dǎo)更友好xq_ * freqs_cis就是復(fù)平面上的一次旋轉(zhuǎn)。但復(fù)數(shù)在有些后端比如某些半精度推理腳本上會帶來額外的類型兼容問題所以在工程上rotate_half版本更常用。3.4 接入多頭注意力把上面的片段串進一個常規(guī)多頭注意力模塊核心就三行cos, sin precompute_rope_cache(head_dim, max_seq_len) q_embed, k_embed apply_rotary_pos_emb(q, k, cos, sin) attn_weights torch.matmul(q_embed, k_embed.transpose(-2, -1)) / math.sqrt(head_dim)這里有一個很容易搞錯的細節(jié)只旋轉(zhuǎn) Q 和 KV 完全不旋轉(zhuǎn)。V 只是被注意力分?jǐn)?shù)加權(quán)求和的對象位置關(guān)系已經(jīng)由權(quán)重表達了再給 V 加位置反而會造成信息重復(fù)。4. 和絕對位置、相對位置、ALiBi 的橫向?qū)Ρ?.1 一張表看清主流方案位置編碼方案作用位置相對位置依賴訓(xùn)練后外推能力KV Cache 負擔(dān)典型場景正弦絕對編碼加到詞嵌入上隱式需要模型自己學(xué)弱超出訓(xùn)練長度迅速退化無早期 Transformer可學(xué)習(xí)絕對編碼加到詞嵌入上隱式很弱訓(xùn)練長度是最硬上限無經(jīng)典雙向編碼器T5 相對 bias加到注意力 logits 上顯式中等超出 bias 表范圍即失效需要緩存 bias 表T5 類模型ALiBi對注意力 logits 做距離懲罰顯式較好但沒有旋轉(zhuǎn)幾何的深層性質(zhì)無長文本推理類模型RoPE作用于 Q、K 向量本身顯式由旋轉(zhuǎn)推導(dǎo)中等但配合 base 調(diào)整能顯著增強無額外狀態(tài)現(xiàn)代大模型默認(rèn)方案4.2 為什么加法式方案輸在了結(jié)構(gòu)上加法式位置編碼的問題在于位置信息是“疊”在內(nèi)容上的。模型對同一個位置的向量施加的偏置會被投到所有下游計算里最后注意力分?jǐn)?shù)確實包含了位置信號但這個信號和內(nèi)容信號糾纏在一起無法分離。相對位置 bias 和 ALiBi 把位置直接疊加到注意力得分上已經(jīng)比加法式強很多但它們?nèi)匀皇且环N“修正項”沒有改變 Q、K 本身的幾何結(jié)構(gòu)。RoPE 則是在計算 Q、K 的點積之前先把它們做了旋轉(zhuǎn)。旋轉(zhuǎn)之后Q、K 的內(nèi)容信息一點沒丟但點積的表達式里自動出現(xiàn)了相對位置。這就是“乘法式”和“加法式”的本質(zhì)區(qū)別RoPE 把位置信息變成了向量之間的夾角關(guān)系而不是一個附加數(shù)值。4.3 長上下文場景下誰更省心從工程角度看RoPE 在 KV Cache 上的優(yōu)勢也很明顯?？蓪W(xué)習(xí)位置編碼在推理時需要為每個位置維護獨立的嵌入向量ALiBi 和 T5 bias 也要在運行時查表或構(gòu)造偏置。RoPE 則只需要在加載 K、V 時把當(dāng)前增量位置的旋轉(zhuǎn)角算好緩存里無需額外存位置狀態(tài)。推理效率上RoPE 幾乎沒有額外開銷這也是大模型普遍選擇它的一個重要原因。當(dāng)然 RoPE 不是萬能藥。它放棄了對模型“某絕對位置有特殊含義”的顯式建模如果你做的是必須依賴絕對位置的合成任務(wù)比如“第 5 個位置一定是某種 token”RoPE 可能需要額外輸入特征來補。但對絕大多數(shù)自然語言序列建模來說相對位置遠比絕對位置重要。5. 工程落地避坑外推、base 和 KV Cache5.1 訓(xùn)練 2048推理 8192為什么直接崩很多人在自己的模型上把訓(xùn)練序列長度定為 2048用了 RoPE 之后天真地以為“相對位置編碼應(yīng)該能外推”結(jié)果把輸入改成 8192 一跑困惑度直接飆升。原因是 RoPE 的頻率是按訓(xùn)練時見過的位置范圍設(shè)計的。位置 4096 對應(yīng)的最低頻通道角度可能是訓(xùn)練時從未出現(xiàn)過的值模型雖然在結(jié)構(gòu)上是相對位置的但訓(xùn)練到的統(tǒng)計特性并沒有覆蓋這么大的旋轉(zhuǎn)角度范圍。絕對位置雖然被旋轉(zhuǎn)“隱藏”了卻仍然會通過頻率范圍影響模型。5.2 調(diào) base 才是外推的真正入口長度外推最常見的方案是修改頻率公式里的 base。NTK-aware 放縮是目前最通用的經(jīng)驗法則假設(shè)你打算把上下文從訓(xùn)練長度 L 擴展到目標(biāo)長度L令放縮比s L / L那么新的 base 大致取new_base base × s^(d / (d - 2))這里的d是 head_dim。為什么分母是d-2因為頻率最低的那個通道i d/2-1對應(yīng)的角頻率是base^(-(d-2)/d)。要讓它在目標(biāo)長度內(nèi)也至少能轉(zhuǎn)滿必要的圈數(shù)就必須按這個冪次抬高 base。這個公式本質(zhì)上是保證最高頻率不變、最低頻率按比例壓低讓整個頻譜能均勻覆蓋到更長的序列。實際調(diào)參時我不會一次性把 base 從 10000 調(diào)到百萬級。比如從 2048 擴展到 8192s4head_dim64那么新的 base 大約是10000 × 4^(64/62) ≈ 10000 × 4.09 ≈ 40900這個量級是相對溫和的短序列上的表現(xiàn)也不會明顯退化。5.3 動態(tài)插值和 YaRN 的思路NTK-aware 的 base 改法只針對“目標(biāo)長度”單一值。如果推理時輸入長度是 4096過一會兒變成 8192再變成 32768固定 base 就不太夠用了。動態(tài) NTK 的思路是隨當(dāng)前序列長度實時計算 base長度越長 base 越大避免最壞情況出現(xiàn)。更精致的方案是 YaRN。它的思路是對不同頻率通道區(qū)別對待高頻通道波長很短訓(xùn)練中已經(jīng)見過大量完整周期盡量保持不動低頻通道波長很長訓(xùn)練中可能只見過不足一圈按放縮比線性拉伸中間頻段做平滑過渡。這樣做的物理直覺是高頻通道編碼的是局部精細關(guān)系微調(diào)它會破壞既有能力低頻通道編碼的是全局粗粒度關(guān)系拉伸它代價小、收益大。我實測下來YaRN 在 2 倍到 4 倍長度擴展時效果不錯但每次換配置都必須跑一遍驗證集對比困惑度沒有銀彈。5.4 KV Cache 上最容易犯的錯RoPE 在自回歸推理時KV Cache 里緩存的 Key 必須是已經(jīng)做過旋轉(zhuǎn)的 Key。當(dāng)你生成第t1個 token 時新 Query 要用位置t1的 cos/sin 旋轉(zhuǎn)新 Key 也要用位置t1的 cos/sin 旋轉(zhuǎn)但緩存里舊的 Key 絕不能再用當(dāng)前位置的 cos/sin 重新乘一遍。如果你忘記區(qū)分“位置索引”和“緩存序號”就會出現(xiàn)一個經(jīng)典 bug緩存里的 Key 被反復(fù)旋轉(zhuǎn)注意力分?jǐn)?shù)完全錯亂。我在調(diào)試解耦注意力時遇到過類似問題最后是打印每個位置 Key 的 norm 和夾角才發(fā)現(xiàn)不對。正確做法是給 RoPE 模塊傳入position_offset讓新 token 的旋轉(zhuǎn)角從offset開始取而緩存區(qū)只負責(zé)追加不做重復(fù)旋轉(zhuǎn)。5.5 推薦的外推流程我自己現(xiàn)在做長上下文微調(diào)時一般先跑一輪短腳本把目標(biāo)長度和訓(xùn)練長度的困惑度差量測出來。如果差量不大直接微調(diào)少量 step 就能恢復(fù)如果差量大先按 NTK 公式改 base再評估。直白的經(jīng)驗是base 改小了短序列受損base 改大了長序列不平滑需要找一個在兩者之間平衡好的點并帶著目標(biāo)長度一起做幾次短訓(xùn)練來校準(zhǔn)。6. 我實際調(diào)試 RoPE 時的一些經(jīng)驗清單最后分享幾條零散的實操體會都是常規(guī)文檔里不太會寫的東西。第一RoPE 的實現(xiàn)通常不需要學(xué)習(xí)參數(shù)但它極度依賴 dtype。在 float16 下累計幾百個位置的 cos/sin 表精度損失會在長序列上累積。如果發(fā)現(xiàn) 4096 長度下中間層輸出有奇怪抖動優(yōu)先把 RoPE 的緩存表用 float32 存儲計算時再轉(zhuǎn)換。第二head_dim 必須是偶數(shù)。如果你的注意力頭維度是奇數(shù)要么補一個維度要么用 padding 截斷。實際模型里 head_dim 幾乎都是 2 的冪但自己搭小實驗時容易忽略。第三位置頻率的順序會影響訓(xùn)練效果。有的實現(xiàn)把高頻維度排在前面有的做交錯排列。高頻在前、低頻在后的排列是最常見的但如果你從開源代碼里抄實現(xiàn)一定要保持預(yù)計算表和rotate_half的分組方式一致。混用不同的分組會導(dǎo)致模型權(quán)重不匹配加載即失效。第四微調(diào)時如果只增加了上下文長度但沒有增加數(shù)據(jù)多樣性外推效果仍然有限。RoPE 只負責(zé)“結(jié)構(gòu)上能支持更長的相對距離”模型要真正學(xué)會利用這些長距離信息數(shù)據(jù)里必須要有足夠的跨長距離依賴樣本。否則你可以把位置編碼調(diào)得再完美模型也只是把注意力均勻攤開并不會產(chǎn)出更智能的長文能力。第五如果只是為了快速驗證 RoPE 是否導(dǎo)致 bug最直接的方法是構(gòu)造一個固定權(quán)重的小模型分別輸入“位置 0 到 256”和“位置 1 到 257”對比兩組輸出的相對差異是否一致。因為 RoPE 是相對位置驅(qū)動的好的實現(xiàn)應(yīng)當(dāng)保證這兩組結(jié)果幾乎只在絕對位置分量上有微小差別。這個小實驗?zāi)軒湍惆褜崿F(xiàn)問題和模型問題快速分開。RoPE 并不是什么復(fù)雜的高深理論它本質(zhì)上就是一個聰明的旋轉(zhuǎn)矩陣應(yīng)用只是把“位置”做進了 Q、K 的夾角里。理解清楚旋轉(zhuǎn)的幾何含義之后不管是在自定義注意力里手寫位置編碼還是在長上下文模型里調(diào)外推你都會比只看文檔照抄的人多一層底氣和判斷力。