♟️

LLM で使われる位置情報のベクトル化について調べてみる

に公開
機械学習
LLM
tech

こんにちは、初めましての方は初めまして。かわらです。転職をして 11 月から所属が変わりました。ただ、分野は変わらない（機械学習がメインの IT エンジニア）ので、これからも気の向くままに機械学習などの技術を調べて zenn に書き残していきたいと思っています。

最近、位置情報のベクトル化の技術の一つに Wavelet-based positional encoding という技術があるのを知りました。気になってどういう技術なのかを調べていたのですが、調べているうちに「他の位置情報の表現方法についても知りたいな」と思い始め、どうせ調べるならまとめてアウトプットしようと思いこの記事を書くことにしました。生成 AI の技術記事に比べると地味であまり面白くない内容になっていますが、興味のある手法の部分だけでも読んでいただければと思います。また、理解をしようとしながらこの記事を書いているため、間違っている部分もあると思います。間違いを見つけた際はコメントなどでこっそり教えてください。

位置情報の表現について

前提として、自己注意機構はその構造的な制約によって「（多重）集合から（多重）集合への写像」となります（位置情報がないと、前と後ろに同じ単語がある場合に内積が同じになり、アテンションの重みが同じになるため）。集合なので、位置などを考慮することが出来ません。その問題を解決するために、Transformer や GPT のような自己注意機構を取り入れたモデルでは位置情報をベクトルにして活用しています。以下では、位置表現についてよく聞く手法をまとめていきます。

Position encoding

Posiion encoding には、大きく分けて絶対位置によるベクトル化と相対位置によるベクトル化の二種類があります。

  • 絶対位置：文の先頭から「1 番目、2 番目、…」としてその位置に対応したベクトルを割り当てる方法。
  • 相対位置：今注目している単語を基準にした相対的な位置に対応したベクトルを割り当てる方法。

以下では、それぞれに関して説明をしていきます。

Absolute position encoding

言わずと知れた、Transformer の論文[1]で使われている絶対位置のベクトル化の方法です。各位置 pos の位置ベクトルを、位置ベクトルの次元 i とベクトルの次元数 d_{model} を用いて以下のように計算します。

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

上の式から分かるように、偶数の次元（2i）には \sin 関数で、奇数の次元（2i+1）には \cos 関数で値を指定しているので、i は実際には 1<=i<=\frac{d_{model}}{2} となります。i によって周波数が変わり、かつ pos によって \sin や \cos の開始位置が変わります。これによってそれぞれの位置のベクトルがそれぞれ変化するので、各位置を表すベクトルの表現として扱えるようになります。

Relative position representation

Relative position representation は Self-Attention with Relative Position Representations で提案された手法で、各単語の相対位置をベクトルにして足し合わせて計算することで位置関係を表現しています。自己注意機構の計算の際に、その相対位置を表すベクトルを足し合わせて「注目している単語からの相対位置」の情報を付与します。ここで「自己注意機構の計算の際に」と書いた部分が注意で、これらの相対位置のベクトルは最初に足すのではなく注意の重みを計算するときに使用します（Absolute position encoding は単語埋め込みに足しています）なので、絶対位置のベクトルと合わせて使うことも出来ます。式での説明は以下のトグルに書いたので、興味がある方は開いてみてください。

式による説明

通常の Transformer モデルでは、入力 x_i から出力 z_i を用いて以下のように計算を行います（ここではマルチヘッドアテンションではなく、一つだけのアテンションを考えていると思ってください。マルチヘッドアテンションになっても、仕組み自体は変わりません）

e_{ij} = \frac{(x_iW^Q)(x_jW^K)^\mathsf{T}}{\sqrt{d_z}} \\ \alpha_{ij} = \frac{\exp e_{ij}}{\sum^n_{k=1}\exp e_{ik}} \\ z_i = \sum^n_{j=1} \alpha_{ij}(x_jW^V)

e_{ij} が i 番目の単語と j 番目の単語の関連度を表し、それを 0~1 の重みにしてから加重平均をとり出力としています。ここで、位置情報がない場合は同じ単語に対する e_{ij} が同じ値になってしまいます。そこで、周辺単語（j 番目の単語）に対してその相対位置を表す a_{ij} を足し合わせることでそれぞれの位置を明示的に取り入れます。変更した式は以下になります。

e_{ij} = \frac{(x_iW^Q)(x_jW^K + a_{ij}^K)^\mathsf{T}}{\sqrt{d_z}} \\ z_i = \sum^n_{j=1} \alpha_{ij}(x_jW^V + a_{ij}^V)

ここで、a_{ij} は相対位置 j-i に対応したベクトルだと思ってもらえればいいです。それぞれの相対位置を一つの単語とみなして単語埋め込みと同じようにベクトルにして、学習をしていきます。これによって、相対位置を表すベクトル表現が得られ、かつ注意機構の計算時に相対位置を考慮することが出来るようになります。

