#

RoPer はジョルジュ・ハリク (@gharik) の作品で、この実装は彼のオリジナルコードに基づいています。

相対距離付き回転式位置埋め込み（RoPer）

ロータリー・ポジション・エンベディング（RoPE）では、アテンション・スコアの計算に相対的な位置が含まれます。ただし、埋め込み自体は、因果関係から暗黙的に取得できるものを除いて、位置情報を取得しません

。

RoPer は値の埋め込みに相対位置情報を明示的に追加します。具体的には、注目したトークンの相対的な位置を追加します。同じ回転位置埋め込みを使用して注目している値を回転させ、加重合計を取った後、最後の埋め込みを反対方向に回転させます。これは、現在の位置を基準に（注意の前に）各値を回転させるのと同じです。

以下は、算術加算で RoPer を使用してトランスフォーマーモデルをトレーニングするためのトレーニングコードです。RoPe よりも大幅に改善されていることがわかります。

埋め込みの相対距離

どんな頭でも、 $n$ ポジションからポジションへと注目し $i$ 、 $v_{i}$ ポジションに価値を埋め込むようにしましょう $a_{n, i}$ 。 $i$ 個々の機能を次のように表してみましょう

。

v_{i}^{(1)}, v_{i}^{(2)}, \dots

通常は、値の埋め込みの加重和を取ります。

$o_{n}^{(j)} = i \sum a_{n, i} v_{i}^{(j)}$

これにより、 $i$ $o_{n}^{(j)}$ 位置に関する距離情報が最終結果に明示的に追加されることはありません。

RoPer は RoPe などの機能を組み合わせて変換します。ペアの場合 $v_{m}^{(1)}$ $v_{m}^{(2)}$ 、次の方法で変換されます $R o PE (v_{m}^{(1)}, v_{m}^{(2)}, m)$ 。変換された機能を寄付しましょう。 $\overset{v}{^}_{m}^{(1)}, \overset{v}{^}_{m}^{(2)}$ 次に、 $\overset{o}{^}_{n}^{(j)}$ 加重和をと逆方向に回転させます。 $R o PE (\overset{o}{^}_{n}^{(1)}, \overset{o}{^}_{n}^{(2)}, - n)$ 注意してください $- n$ .

次の点に注意してください。

R o PE (x_{m}^{(1)}, x_{m}^{(2)}, m) = (c o s m θ s i n m θ - s i n m θ c o s m θ) (x_{m}^{(1)} x_{m}^{(2)}) = (x_{m}^{(1)} c o s m θ - x_{m}^{(2)} s i n m θ x_{m}^{(2)} c o s m θ + x_{m}^{(1)} s i n m θ)

変換後の最終出力は、

R o PE (\overset{o}{^}_{n}^{(1)}, \overset{o}{^}_{n}^{(2)}, - n) (\overset{o}{^}_{n}^{(1)} c o s n θ + \overset{o}{^}_{n}^{(2)} s i n n θ \overset{o}{^}_{n}^{(2)} c o s n θ - \overset{o}{^}_{n}^{(1)} s i n n θ) =

その点に注意してください $s i n (- n θ) = - s i n n θ$ 。

$\overset{o}{^}_{n}^{(1)} c o s n θ + \overset{o}{^}_{n}^{(2)} s i n n θ$ 第1学期を広げてみよう

\overset{o}{^}_{n}^{(1)} c o s n θ + \overset{o}{^}_{n}^{(2)} s i n n θ i \sum a_{n, i} \overset{v}{^}_{i}^{(1)} c o s n θ + i \sum a_{n, i} \overset{v}{^}_{i}^{(2)} s i n n θ i \sum a_{n, i} (v_{i}^{(1)} c o s i θ - v_{i}^{(2)} s i n i θ) c o s n θ i \sum a_{n, i} (v_{i}^{(2)} c o s i θ + v_{i}^{(1)} s i n i θ) s i n m θ i \sum a_{n, i} v_{i}^{(1)} (c o s i θ c o s n θ + s i n i θ s i n n θ) i \sum a_{n, i} v_{i}^{(2)} (c o s i θ s i n n θ - s i n i θ c o s n θ) i \sum a_{n, i} v_{i}^{(1)} c o s (i - n) θ - i \sum a_{n, i} v_{i}^{(2)} s i n (i - n) θ i \sum a_{n, i} v_{i}^{(1)} c o s (i - n) θ - i \sum a_{n, i} v_{i}^{(2)} s i n (i - n) θ = = + = + = =

同様に、2 番目の項が次の式と等しいことがわかります。

$i \sum a_{n, i} v_{i}^{(1)} c o s (i - n) θ + i \sum a_{n, i} v_{i}^{(2)} s i n (i - n) θ$

これにより、

R o PE (\overset{o}{^}_{n}^{(1)}, \overset{o}{^}_{n}^{(2)}, - n) (\sum_{i} a_{n, i} v_{i}^{(1)} c o s (i - n) θ - \sum_{i} a_{n, i} v_{i}^{(2)} s i n (i - n) θ \sum_{i} a_{n, i} v_{i}^{(1)} c o s (i - n) θ + \sum_{i} a_{n, i} v_{i}^{(2)} s i n (i - n) θ) i \sum a_{n, i} R o PE (v_{i}^{(1)}, v_{i}^{(1)}, (i - n) θ) = =

