Transformer の位置情報

Transformer は順序に依存しない自己注意機構を採用しているため、入力シーケンス内の各トークンの位置情報を明示的にモデルに入れ込む必要があります。

そして、この位置情報の入れ込み方にいくつかバリエーションがあります。

Absolute Positional Embeddings (Positional Encoding)
Relative Positional Embeddings (Relative PE)
Rotary Positional Embeddings (RoPE)
Attention with Linear Biases (ALiBi)
Transformer without Positional Embeddings (NoPE)

それぞれ解説していきます。

Absolute Positional Embeddings (APE) あるいは Positional Encoding

Attention is all you need の最初の Transformer に採用されている Positional Embeddings。
最初は Positional Encoding という名前でした。

APE では、sin と cos 関数を使用し、位置 $pos$ と次元 $i$、モデルの埋め込み次元 $d_{model}$ に対して、次のように計算されます。

$PE(pos, 2i) = \sin(pos / 10000^{2i/d_{model}})$
$PE(pos, 2i+1) = \cos(pos / 10000^{2i/d_{model}})$

APE は実装が容易ですが、シーケンスの長さが変動する場合（例：学習した文より長い文を扱う場合）に性能が悪くなるという欠点がありました。

Relative Positional Embeddings (Relative PE)

APE は、シーケンスの長さが変わる場合に性能が悪くなるという欠点があったので、相対的な位置を見る Relative PE が提案されました。

Relative PE では、相対的な位置情報を self-attention 機構に組み込んでいます。

変更前
- $ e_{ij} = \frac{(x_iW^Q)(x_jW^K)^T}{\sqrt{d_z}} $
変更後
- $ e_{ij} = \frac{(x_iW^Q)((x_jW^K + a^K_{ij})^T)}{\sqrt{d_z}} $
  - ここで $a^K_{ij}$ は入力要素 $x_i$ と $x_j$ 間の相対位置を表すベクトル

これによって、シーケンスの長さが変わっても対応できるようになった一方、計算コストが高くなるという欠点が出てきました。

Rotary Positional Embeddings (RoPE)

RoPE は、Relative PE と同じように、APE の「シーケンスの長さが変わる場合、性能が悪くなる」という欠点に対応するため提案されました。

RoPE では以下のような回転行列を使います。

R_i = \begin{bmatrix} \cos(\theta_i) & -\sin(\theta_i) \\ \sin(\theta_i) & \cos(\theta_i) \end{bmatrix}

$θ_i$は位置$i$に応じた回転の角度です。回転を利用することで、トークン間の相対距離が増加するにあたって内積が減衰するので相対位置を表現しています。

RoPE は LLaMA 系のモデルなど広く利用されており、現在の主流な位置エンベディングです。

Attention with Linear Biases (ALiBi)

シーケンスの長さの変化に対応すべく Relative PE や RoPE が提案されたものの、Relative PE は効果あるけどメモリも食う上に遅く、RoPE は Relative PE ほど効果がありませんでした。
それならもっと早くて性能も良いものを作ろうとしてできたのが ALiBi です。

ALiBiでは、位置情報として、クエリキーペア間の距離に基づいた静的なバイアス $M$ を加算します。

 \text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T + M}{\sqrt{d_k}}\right)V

ここで、$Q, K, V$ はそれぞれクエリ、キー、バリューの行列、$d_k$ はキーベクトルの次元です。
静的なバイアス $M$ は行列で表され、その要素 $M_{ij}$ は $-m \times |i-j|$（ここで $m$ はハイパーパラメータ）によって計算されます。この式では遠くのペアよりも近いペアのスコアが相対的に高くなります。

ALiBi は BLOOM に利用されています。

Transformer without Positional Embeddings (NoPE)

位置情報を入れてもそれほど効果が見られないなら、いっそのこと抜いてしまおうというのが NoPE です。

もちろん、位置情報を抜くことで、シーケンス内のトークン間の位置関係をモデルが把握できず、これにより、位置情報が重要タスクにおいては、パフォーマンスの低下が起きる可能性があります。

ただし、位置に依存する Absolute Positional Embeddings と比較すると、固定された位置パターンに依存しなくなるため、任意のシーケンス長に対応できますし、そもそも計算工程を一つ省くため計算量が少なくなります。

NoPE は比較的新しい手法のため、NoPE を使った有名モデルはまだ見たことがありませんが、メリット・デメリットを理解して使う分には有効な手法になる可能性があります。