0.前言

??本文是近期位置編碼相關(guān)內(nèi)容的學(xué)習(xí)記錄，之前遇到位置編碼的內(nèi)容都是直接跳過的，在看了近期一些模型還有蘇建林老師的博客內(nèi)容后發(fā)現(xiàn)位置編碼也是一個很重要的內(nèi)容。
??直接從最新的看會很迷惑這些位置編碼的代碼是在做什么神奇的操作。。。。。。以及為什么是這樣，所以本文從最早的開始記錄，也是我學(xué)習(xí)的過程。

1. sinusoidal編碼

??這部分推薦看蘇建林老師的博客，想不明白的地方可以截圖不停地追問通義千問。

1.0 數(shù)學(xué)知識——復(fù)數(shù)

1.0.1 復(fù)數(shù)乘法、共軛復(fù)數(shù)

??在數(shù)學(xué)中，復(fù)數(shù)可以被表示為 $a+bi$ 的形式，其中 $a$ 和 $b$ 是實數(shù)，$i$ 是虛數(shù)單位（滿足 $i^2=?1$)。復(fù)數(shù)可以在二維平面上用向量表示，橫軸代表實部，縱軸代表虛部。
??假設(shè)我們有兩個二維向量 $[x_1,y_1]$ 和 $[x_2,y_2]$，我們可以將它們視為兩個復(fù)數(shù) $z_1=x_1+y_{1}i$ 和 $z_2=x_2+y_2$。
??復(fù)數(shù)的乘法遵循特定的規(guī)則：

$$z_1?z_2=(x_1+y_{1}i)?(x_2+y_{2}i)=x_1{x}_2?y_1{y}_2+(x_1{y}_2+x_2{y}_1)i$$

??如果我們想要計算兩個復(fù)數(shù)的內(nèi)積，并且只關(guān)心結(jié)果的實部，那么我們可以使用共軛的概念。給定一個復(fù)數(shù) $z=a+bi$，其共軛定義為$z^{?}=a?bi$?；楣曹椀膬蓚€復(fù)數(shù)相乘，結(jié)果為模長平方。
$$z?z^{?}=(a+bi)?(a?bi)=a(a)+a(?bi)+(bi)a+(bi)(?bi)=a^2?abi+abi?b^2{i}^2=a^2+b^2$$

??使用共軛可以幫助我們“消去”虛部，使得最終結(jié)果成為實數(shù)。對于兩個復(fù)數(shù)$z_1$ 和 $z_2$，它們的共軛為
$$z_1?z_2^{?}=(x_1+y_{1}i)?(x_2?y_{2}i)=x_1{x}_2+y_1{y}_2+(x_2{y}_1?x_1{y}_2)i$$

它們的內(nèi)積可以定義為：

$$?z_1,z_2?=x_1{x}_2+y_1{y}_2=\text{Re}[z_1?z_2^{?}]$$

換句話說，給定兩個復(fù)數(shù) $z_1=x_1+y_{1}i$ 和 $z_2=x_2+y_{2}i$，它們作為二維向量的內(nèi)積可以通過公式 $?z_1,z_2?=x_1{x}_2+y_1{y}_2$ 來計算。

1.0.2 復(fù)數(shù)的指數(shù)表示

??復(fù)數(shù)還有指數(shù)表示形式，它基于歐拉公式（Euler’s formula），將復(fù)數(shù)與三角函數(shù)和指數(shù)函數(shù)聯(lián)系起來。歐拉公式表述為：

$$e^{i\theta }=\cos (\theta )+i\sin (\theta )$$

這里，$e$ 是自然對數(shù)的底數(shù)，而 $\theta$ 是以弧度為單位的角度。

??對于任意一個非零復(fù)數(shù) $z=a+bi$，可以將其轉(zhuǎn)換為極坐標(biāo)形式（polar form）來表示，即通過它的模（magnitude）和輻角（argument）來描述：

• 模（或絕對值）：$r=|z|=\sqrt{a^2+b^2}$
• 輻角（或幅角）：$\theta =\arg (z)$，是實軸正方向到從原點到復(fù)數(shù)點連線之間的夾角。

因此，任何非零復(fù)數(shù)都可以寫成：

$$z=r(\cos (\theta )+i\sin (\theta ))$$

利用歐拉公式，這個表達式可以簡化為指數(shù)形式：

$$z=r{e}^{i\theta }$$

這里，$r$ 代表復(fù)數(shù)的長度或大小，而$e^{i\theta }$描述了該復(fù)數(shù)的方向。

1.1 sinusoidal編碼來歷