つまり、現在の位置を基準にローテーションされた値の加重平均です。

これは、算術加算タスクにRoPerを使用する実験です。

118from typing import Optional
119
120import torch
121
122from labml_nn.transformers.rope import RotaryPositionalEmbeddings, RotaryPEMultiHeadAttention

#

反対方向に回転するRoPEモジュール

これはRoPE回転の実装を継承し、方向を変えます。

125class ReverseRotaryPositionalEmbeddings(RotaryPositionalEmbeddings):

#

x キーまたは形状のあるクエリの先頭にあるテンソルです [seq_len, batch_size, n_heads, d]

132    def forward(self, x: torch.Tensor):

#

$c o s$ $s i n$ キャッシュと値

137        self._build_cache(x)

#

機能を分割して、一部の機能セットにのみロータリー埋め込みを適用することもできます。

140        x_rope, x_pass = x[..., :self.d], x[..., self.d:]

#

計算 $[- x^{(\frac{d}{2} + 1)}, - x^{(\frac{d}{2} + 2)}, ..., - x^{(d)}, x^{(1)}, x^{(2)}, ..., x^{(\frac{d}{2})}]$

144        neg_half_x = self._neg_half(x_rope)

#

計算

(x_{m}^{(i)} c o s - m θ_{i} - x_{m}^{(i + \frac{d}{2})} s i n - m θ_{i} x_{m}^{(i + \frac{d}{2})} c o s - m θ_{i} + x_{m}^{(i)} s i n - m θ_{i}) = (x_{m}^{(i)} c o s m θ_{i} + x_{m}^{(i + \frac{d}{2})} s i n m θ_{i} x_{m}^{(i + \frac{d}{2})} c o s m θ_{i} - x_{m}^{(i)} s i n m θ_{i})

にとって $i \in 1, 2, ..., \frac{d}{2}$

160        x_rope = (x_rope * self.cos_cached[:x.shape[0]]) - (neg_half_x * self.sin_cached[:x.shape[0]])

#

163        return torch.cat((x_rope, x_pass), dim=-1)

#

回転式位置埋め込みによるマルチヘッドアテンション

オリジナルのトランスフォーマーのマルチヘッドアテンションを無効にします。

166class RotaryValuePEMultiHeadAttention(RotaryPEMultiHeadAttention):

#

173    def __init__(self, heads: int, d_model: int,
174                 rope_percentage: float = 0.5, rope_value_percentage: float = 0.5,
175                 dropout_prob: float = 0.0):
176        super().__init__(heads, d_model, rope_percentage, dropout_prob)

#

ロータリーポジショナル埋め込みレイヤー

179        d_rope_value = int(self.d_k * rope_value_percentage)
180
181        self.value_rotary_pe = RotaryPositionalEmbeddings(d_rope_value)
182        self.value_reverse_rotary_pe = ReverseRotaryPositionalEmbeddings(d_rope_value)

#

query 、key value およびは、クエリ、キー、および値のベクトルのコレクションを格納するテンソルです。形があります[seq_len, batch_size, d_model] 。

mask [seq_len, seq_len, batch_size] 形状があり、バッチの場合b 、mask[i, j, b] i その位置のクエリがその位置のキー値にアクセスできるかどうかを示します。j

184    def forward(self, *,
185                query: torch.Tensor,
186                key: torch.Tensor,
187                value: torch.Tensor,
188                mask: Optional[torch.Tensor] = None):

#

query 、key value そして形がある [seq_len, batch_size, d_model]

200        seq_len, batch_size, _ = query.shape
201
202        if mask is not None:
203            mask = self.prepare_mask(mask, query.shape, key.shape)

#

query key value 注意力計算の準備をして[seq_len, batch_size, heads, d_k] これで形ができあがります。

207        query = self.query(query)
208        key = self.key(key)
209        value = self.value(value)

#

$Q K^{⊤}$ アテンションスコアを計算します。[seq_len, seq_len, batch_size, heads] これにより形状のテンソルが得られます

。

213        scores = self.get_scores(query, key)

#

スケールスコア $\frac{Q K ^{⊤}}{d _{k}}$

216        scores *= self.scale

#

マスクを適用

219        if mask is not None:
220            scores = scores.masked_fill(mask == 0, float('-inf'))

#

$so f t ma x$ キーシーケンス次元に沿って注目 $se q so f t ma x (\frac{Q K ^{⊤}}{d _{k}})$

224        attn = self.softmax(scores)

#

ドロップアウトを適用

227        attn = self.dropout(attn)

#

位置情報を含むように、加重合計を取る前に値の埋め込みをローテーションしてください

230        value = self.value_rotary_pe(value)

#

値による乗算 $se q so f t ma x (\frac{Q K ^{⊤}}{d _{k}}) V$

234        x = torch.einsum("ijbh,jbhd->ibhd", attn, self.value_rotary_pe(value))

#

各埋め込み部品が相対的な位置を保持するように、反対方向に回転させます

237        x = self.value_reverse_rotary_pe(x)

#

他の計算に注意を向けておく

240        self.attn = attn.detach()

#

複数のヘッドを連結

243        x = x.reshape(seq_len, batch_size, -1)

#

出力レイヤー

246        return self.output(x)