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