??之所以要位置編碼，在沒有掩碼的情況下，attention函數(shù)$f(x)$是對稱的，比如對于輸入的Q序列里面的兩個向量$x_m$和$x_n$調(diào)換位置( f 1 = { x 1 , . . . , x m , . . . , x n , . . . } f_1=\{x_1,...,x_m,...,x_n,...\} f1?={x1?,...,xm?,...,xn?,...}和 f 2 = { x 1 , . . . , x n , . . . , x m , . . . } f_2=\{x_1,...,x_n,...,x_m,...\} f2?={x1?,...,xn?,...,xm?,...})，有$f_1=f_2$，從結(jié)果上區(qū)分不出輸入是$x_m$還是$x_n$。
??所以要讓attention的$Q?K$這個乘法過程中，$Q$和$K$分別帶上位置信息，每個位置的向量加上一個和位置信息相關(guān)的向量$p$，變成例如 { x 1 + p 1 , . . . , x m + p m , . . . , x n + p n } \{x_1+p_1,...,x_m+p_m,...,x_n+p_n\} {x1?+p1?,...,xm?+pm?,...,xn?+pn?}，$p_m$是位置編碼向量。
??在只考慮m,n這兩個位置的位置編碼情況下，泰勒展開后發(fā)現(xiàn)只有$p_m^T\mathcal{H}{p}_n$這一項同時包含$p_m$和$p_n$。在最簡單的情況下，取$\mathcal{H}=\mathcal{I}$，此時$p_m^T\mathcal{H}{p}_n=p_m^T{p}_n=?p_m,p_n?$。希望這一項能夠表示m和n的相對位置，最好能有一個函數(shù)$g(?)$使得
$$?p_m,p_n?=g(m?n)$$
??為了方便理解，先考慮2維的情況，假如$Q$是2維的，借助復(fù)數(shù)作為工具進行計算，有$$?p_m,p_n?=\text{Re}[p_m?p_n^{?}]$$
??假設(shè)有復(fù)數(shù)$q_{m?n}$讓上式成立，$p_m?p_n^{?}=q_{m?n}$。用復(fù)數(shù)的指數(shù)形式表示，假設(shè)$p_m=r_m{e}^{i{?}_m}$, $p_n^{?}=r_n{e}^{?i{?}_n}$, $q_{m?n}=R_{m?n}{e}^{i\Phi m?n}$