長文は短い文と比較してデータに多く含まれておらず、そのため絶対位置によるベクトル化は長文に対して性能が低下すると言われています。例えば 1 番目の単語はどの文でも出てあっても出てきますが、10000 番目の単語は 10000 単語含まれている文でしか出てこず、そのような文はほとんどないことは想像できるかと思います。

そのような問題に対して、相対位置によるベクトル化はどのような場所であっても相対的な位置だけを考慮すればいいため、長文に対してもその性能は低下しないと考えられ、確かに論文によると絶対位置に比べてその性能は向上しています。しかし、「相対位置に対応したベクトル」をモデル内に持っていないといけないため、パラメータが増加するという欠点があります（まあこのデメリットは絶対位置と同じように固定した値にすることで改善可能かなとも思いますが）

絶対位置と相対位置の優劣について

Position Information in Transformers: An Overview によると、絶対位置と相対位置について長文における精度の差はあまりないと評価されています。実際に絶対位置のベクトルを使っていても、三角関数の性質によって相対位置を表現することができます。また、相対位置のベクトルはパラメータが増える、考慮する相対位置の最小、最大値を人手で決めないといけないといったデメリットがあります。そのためか、あまり使われているモデルはないかと思います。ですが、ここでは「注意機構の部分で足すという方法があるということを知っていても損はないだろう」と判断して書くことにしました。

また絶対位置と相対位置は、元の論文ではそれぞれ「単語埋め込みに」「自己注意機構の前のベクトルに」対して使用されています。ですが、もちろん単語埋め込みに対して相対位置のベクトルを用いてもいいと思いますし、自己注意機構の計算時に絶対位置のベクトルを用いてもいいと思います（もしかしたらそういう研究は既にあるかもしれません）実際に実験をしてどういう構造がいいかを模索するといいかと思います。

Rotary position embedding

Rotary position embedding (RoPE) は ReFormer の論文[2] で提案された手法です。相対位置のベクトル化と同じように、長い文に対しても性能が維持できる手法として提案されました。回転行列を用いて位置を表現することで、絶対位置と相対位置を両方表現できることが特徴になります。

式による説明

RoPE は、クエリとキーのベクトルをその位置に応じて回転させることで位置を表現します。ここで2次元ベクトル (x, y) を角度 \theta だけ回転させる操作を考えると、回転行列を使って以下のように表現できます。

\begin{pmatrix} \cos\theta & -\sin\theta \\ \sin\theta & \cos\theta \end{pmatrix} \begin{pmatrix} x \\ y \end{pmatrix}

RoPE ではこの考えを拡張して、各位置 m に対して角度 m\theta で回転させます。式で示すと以下になります。

f_{\{q,k\}}(\mathbf{x}_m, m) = \mathbf{R}^d_{\Theta,m}\mathbf{W}_{\{q,k\}}\mathbf{x}_m \\ \mathbf{R}^d_{\Theta,m} = \begin{pmatrix} \cos m\theta_1 & -\sin m\theta_1 & 0 & 0 & \cdots & 0 & 0 \\ \sin m\theta_1 & \cos m\theta_1 & 0 & 0 & \cdots & 0 & 0 \\ 0 & 0 & \cos m\theta_2 & -\sin m\theta_2 & \cdots & 0 & 0 \\ 0 & 0 & \sin m\theta_2 & \cos m\theta_2 & \cdots & 0 & 0 \\ \vdots & \vdots & \vdots & \vdots & \ddots & 0 & 0 \\ 0 & 0 & 0 & 0 & \cdots & \cos m\theta_{d/2} & -\sin m\theta_{d/2} \\ 0 & 0 & 0 & 0 & \cdots & \sin m\theta_{d/2} & \cos m\theta_{d/2} \\ \end{pmatrix}

ここで \Theta=\{ \theta_i = 10000^{-2(i-1)/d}, i\in[1,2,...,d/2 ] \} です。式を見ると分かるように、上から順に二つずつの要素を回転させていることが分かります。Absolute position encoding では位置表現を単語埋め込みに足し合わせていましたが、RoPE では単語埋め込みの要素を直接回転させます。論文の図を以下に引用しておきます。この図をみると、それぞれ二要素ずつ回転させていることが分かりやすいかと思います。


Reformer の論文図 1 より引用

また、m 番目のクエリと n 番目のキーの内積を考えると以下になります。これにより、RoPE で相対位置 n-m を考慮した計算も出来ていることが分かります。さらに、これは内積だけであり通常の Transformer の処理と同じであるため、特に工夫を凝らさず自然と相対位置を計算していることも分かります。

\mathbf{q}^{\mathsf{T}}_m\mathbf{k}_n = (\mathbf{R}^d_{\Theta,m}\mathbf{W}_{q}\mathbf{x}_m)(\mathbf{R}^d_{\Theta,n}\mathbf{W}_{k}\mathbf{x}_n) = \mathbf{x}^{\mathsf{T}}\mathbf{W}_{q}\mathbf{R}^d_{\Theta,n-m}\mathbf{W}_{k}\mathbf{x}_n

