位置編碼深度解析:原理、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ì)稱(chēng)的它不關(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ù)天然只依賴(lài)兩個(gè) token 之間的相對(duì)距離而不依賴(lài)它們各自的絕對(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)題類(lèi)比成語(yǔ)義理解就很好懂了。同樣一句話(huà)“張三打了李四”和“李四打了張三”用的詞一模一樣唯一的差別在詞序。如果模型沒(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)成只依賴(lài)相對(duì)位置的形式。換句話(huà)說(shuō)模型要學(xué)出“距離越近關(guān)系越強(qiáng)”這種歸納偏置需要靠大量數(shù)據(jù)硬學(xué)而不是由結(jié)構(gòu)天然保證。后來(lái)也有人用可學(xué)習(xí)位置編碼、相對(duì)位置偏置、T5 式相對(duì) bias、ALiBi 線(xiàn)性距離懲罰等方案。這些方案各有進(jìn)展但要么仍然是在加法層面打轉(zhuǎn)要么雖然顯式建模了相對(duì)位置卻沒(méi)有和 Q、K 的旋轉(zhuǎn)幾何結(jié)合起來(lái)無(wú)法同時(shí)滿(mǎn)足“形式簡(jiǎn)單”“可外推”“不增加推理狀態(tài)”這三個(gè)要求。這才輪到 RoPE 出場(chǎng)。1.3 我們到底想要怎樣一種位置編碼站在工程角度理想位置編碼應(yīng)該滿(mǎn)足四個(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í)滿(mǎn)足這四點(diǎn)。它的核心思路是給每個(gè) Query 和 Key 單獨(dú)施加一個(gè)只和它自身位置相關(guān)的旋轉(zhuǎn)變換再計(jì)算兩者的點(diǎn)積。三角函數(shù)和旋轉(zhuǎn)矩陣恰好滿(mǎ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)的位置信息。這和小波分析里的多尺度分解思路類(lèi)似不同通道各司其職。完整的高維變換寫(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 名稱(chēng)里“旋轉(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ě)死的而是由頻率分布自然形成的類(lèi)似于原始正弦編碼里的衰減行為但尺度更平滑。實(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ù)類(lèi)型。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)額外的類(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ì)位置依賴(lài)訓(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 類(lèi)模型ALiBi對(duì)注意力 logits 做距離懲罰顯式較好但沒(méi)有旋轉(zhuǎn)幾何的深層性質(zhì)無(wú)長(zhǎng)文本推理類(lèi)模型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ì)也很明顯??蓪W(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)銷(xiāo)這也是大模型普遍選擇它的一個(gè)重要原因。當(dāng)然 RoPE 不是萬(wàn)能藥。它放棄了對(duì)模型“某絕對(duì)位置有特殊含義”的顯式建模如果你做的是必須依賴(lài)絕對(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)滿(mǎ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ò)不足一圈按放縮比線(xiàn)性拉伸中間頻段做平滑過(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ò)類(lèi)似問(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ù)但它極度依賴(lài) 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的分組方式一致。混用不同的分組會(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)距離依賴(lài)樣本。否則你可以把位置編碼調(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ì)比只看文檔照抄的人多一層底氣和判斷力。