解方程：
$$r_m{r}_n{e}^{i({?}_m?{?}_n)}=R_{m?n}{e}^{i{\Phi }_{m?n}}$$
整理后有：
$$\{\begin{array}{l}{r}_m{r}_n=R_{m?n}\\ {?}_m?{?}_n={\Phi }_{m?n}\end{array}$$

• 解第一個條件 對于 $r_m{r}_n=R_{m?n}$
  當(dāng) $n=m$ 時，可以得到 $r_m^2=R_0$。這意味著 $r_m$ 是一個常數(shù)（因為 $R_0$ 是一個固定值），為了簡化，設(shè)$r_m=1$。

• 解第二個條件，對于 ${?}_m?{?}_n={\Phi }_{m?n}$
  首先，令 $n=0$，則有${?}_m?{?}_0={\Phi }_m$，如果我們假設(shè)${?}_0=0$（不失一般性，因為角度是相對的），那么 ${?}_m={\Phi }_m$。接著，令 $n=m?1$，則有${?}_m?{?}_{m?1}={\Phi }_1$，這里 ${\Phi }_1$ 是一個固定的相位差。由于 ${\Phi }_{m?n}$ 表示的是相對位置信息，因此 ${\Phi }_1$實際上是一個常數(shù)。這意味著 { ? m } \{\phi_m\} {?m?} 形成了一個等差數(shù)列，其中每一項之間的差值為 ${\Phi }_1$。用數(shù)學(xué)語言來說，就是存在一個常數(shù) $\theta$（在這里 $\theta ={\Phi }_1$）使得${?}_m=m\theta$。

??最終，通過解方程可以得到隱向量維度是2維的情況下，位置編碼的一個解，通過歐拉公式表示為cos和sin的形式：
$$p_m=e^{im\theta }?p_m=\left(\begin{array}{c}\cos (m\theta )\\ \sin (m\theta )\end{array}\right)$$
當(dāng)$Q$向量的隱向量維度是$d$維時，位置$m$的$Q_m$對應(yīng)的編碼向量$p_m=\left(\begin{array}{c}\cos (m{\theta }_0)\\ \sin (m{\theta }_0)\\ \cos (m{\theta }_1)\\ \sin (m{\theta }_1)\\ ...\\ \cos (m{\theta }_{d/2?1})\\ \sin (m{\theta }_{d/2?1})\end{array}\right)$
??這里面需要注意的是，$d$指的是隱向量的維度，$m$指的是向量是在第$m$個。在《Attention is All You Need》，位置編碼的計算公式如下：
$$\{\begin{array}{l}{p}_{k,2i}=\sin \left(\frac{k}{10000^{2i/d}}\right)\\ p_{k,2i+1}=\cos \left(\frac{k}{10000^{2i/d}}\right)\end{array}$$
這里，$p_{k,2i}$ 和 $p_{k,2i+1}$ 分別表示位置 $k$ 的編碼向量的第$2i$ 和 $2i+1$ 個分量，$k$ 是位置索引(對應(yīng)上面的推導(dǎo)的$m$)， $i$ 是向量維度的索引，$d$ 是向量的總維度。對應(yīng)位置的位置編碼會和在attention運算前$Q$和$K$相加。

1.2 代碼實現(xiàn)

??看了下代碼，之前不少多模態(tài)模型的位置編碼都是學(xué)習(xí)式的，而且是直接位置編碼和$q$、$k$相加?！禔ttention is all you need》有一份別人實現(xiàn)的pytorch代碼里面是sinusoidal編碼，并且完全遵循了上面公式的實現(xiàn)方式，sinusoidal編碼也是和$q$、$k$相加：

#https://github.com/jadore801120/attention-is-all-you-need-pytorch/blob/master/transformer/Models.py
class PositionalEncoding(nn.Module):def __init__(self, d_hid, n_position=200):super(PositionalEncoding, self).__init__()# Not a parameterself.register_buffer('pos_table', self._get_sinusoid_encoding_table(n_position, d_hid))def _get_sinusoid_encoding_table(self, n_position, d_hid):''' Sinusoid position encoding table '''# TODO: make it with torch instead of numpydef get_position_angle_vec(position):return [position / np.power(10000, 2 * (hid_j // 2) / d_hid) for hid_j in range(d_hid)]sinusoid_table = np.array([get_position_angle_vec(pos_i) for pos_i in range(n_position)])sinusoid_table[:, 0::2] = np.sin(sinusoid_table[:, 0::2])  # dim 2isinusoid_table[:, 1::2] = np.cos(sinusoid_table[:, 1::2])  # dim 2i+1return torch.FloatTensor(sinusoid_table).unsqueeze(0)def forward(self, x):return x + self.pos_table[:, :x.size(1)].clone().detach()  # 直接相加

2. Rotary Positional Embedding (RoPE) ——旋轉(zhuǎn)位置編碼

• 運算前含有絕對位置的信息，運算后的結(jié)果含有相對位置的信息

2.1 RoPE來歷

??RoPE是蘇建林老師在博客里面提出來的一種位置編碼方式，提出的背景、證明等可以參考其博客空間。這部分如果有比較晦澀難懂的地方也是可以直接截博客里面的圖片問通義千問，通義千問可以在看圖之后進行非常仔細(xì)的解答。
??RoPE用絕對編碼的方式，在計算$Q$和$K$的內(nèi)積時，又讓結(jié)果能帶入$Q$和$K$的相對信息。對于位置為$m$的$q_m$和位置為n的$k_n$分別乘以絕對位置編碼$e^{im\theta }$和$e^{in\theta }$，得到$q_m{e}^{im\theta }$和$q_n{e}^{in\theta }$，在進行內(nèi)積運算，會發(fā)現(xiàn)運算結(jié)果含有相對信息
$$?q_m{e}^{im\theta },k_n{e}^{in\theta }?=\mathrm{Re}\left[(q_m{e}^{im\theta })(k_n{e}^{in\theta }{)}^{?}\right]=\mathrm{Re}\left[q_m{k}_n^{?}{e}^{i(m?n)\theta }\right]$$
??最簡單的情況下假如$Q$向量的隱向量維度$d=2$，這個操作對于位置$m$的$q_m$向量進行了一個旋轉(zhuǎn)操作
$$q_m{e}^{im\theta }=\left(\begin{array}{cc}\cos m\theta & ?\sin m\theta \\ \sin m\theta & \cos m\theta \end{array}\right)\left(\begin{array}{c}{q}_m^0\\ q_m^1\end{array}\right)$$
??通用的情況下，對于位置在$m$(可以說position_id=m)的$Q$向量$q_m$，它的旋轉(zhuǎn)位置編碼的計算過程為：
$$\left(\begin{array}{c}{q}_0\\ q_1\\ q_2\\ q_3\\ ?\\ q_{d?2}\\ q_{d?1}\end{array}\right)?\left(\begin{array}{c}\cos m{\theta }_0\\ \cos m{\theta }_0\\ \cos m{\theta }_1\\ \cos m{\theta }_1\\ ?\\ \cos m{\theta }_{d/2?1}\\ \cos m{\theta }_{d/2?1}\end{array}\right)+\left(\begin{array}{c}?q_1\\ q_0\\ ?q_3\\ q_2\\ ?\\ ?q_{d?1}\\ q_{d?2}\end{array}\right)?\left(\begin{array}{c}\sin m{\theta }_0\\ \sin m{\theta }_0\\ \sin m{\theta }_1\\ \sin m{\theta }_1\\ ?\\ \sin m{\theta }_{d/2?1}\\ \sin m{\theta }_{d/2?1}\end{array}\right)$$

2.2 代碼實現(xiàn)

• RoPE是相乘進行的位置編碼，不是相加

2.2.1 GPT-J風(fēng)格的1D-RoPE實現(xiàn)

??看代碼這部分比較讓人頭大，如果是完全按照上面公式來實現(xiàn)的是一目了然的，首先看這種實現(xiàn)方式，被稱為GPT-J。在Meta官方實現(xiàn)的llama代碼里面，可以找到這種實現(xiàn)方式。當(dāng)然，這里也不是使用提到的乘法方式，而是使用了復(fù)數(shù)運算。一個復(fù)數(shù)對應(yīng)2個實數(shù)，所以如果是$q$轉(zhuǎn)為了復(fù)數(shù)，維度只有$d/2$，最后變成實數(shù)時回到$d$維。以及需要注意$q_0$和$q_1$對應(yīng)的是$m{\theta }_0$，所以freqs_cis的維度只有$q$和$k$的一半就夠了。

# https://github.com/meta-llama/llama/blob/main/llama/model.py# 下面這個函數(shù)是要預(yù)先把從0到最大長度的位置編碼需要使用的角度算好
def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0):"""Precompute the frequency tensor for complex exponentials (cis) with given dimensions.This function calculates a frequency tensor with complex exponentials using the given dimension 'dim'and the end index 'end'. The 'theta' parameter scales the frequencies.The returned tensor contains complex values in complex64 data type.Args:dim (int): Dimension of the frequency tensor.end (int): End index for precomputing frequencies.theta (float, optional): Scaling factor for frequency computation. Defaults to 10000.0.Returns:torch.Tensor: Precomputed frequency tensor with complex exponentials."""## 因為維度為2i、2i+1的mθ相同，所以是(0, dim, 2)freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim)) # 算θ值t = torch.arange(end, device=freqs.device)  # end是最大長度，對應(yīng)一個個位置mfreqs = torch.outer(t, freqs).float()  #這個是算m*θ行數(shù)為t，列數(shù)為dim//2，每行對應(yīng)一個q向量freqs_cis = torch.polar(torch.ones_like(freqs), freqs)  # 變成復(fù)數(shù)形式，幅度為1，角度為freqsreturn freqs_cis# 在每個attention block中有
xq = xq.view(bsz, seqlen, self.n_local_heads, self.head_dim)
xk = xk.view(bsz, seqlen, self.n_local_kv_heads, self.head_dim)
xv = xv.view(bsz, seqlen, self.n_local_kv_heads, self.head_dim)
xq, xk = apply_rotary_emb(xq, xk, freqs_cis=freqs_cis)
......
score = torch.matmul(xq,xk.transpose(2,3)) # 位置編碼后直接計算attention分?jǐn)?shù)def apply_rotary_emb(xq: torch.Tensor,xk: torch.Tensor,freqs_cis: torch.Tensor,
) -> Tuple[torch.Tensor, torch.Tensor]:"""Apply rotary embeddings to input tensors using the given frequency tensor.This function applies rotary embeddings to the given query 'xq' and key 'xk' tensors using the providedfrequency tensor 'freqs_cis'. The input tensors are reshaped as complex numbers, and the frequency tensoris reshaped for broadcasting compatibility. The resulting tensors contain rotary embeddings and arereturned as real tensors.Args:xq (torch.Tensor): Query tensor to apply rotary embeddings.xk (torch.Tensor): Key tensor to apply rotary embeddings.freqs_cis (torch.Tensor): Precomputed frequency tensor for complex exponentials.Returns:Tuple[torch.Tensor, torch.Tensor]: Tuple of modified query tensor and key tensor with rotary embeddings.     """# 把q向量看成復(fù)數(shù)，2個2個一組看成一個復(fù)數(shù)，例如(q0,q1)->(qc_0)xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))freqs_cis = reshape_for_broadcast(freqs_cis, xq_)xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3) # 做乘法xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)return xq_out.type_as(xq), xk_out.type_as(xk)