RoPE は相対位置を扱えるため、絶対位置では苦手だった長文に対しても性能が維持できます。また、単語の位置だけから位置のベクトルを作成できるので、追加のパラメータが不要になります。この手法は実際に Llama や GPT-NeoX など多くの大規模言語モデルで採用されており、現在の主流の位置のベクトル化の手法の一つとなっています。

Attention with Linear Biases

Attention with Linear Biases (ALiBi) は Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation で提案された手法です。その名の通り、注意スコアに線形のバイアスを足すというシンプルなアイデアで位置情報を表現しています。

結構シンプルなのでここでは式の説明をトグルにせずにそのまま書いておきます（興味なければ飛ばしてください）

自己注意機構の注意の重みは、\mathrm{softmax}(\mathrm{\mathbf{q}}_i \mathrm{\mathbf{K}}^\mathsf{T}) と表すことができます。ここで i 番目のクエリベクトルを \mathrm{\mathbf{q}}_i としています。ALiBi ではこの自己注意機構の重みの計算時に、以下の式のようにそれぞれの相対位置に応じた定数を足し合わせます。

\mathrm{softmax}(\mathrm{\mathbf{q}}_i \mathrm{\mathbf{K}}^\mathsf{T} + m\cdot [-(i-1),\cdots,-2,-1,0])

ここで m は各ヘッドごとに異なる定数です。下に論文から引用した図も載せておきます。


AliBi の論文図 3 より引用

図を見ると分かりやすいですが、クエリとキーの相対位置に基づいて定数を足し合わせているだけです。そのため、追加のパラメータを持たずに位置情報を表現できることが分かります。このような単純な計算にも関わらず、論文の実験では他の手法と比較して長文で性能が維持できていることが示されています。

この手法と結果を見ると、とりあえず位置毎に区別できればどのような値でも一定の効果が出るように思えますね。「区別できる値を使えば」そう思っていたのですが、次に紹介する手法ではなんと位置情報を使いません！（ネタバレ）

No positional encoding

No positional encoding (NoPE)[3][4] はその名の通り、明示的な位置表現を使わない手法です。明示的な位置表現を使わないと、初めの方に説明したように自己注意機構は「集合から集合への写像」となるため、順序を考慮できなくなります。「え、じゃあトークンの位置関係が分からなくなってしまうのでは？」と疑問に思うかもしれません。

論文によると、明示的な位置符号化がなくてもモデルは位置情報をある程度考慮できることが示されています。Causal Attention を使う場合、「自分より前のトークンしか見られない」という制約自体が位置情報を暗黙的に与えていると考えられるらしいです。

論文の実験結果では、NoPE を使ったモデルでもある程度の性能を達成できることが示されています。また、タスクによっては文が長くなっても精度が低下しないことも示されています。ただ、やはり NoPE だけでは位置情報を表すことが難しいのか、実際のモデルでは層毎に NoPE や RoPE, Relative position encoding を取り入れているようです。

Wavelet-based positional representation (WPR)

WPR[5] はウェーブレット変換を使用して位置を表現します。正直あんまり理解できている気はしていないのですが、とりあえずメモを残します。「これを理解するついでに他の手法を調べたはずなのに分かってないの？」というツッコミは無視します。

ウェーブレット変換自体は信号処理の分野でよく使われる手法で、信号を異なる周波数成分に分解する技術です。Wavelet-based positional encoding は、この性質を利用して位置情報を複数の成分で表現します。論文中では RoPE がウェーブレット変換の一種であることを示しています。つまり、RoPE よりも柔軟な位置情報を扱うことが出来ます。

実験では WPR が他の手法と比較して一貫していい性能を発揮できてました。また、長文に対しても他の手法と比較してうまく情報をとらえられていることが分かりました。またこの結果は大規模なモデルに対しても有効で、具体的には Llama-3-8B モデルでも RoPE と比較して性能が向上しており、この手法の有効性が伺えます。

WPR を活用した LLM はまだ聞きませんが、そのうち WPR を活用して長文に対しての性能がより向上した LLM が出てくるかもしれません。いや、もしかしたらそのモデルを開発するのは君かもしれない！（投げやり）

まとめ

この記事では Transformer やその派生モデルにおける位置情報のベクトル化についてまとめました。個人的には NoPE が結構興味深く、「位置表現は絶対に必要」という固定概念を揺さぶられた気分です。位置表現、なかなかに奥深い… また、こうやって並べてみると、位置表現は「長文の精度」を改善する試みだととらえることもできて、特に最近の長文ブーム（？）ではより一層この「位置表現」の工夫が LLM の性能に効いてくるのかなと思います。

以上、かわらでした。何か間違いや補足などあればコメントでも DM でも飛ばしてもらえるとありがたいです。

脚注

  1. https://arxiv.org/abs/1706.03762 ↩︎

  2. https://arxiv.org/abs/2104.09864 ↩︎

  3. https://arxiv.org/abs/2203.16634 ↩︎

  4. https://arxiv.org/abs/2305.19466 ↩︎

  5. https://arxiv.org/abs/2502.02004 ↩︎

GitHubで編集を提案

Discussion