??上面的代碼是用復(fù)數(shù)乘法實現(xiàn)的，可能不是特別直觀，考慮最簡單的$d=2$的情形，這種情況下令$q=(q_0,q_1)$，這兩個向量要旋轉(zhuǎn)的角度是${\theta }_0$。
??首先，apply_rotary_emb()函數(shù)里面的view_as_complex是讓$q_0$和$q_1$組成了一個復(fù)數(shù)$q_c=q_0+i?q_1$。
??假設(shè) freqs_cis 對應(yīng)于這個位置和頻率分量的旋轉(zhuǎn)因子為 $e^{i{\theta }_0}=\cos ({\theta }_0)+i\sin ({\theta }_0)$，即[$cos({\theta }_0)$,$sin({\theta }_0)$]，注意預(yù)先計算的函數(shù)precompute_freqs_cis()里面最后也是以復(fù)數(shù)形式表示的，這個cos和sin變成了一個復(fù)數(shù)，也就是freqs_cis[0] = $cos({\theta }_0)+i?sin({\theta }_0)$。

??對 $q_0$ 和$q_1$進行旋轉(zhuǎn)，需要執(zhí)行復(fù)數(shù)乘法xq_ * freqs_cis：
$$q_m{e}^{im{\theta }_0}=(q0+i?q1)\times (\cos ({\theta }_0)+i?\sin ({\theta }_0))$$

根據(jù)復(fù)數(shù)乘法公式：

$$(a+bi)\times (c+di)=(ac?bd)+i(ad+bc)$$

上述表達式展開為：

$$(q0?\cos ({\theta }_0)?q1?\sin ({\theta }_0))+i(q0?\sin ({\theta }_0)+q1?\cos ({\theta }_0))$$

因此，旋轉(zhuǎn)后的結(jié)果是一個新的復(fù)數(shù)，其實部和虛部分別為：

• 實部：$q0?\cos ({\theta }_0)?q1?\sin ({\theta }_0)$
• 虛部：$q0?\sin ({\theta }_0)+q1?\cos ({\theta }_0)$

這和1D-RoPE的結(jié)果是一致的：

$$(q0?\cos ({\theta }_0)?q1?\sin ({\theta }_0),q0?\sin ({\theta }_0)+q1?\cos ({\theta }_0))$$

2.2.2 GPT-NeoX style的1D-RoPE

??在eleuther的官方實現(xiàn)以及transformer的llama代碼里面，使用的是另一種風(fēng)格的旋轉(zhuǎn)位置編碼的代碼，沒有用到復(fù)數(shù)計算。重點關(guān)注forward、rotate_half以及embedding函數(shù)。

# https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.pyclass LlamaRotaryEmbedding(nn.Module):def __init__(self, config: LlamaConfig, device=None):super().__init__()# BC: "rope_type" was originally "type"if hasattr(config, "rope_scaling") and config.rope_scaling is not None:self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type"))else:self.rope_type = "default"self.max_seq_len_cached = config.max_position_embeddingsself.original_max_seq_len = config.max_position_embeddingsself.config = configself.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)self.register_buffer("inv_freq", inv_freq, persistent=False)self.original_inv_freq = self.inv_freqdef _dynamic_frequency_update(self, position_ids, device):"""dynamic RoPE layers should recompute `inv_freq` in the following situations:1 - growing beyond the cached sequence length (allow scaling)2 - the current sequence length is in the original scale (avoid losing precision with small sequences)"""seq_len = torch.max(position_ids) + 1if seq_len > self.max_seq_len_cached:  # growthinv_freq, self.attention_scaling = self.rope_init_fn(self.config, device, seq_len=seq_len)self.register_buffer("inv_freq", inv_freq, persistent=False)  # TODO joao: may break with compilationself.max_seq_len_cached = seq_lenif seq_len < self.original_max_seq_len and self.max_seq_len_cached > self.original_max_seq_len:  # reset# This .to() is needed if the model has been moved to a device after being initialized (because# the buffer is automatically moved, but not the original copy)self.original_inv_freq = self.original_inv_freq.to(device)self.register_buffer("inv_freq", self.original_inv_freq, persistent=False)self.max_seq_len_cached = self.original_max_seq_len@torch.no_grad()def forward(self, x, position_ids):if "dynamic" in self.rope_type:self._dynamic_frequency_update(position_ids, device=x.device)# Core RoPE blockinv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)position_ids_expanded = position_ids[:, None, :].float()# Force float32 (see https://github.com/huggingface/transformers/pull/29285)device_type = x.device.typedevice_type = device_type if isinstance(device_type, str) and device_type != "mps" else "cpu"with torch.autocast(device_type=device_type, enabled=False):freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)emb = torch.cat((freqs, freqs), dim=-1)cos = emb.cos()sin = emb.sin()# Advanced RoPE types (e.g. yarn) apply a post-processing scaling factor, equivalent to scaling attentioncos = cos * self.attention_scalingsin = sin * self.attention_scalingreturn cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)  # 這個更像原始的RoPE，沒有變成復(fù)數(shù)，就是分開了cos和sindef rotate_half(x):"""Rotates half the hidden dims of the input."""x1 = x[..., : x.shape[-1] // 2]x2 = x[..., x.shape[-1] // 2 :]return torch.cat((-x2, x1), dim=-1)def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):"""Applies Rotary Position Embedding to the query and key tensors.Args:q (`torch.Tensor`): The query tensor.k (`torch.Tensor`): The key tensor.cos (`torch.Tensor`): The cosine part of the rotary embedding.sin (`torch.Tensor`): The sine part of the rotary embedding.position_ids (`torch.Tensor`, *optional*):Deprecated and unused.unsqueeze_dim (`int`, *optional*, defaults to 1):The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] andsin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, notethat cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q andk have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makescos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k havethe shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.Returns:`tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding."""cos = cos.unsqueeze(unsqueeze_dim)sin = sin.unsqueeze(unsqueeze_dim)q_embed = (q * cos) + (rotate_half(q) * sin)  # 像是原始公式里面的cos和sin操作k_embed = (k * cos) + (rotate_half(k) * sin)return q_embed, k_embed

??一眼看下來，會發(fā)現(xiàn)不對啊，如果沒有rotate_half是可以理解的，rotate_half之后，和原始公式對不上了。原來的是比如$(q_0,q_1)$是在一組的，得到$(q0?\cos ({\theta }_0)?q1?\sin ({\theta }_0),q0?\sin ({\theta }_0)+q1?\cos ({\theta }_0))$?，F(xiàn)在的$q_0$和$q_{d//2}$咋在一起了。而且神奇的是，meta版本代碼訓(xùn)練的模型，能用transformer版本的代碼加載。
??github上有一個issue解釋這個問題，具體參考這個issue。首先結(jié)論是，這是2種RoPE方式，這兩種方式結(jié)果肯定不一樣。但是，如果加入一些轉(zhuǎn)換的代碼，轉(zhuǎn)換后能和另一種方式能對上。
??meta的GPT-J風(fēng)格的模型，要用transformer加載，需要先把$W_q$矩陣$W_k$矩陣進行一些轉(zhuǎn)換，有一個permute函數(shù)專門進行這個操作。

# https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/convert_llama_weights_to_hf.pydef permute(w, n_heads, dim1=dim, dim2=dim):return w.view(n_heads, dim1 // n_heads // 2, 2, dim2).transpose(1, 2).reshape(dim1, dim2)f"model.layers.{layer_i}.self_attn.q_proj.weight": permute(loaded[f"layers.{layer_i}.attention.wq.weight"], n_heads=n_heads)........
## 

??這個函數(shù)的效果在官方論壇里面有，仔細(xì)比對前后，還真能對上。。。。把$W_q$和$W_k$矩陣?yán)锩嬖匚恢米兞耸钦鏇]想到的。

??為什么用這種實現(xiàn)方式，而不是GPT-J的實現(xiàn)方式，官方的解答是，一方面這種方式開銷小，另一方面最重要的原因是GPT-J涉及版權(quán)的問題。

3. 二維旋轉(zhuǎn)位置編碼(2D-RoPE)

??一維旋轉(zhuǎn)位置編碼1D-RoPE，或是二維的2D-RoPE，這個維度指的是有幾維的位置信息，也就是position_id的維度。對于文本，是一維序列，所以只有$x$軸上的信息，position_id =
{0,1,2,3…}，也就是之前提到的$m$，表示到底這個$q$是在第幾個。對于圖片，ViT會切分為一個個patch，每個patch需要標(biāo)識是在第幾行、第幾列，所以需要{w,h}的形式來表示。如果是視頻，還有時間維度時間幀的信息，需要三維的形式{t,w,h}。
??上面的1D-RoPE擴展到高維度的方式很簡單粗暴，如果要進行X-D RoPE，就把$q$和$k$向量從頭到尾平均分成X份，每一份里面再按照position_id進行1D-RoPE。例如是要進行2D RoPE，就把$q$分成$q[0:d/2]$和$q[d/2:]$這兩份，$k$分成$k[0:d/2]$和$k[d/2:]$這兩份，position_id是[(x,y)]的形式，就讓前半段計算$q?e^{ix\theta }$，后半段計算前半段計算$q?e^{iy\theta }$。
??2D-RoPE里面難的地方可能是position_id的計算，尤其如果輸入不止有圖片，是圖文混雜的情況下。蘇神完整分析了各種可行性方案，在他的博客中可以仔細(xì)閱讀這一部分——多模態(tài)位置 編碼的思考。提到了一種方案，就是如果輸入是圖片，就用(x,y)的形式給出position_id，如果輸入是文本，就讓$x=y$。
??舉例而言，首先，對于patch大小為$w?h$的圖片，如果圖片在開頭：

position_id	第1行	第h行
x	1 … 1	h … h
y	1 … w	1 … w

??如果開頭是一段文本，文本的長度為$L$，這個句子的位置編碼為$\{(1,1),(2,2),...,(L,L)\}$，上面的圖片接在這段文本后面，圖片的編碼變?yōu)?/p>

position_id	第1行	第h行
x	L+1 … 1	L+h … L+h
y	L+1 … L+w	L+1 … L+w

??蘇神提到這種方式看著不完美，沒有對稱。因為句子的最后一個token的位置ID是$(L,L)$，它和圖片的第一個patch的位置ID$(L+1,L+1)$的差距是$(1,1)$，但是如果圖片后面再接一個句子，因為圖片的最后一個token的位置ID是$(L+h,L+w)$，如果$w\ne h$，后面這個句子的位置ID不可能和前面圖片最后一個token的差距也是$(1,1)$，顯得不優(yōu)雅，只有像下面示意這樣$w=h$時才對稱。

??更進階的蘇神提到的位置編碼的方式可以看他的博客尤其多模態(tài)編碼的思考這篇，還有下面朋友們寫的評論，都非常有價值?？吹竭@里應(yīng)該就可以去看多模態(tài)模型的2D-RoPE的相關(guān)的源代碼了。

3.1 ECCV的2D-RoPE論文中的實現(xiàn)

??《Rotary Position Embedding for Vision Transformer》這篇論文在ViT中實現(xiàn)了2D-RoPE，不過看上去好像它的編碼freqs_cis是被＋到x上的，也就是$q+e^{i\theta }$，如果我理解錯了請一定要指出來~

#https://github.com/naver-ai/rope-vit/blob/main/models/vit_rope.py# 計算position id
def init_t_xy(end_x: int, end_y: int):t = torch.arange(end_x * end_y, dtype=torch.float32)t_x = (t % end_x).float()t_y = torch.div(t, end_x, rounding_mode='floor').float()return t_x, t_y# 計算RoPE的角度mθ
def compute_axial_cis(dim: int, end_x: int, end_y: int, theta: float = 100.0):freqs_x = 1.0 / (theta ** (torch.arange(0, dim, 4)[: (dim // 4)].float() / dim))freqs_y = 1.0 / (theta ** (torch.arange(0, dim, 4)[: (dim // 4)].float() / dim))t_x, t_y = init_t_xy(end_x, end_y)freqs_x = torch.outer(t_x, freqs_x)freqs_y = torch.outer(t_y, freqs_y)freqs_cis_x = torch.polar(torch.ones_like(freqs_x), freqs_x)freqs_cis_y = torch.polar(torch.ones_like(freqs_y), freqs_y)return torch.cat([freqs_cis_x, freqs_cis_y], dim=-1)class rope_vit_models(vit_models):self.compute_cis = partial(compute_axial_cis, dim=embed_dim//num_heads, theta=rope_theta)freqs_cis = self.compute_cis(end_x = img_size // patch_size, end_y = img_size // patch_size)self.freqs_cis = freqs_cisif self.freqs_cis.shape[0] != x.shape[1] - 1:freqs_cis = self.compute_cis(end_x = W // self.patch_size, end_y = H // self.patch_size)else:freqs_cis = self.freqs_cisfreqs_cis = freqs_cis.to(x.device)for i , blk in enumerate(self.blocks):x = blk(x, freqs_cis=freqs_cis)#但是根據(jù)這個好像這個編碼freqs_cis是被＋到x上的
#https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/vision_transformer.py
#x = x + pos_embed

3.2 qwen2-vl的實現(xiàn)

??首先需要明確一下，qwen2-vl的論文里面重點是說M-RoPE的優(yōu)點，但是里面也寫道了qwen2-vl重新訓(xùn)練了ViT，ViT使用的是2D-RoPE。2D-RoPE是在純ViT部分使用的，圖片首先被CNN變成patch，然后進ViT的block，在ViT的block里面使用2D-RoPE。出了ViT之后，圖片向量填充到文本向量預(yù)留的位置里面之后，對這個圖文混合向量才是使用M-RoPE。
??qwen2-vl的位置編碼風(fēng)格是GPT-NeoX的風(fēng)格的，所以會有rotate_half()函數(shù)，首先看角度生成和最后的乘法部分$q?e^{i\theta }$，看這部分的時候會疑惑position_id的代碼去哪里了，稍后再看。

# modeling_qwen2_vl.py
# 里面涉及到的是apply_rotary_pos_emb_vision函數(shù)，以q為例子
# 輸入為q和位置信息rotary_pos_emb#1. rotary_pos_emb的值為θ角
class VisionRotaryEmbedding(nn.Module):def __init__(self, dim: int, theta: float = 10000.0) -> None:super().__init__()inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float) / dim))self.register_buffer("inv_freq", inv_freq, persistent=False)def forward(self, seqlen: int) -> torch.Tensor:seq = torch.arange(seqlen, device=self.inv_freq.device, dtype=self.inv_freq.dtype)freqs = torch.outer(seq, self.inv_freq)return freqs
head_dim = config.embed_dim // config.num_heads
self.rotary_pos_emb = VisionRotaryEmbedding(head_dim // 2) # 二維的rope，所以只需要一半# Copied from transformers.models.llama.modeling_llama.rotate_half
def rotate_half(x):"""Rotates half the hidden dims of the input."""x1 = x[..., : x.shape[-1] // 2]x2 = x[..., x.shape[-1] // 2 :]return torch.cat((-x2, x1), dim=-1)def apply_rotary_pos_emb_vision(tensor: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor:orig_dtype = tensor.dtypetensor = tensor.float()cos = freqs.cos()sin = freqs.sin()cos = cos.unsqueeze(1).repeat(1, 1, 2).unsqueeze(0).float() # 先在第1維增加一個維度，變成[seqlen,1,dim//4] repeat(1, 1, 2) 表示在第0維不重復(fù)，在第1維不重復(fù)，在第2維重復(fù)2次，變成[[0,1,2,3,0,1,2,3]]。sin = sin.unsqueeze(1).repeat(1, 1, 2).unsqueeze(0).float()output = (tensor * cos) + (rotate_half(tensor) * sin)  # 位置編碼是q*e^iθoutput = output.to(orig_dtype)return output## 2. ViT的block中，q、k被加入位置信息
class VisionAttention(nn.Module):def __init__(self, dim: int, num_heads: int = 16) -> None:super().__init__()self.num_heads = num_headsself.head_dim = dim // num_headsself.qkv = nn.Linear(dim, dim * 3, bias=True)self.proj = nn.Linear(dim, dim)def forward(self, hidden_states: torch.Tensor, cu_seqlens: torch.Tensor, rotary_pos_emb: torch.Tensor = None) -> torch.Tensor:seq_length = hidden_states.shape[0]q, k, v = self.qkv(hidden_states).reshape(seq_length, 3, self.num_heads, -1).permute(1, 0, 2, 3).unbind(0)q = apply_rotary_pos_emb_vision(q.unsqueeze(0), rotary_pos_emb).squeeze(0)k = apply_rotary_pos_emb_vision(k.unsqueeze(0), rotary_pos_emb).squeeze(0)attention_mask = torch.full([1, seq_length, seq_length], torch.finfo(q.dtype).min, device=q.device, dtype=q.dtype)for i in range(1, len(cu_seqlens)):attention_mask[..., cu_seqlens[i - 1] : cu_seqlens[i], cu_seqlens[i - 1] : cu_seqlens[i]] = 0q = q.transpose(0, 1)k = k.transpose(0, 1)v = v.transpose(0, 1)attn_weights = torch.matmul(q, k.transpose(1, 2)) / math.sqrt(self.head_dim)attn_weights = attn_weights + attention_maskattn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(q.dtype)attn_output = torch.matmul(attn_weights, v)attn_output = attn_output.transpose(0, 1)attn_output = attn_output.reshape(seq_length, -1)attn_output = self.proj(attn_output)return attn_output

??position_id的代碼在qwen2-vl的model里面定義，可以看到它也有置換函數(shù)，如果不置換，可以打印出來，以及取pos_id之后打印一下看看，是很規(guī)整的$(0,0),(0,1),....(0,w?1),(1,0),...,(h?1,w?1)$的形式，并且是相乘的形式

class Qwen2VisionTransformerPretrainedModel(Qwen2VLPreTrainedModel):def __init__(self, config) -> None:self.spatial_merge_size = config.spatial_merge_sizehead_dim = config.embed_dim // config.num_headsself.rotary_pos_emb = VisionRotaryEmbedding(head_dim // 2)# grid_thw是一個[[t,h,w]]的形式，如果就一張圖這里t=1def rot_pos_emb(self, grid_thw):pos_ids = []for t, h, w in grid_thw:hpos_ids = torch.arange(h).unsqueeze(1).expand(-1, w)hpos_ids = hpos_ids.reshape(h // self.spatial_merge_size,self.spatial_merge_size,w // self.spatial_merge_size,self.spatial_merge_size,)hpos_ids = hpos_ids.permute(0, 2, 1, 3)  # 可以打印一下不置換的結(jié)果hpos_ids = hpos_ids.flatten()wpos_ids = torch.arange(w).unsqueeze(0).expand(h, -1)wpos_ids = wpos_ids.reshape(h // self.spatial_merge_size,self.spatial_merge_size,w // self.spatial_merge_size,self.spatial_merge_size,)wpos_ids = wpos_ids.permute(0, 2, 1, 3)wpos_ids = wpos_ids.flatten()pos_ids.append(torch.stack([hpos_ids, wpos_ids], dim=-1).repeat(t, 1)) # x,y，重復(fù)t份pos_ids = torch.cat(pos_ids, dim=0)max_grid_size = grid_thw[:, 1:].max()rotary_pos_emb_full = self.rotary_pos_emb(max_grid_size) rotary_pos_emb = rotary_pos_emb_full[pos_ids].flatten(1) # 根據(jù)patch的形狀來取position_id的(x,y)return rotary_pos_embdef forward(self, hidden_states: torch.Tensor, grid_thw: torch.Tensor) -> torch.Tensor:hidden_states = self.patch_embed(hidden_states)rotary_pos_emb = self.rot_pos_emb(grid_thw)for blk in self.blocks:hidden_states = blk(hidden_states, cu_seqlens=cu_seqlens, rotary_pos_emb=rotary_pos_emb)

4. qwen2-vl提出的M-RoPE

??首先明確，M-RoPE是3D-RoPE，乘法過程等是3D-RoPE的方式，自定義的部分在于position_id的計算方式上?？梢曰仡?D-RoPE里面蘇神在多模態(tài)上討論的不同實現(xiàn)方式，到底怎么排文本和圖片。qwen2-vl的論文中給出了一個編排position_id的圖，可以看到圖片是按順序排的，多了一個時間維度的坐標(biāo)，position_id是3維的$(t,h,w)$。然后對于圖片后面接的文本起始編碼，取圖片的最后一個patch的position_id的各個維度的最大值+1。
??代碼上，需要關(guān)注2個地方，一個是這個position_id如何算的，這個在get_rope_index函數(shù)中定義，函數(shù)比較長，可以看它的注釋，實現(xiàn)的就是上圖的邏輯，計算出每個位置的position_id。

Each embedding sequence contains vision embedding and text embedding or just contains text embedding.
For pure text embedding sequence, the rotary position embedding has no difference with mordern LLMs.Examples:input_ids: [T T T T T], here T is for text.temporal position_ids: [0, 1, 2, 3, 4]height position_ids: [0, 1, 2, 3, 4]width position_ids: [0, 1, 2, 3, 4]For vision and text embedding sequence, we calculate 3D rotary position embedding for vision partand 1D rotary position embeddin for text part.Examples:Assume we have a video input with 3 temporal patches, 2 height patches and 2 width patches.input_ids: [V V V V V V V V V V V V T T T T T], here V is for vision.vision temporal position_ids: [0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2]vision height position_ids: [0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1]vision width position_ids: [0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1]text temporal position_ids: [3, 4, 5, 6, 7]text height position_ids: [3, 4, 5, 6, 7]text width position_ids: [3, 4, 5, 6, 7]Here we calculate the text start position_ids as the max vision position_ids plus 1.

??最后M-RoPE的函數(shù)里面完成嵌入3D RoPE編碼，$q$和$k$分別與角度相乘。不過這里面好像有一個mrope_section，看config里面好像涉及rope_scaling的內(nèi)容，這個學(xué)不動了后面學(xué)學(xué)emmm

query_states, key_states = apply_multimodal_rotary_pos_emb(query_states, key_states, cos, sin, self.rope_scaling["mrope_section"]
)def apply_multimodal_rotary_pos_emb(q, k, cos, sin, mrope_section, unsqueeze_dim=1):mrope_section = mrope_section * 2cos = torch.cat([m[i % 3] for i, m in enumerate(cos.split(mrope_section, dim=-1))], dim=-1).unsqueeze(unsqueeze_dim)sin = torch.cat([m[i % 3] for i, m in enumerate(sin.split(mrope_section, dim=-1))], dim=-1).unsqueeze(unsqueeze_dim)q_embed = (q * cos) + (rotate_half(q) * sin)k_embed = (k * cos) + (rotate_half(k) * sin)return q_embed, k_embed

??借助M-RoPE，論文里面提到qwen2-vl的外推能力挺好。

5. qwen2.5-vl的位置編碼

??最近測試下來qwen2.5-vl的效果沒有qwen2-vl好，可能因為里面用了窗口注意力(qwen-vl里面提到過，說這個效果不如global attention)，qwen2.5能做的任務(wù)比qwen2要好。如果要使用qwen2.5，記得transformer版本安裝方式為

pip install git+https://github.com/huggingface/transformers.git@9985d06add07a4cc691dc54a7e34f54205c04d40

??看上去qwen2.5-vl的位置編碼，和qwen2-vl的主要區(qū)別是position_id這里面，time_id這一維度的計算方式，qwen2.5-vl里面不同的采樣率，會對應(yīng)不同的time_id，和絕對的時間進行對齊。

6.很好的參考資料

1.蘇劍林老師的Transformer升級系列，在他的網(wǎng)站“歸檔”里面進行搜索，可以一章一章的看，例如:

• Transformer升級之路：2、博采眾長的旋轉(zhuǎn)式位置編碼
• Transformer升級之路：4、二維位置的旋轉(zhuǎn)式位置編碼
• “閉門造車”之多模態(tài)思路淺談（三）：位置編碼
• Transformer升級之路：17、多模態(tài)位置編碼的簡單思考

2.eleuther的博客：https://blog.eleuther.ai/rotary-embeddings/

7.TODO

• 天池比賽最近出新的LLM比賽了
• 強化學(xué)習(xí)很多教程云里霧里的，發(fā)現(xiàn)磨菇書非常不錯，代碼還沒看：https://datawhalechina.github.io/easy-rl/#/