位置を回転で入れる

目標

回転位置埋め込み(RoPE)を、標準ライブラリだけで自分で作ります。偶数次元を2つずつペアにまとめ、ペアごとに異なる速度で位置の分だけ回し、そうして回したクエリとキーの内積が2つの位置の差だけに依存することを、数字で確認します。足す方式(絶対位置エンコーディング)を並べて作り、同じ間隔で値が揺らぐ様子を比べ、回転が長さを変えないことと、距離が離れるときにどの成分が残るかを測ります。

なぜ重要なのか

位置ベクトルを足すと、内積を展開したときにq·P[n]とP[m]·kのような片方の位置しか持たない項が残ります。その項は差としてまとめられないため、間隔がまったく同じでも、文の前のほうで測った値と後ろのほうで測った値が異なります。言語で重要なのは、たいてい何語前の単語かなのに、スコアには絶対位置が混ざり込んでいるのです。回転は、その問題を演算そのもので解決します。2つのベクトルを同じ向きにいっしょに回すと、間の角が変わらないので、それぞれ自分の位置の分だけ回してから内積をとると、値が差だけに依存します。これは近似ではなく等式なので、数字で確認できます。このラボは、モデルを呼び出しません。このPodのシステムのPythonにはnumpyがなく、インターネットもありません。mathだけで作り、ここで測った数字だけを使います。そのため、「あるモデルはどんな底を使う」といった話は、ここではしません。採点ツールは、書いておいた説明を信用しません。自分で作ったモジュールを実際に呼び出し、毎回異なる次元と位置で関数を直接叩いて、採点ツールが別に計算した値と、許容誤差の範囲内で照合します。入力は実行のたびに変わるので、値を暗記して入れることはできません。

ステップ

/root/work/tf-rope/rope.pyにDIM = 64・THETA_BASE = 10000.0・demo_vectors()・dot(a, b)・thetas(d, base=THETA_BASE)を作成してください。thetasは、ペアごとに異なる回転速度を返します。
rotate_pair(x0, x1, angle)を追加して、平面上の点1つをangleラジアンだけ反時計回りに回すようにしてください。
apply_rope(vec, pos, base=THETA_BASE)を追加して、ベクトルをペアにまとめ、各ペアをpos * theta_iだけ回した新しいベクトルを返すようにしてください。
rope_score(q, k, m, n, base=THETA_BASE)を作成して、m番目の位置のクエリとn番目の位置のキーをそれぞれ回してから内積をとるようにしてください。
sin_pos(pos, d, base=THETA_BASE)とadd_score(q, k, m, n, base=THETA_BASE)を作成して、足す方式の同じスコアを測るようにしてください。
offset_scan(q, k, offset, starts)を作成して、同じ間隔を複数の開始位置で測るようにしてください。返す値は、(시작자리, 회전 점수, 더하는 점수)(プレースホルダーは開始位置、回転のスコア、足す方式のスコアです)のペアのリストです。
turns(d, delta, base=THETA_BASE)とslow_pairs(d, delta, base=THETA_BASE)を作成して、間隔ごとにペアが何周回るか、まだ1周を超えていないペアがいくつあるかを測るようにしてください。
測った値を、/root/work/tf-rope/rope_report.jsonと、/root/work/tf-rope/rope_report.mdに記録してください。

参考

実行の契約: 採点ツールは、/root/work/tf-rope/rope.pyをPythonモジュールとして読み込み、DIM・THETA_BASE・demo_vectors・dot・thetas・rotate_pair・apply_rope・rope_score・sin_pos・add_score・offset_scan・turns・slow_pairsを直接使います。スクリプトとして実行しないので、if __name__ == "__main__"はなくてかまいません。
demo_vectors()は、(q, k)を返します。乱数ではなく、式で固定します。q[j] = math.cos(0.7 * j + 0.3)、k[j] = math.sin(0.4 * j + 1.1)で、どちらも長さはDIMです。
thetas(d, base)は、長さd // 2のリストで、i番目の値はbase ** (-2 * i / d)です。最初の値はつねに1.0で、後ろへ行くほど小さくなります。次元ごとではなく、ペアごとに1つです。
rotate_pair(1.0, 0.0, math.pi / 2)は、(0.0, 1.0)です。反時計回りに固定します。返す値は、要素が2つのタプルです。
apply_ropeは、(vec[0], vec[1])、(vec[2], vec[3])のように、隣同士でペアにまとめます。i番目のペアの回転角はpos * thetas(len(vec), base)[i]で、ペアごとに角が異なります。受け取ったリストをその場で書き換えず、新しいリストを作ってください。posが0なら、もとのベクトルと同じです。
rope_score(q, k, m, n)は、dot(apply_rope(q, m), apply_rope(k, n))です。√dでは割りません。
sin_pos(pos, d, base)は、長さdのリストで、j番目の値は、angle = pos / (base ** ((2 * (j // 2)) / d))としたときに、jが偶数ならmath.sin(angle)、奇数ならmath.cos(angle)です。
add_score(q, k, m, n)は、qにsin_pos(m, len(q))を、kにsin_pos(n, len(k))を、それぞれ足してから内積をとります。片方にだけ足してはいけません。
offset_scan(q, k, offset, starts)は、startsの位置sごとに、クエリをs + offsetに、キーをsに置きます。順序はstartsと同じです。
turns(d, delta, base)は、長さd // 2のリストで、i番目の値はdelta * theta_i / (2 * math.pi)です。slow_pairsは、その値が1.0より小さいペアの個数(整数)です。
ステップ8のレポートは、demo_vectors()のqとkで測ります。間隔は2、開始位置は3・10・100・4000、ノルムを測る位置は0・1・7・100・4096、周回数を数える間隔は1・16・256・4096です。JSONのキーは、dim・base・offset・starts・rope_scores・rope_spread・add_scores・add_spread・norm_before・norm_max_gap・deltas・slow_pairs・turns_first・turns_lastです。turns_firstは、間隔ごとのペア0の周回数で、turns_lastは、最後のペアの周回数です。
rope_spreadとadd_spreadは、それぞれ4つのスコアの最大値から最小値を引いた値です。norm_beforeはqのノルムで、norm_max_gapは、5つの位置で回したあとのノルムともとのノルムの差のうち、最も大きい絶対値です。
実数を==で比べないでください。採点ツールは、abs(a - b) <= 1e-9 + 1e-6 * abs(b)で判定します。位置が大きくなると角度も大きくなり、誤差が増えるので、このラボは位置を0から4096までだけ判定します。
このPodにはインターネットがありません。pip installはできず、numpyは/opt/onnx-lab/bin/pythonの中にしかないので、システムのPythonではimport numpyができません。mathだけで十分です。
公式ドキュメント: RoFormer — Rotary Position Embedding・Attention Is All You Need・Python math
よくある間違い: ペアではなく次元ごとに角を付ける、すべてのペアにペア0の角を使う、回転の向きを逆にする、クエリだけ回してキーはそのままにする、足す方式で片方にだけ位置ベクトルを足す、周回数を数えるときに2πで割らない、受け取ったリストをその場で書き換える。

ペアごとに異なる回転速度を作る

/root/work/tf-rope/rope.pyにDIM = 64・THETA_BASE = 10000.0・demo_vectors()・dot(a, b)・thetas(d, base=THETA_BASE)を作成してください。thetasは、長さd // 2のリストで、i番目の値はbase ** (-2 * i / d)です。demo_vectors()は、q[j] = math.cos(0.7 * j + 0.3)、k[j] = math.sin(0.4 * j + 1.1)である、長さDIMの2つのリストを返します。

次元ではなくペアごとに1つです。64次元なら32個です。iが0のとき指数が0なので、最初の値は1.0で、後ろへ行くほど指数関数的に小さくなります。割り算を整数で行うと、指数がすべて0になって値がすべて1.0になるので、-2.0 * i / dのように実数にしてください。dotは、zipでペアにして掛けてから足せば1行です。

平面上で1回回す

rotate_pair(x0, x1, angle)を追加してください。点(x0, x1)をangleラジアンだけ反時計回りに回した、要素が2つのタプルを返します。rotate_pair(1.0, 0.0, math.pi / 2)は、(0.0, 1.0)です。

(x0*cos - x1*sin, x0*sin + x1*cos)です。2つの符号のうち片方だけ反転しても時計回りになり、後ろのステップの値がすべて変わります。角が0ならもとの点のままで、どんな角を与えても、原点からの距離が変わってはいけません。cosとsinの2乗を足すと1になるからです。

ベクトルを位置の分だけ回す

apply_rope(vec, pos, base=THETA_BASE)を追加してください。(vec[0], vec[1])、(vec[2], vec[3])のように隣同士でペアにまとめ、i番目のペアをpos * thetas(len(vec), base)[i]だけ回した新しいリストを返します。受け取ったリストはそのままにしてください。

ペアごとに角が異なります。ペア0の角をすべてのペアに使うと、単にまるごと回しただけになり、位置の情報が1重分しか入りません。posが0なら、すべての角が0なので、もとのベクトルと同じになる必要があり、どんな位置を与えても、ベクトルのノルムは変わってはいけません。baseを受け取ってthetasに渡すのを忘れないでください。デフォルト値だけを使うと、別の底で呼び出したときに静かに間違います。

回したクエリとキーのスコアを出す

rope_score(q, k, m, n, base=THETA_BASE)を作成してください。m番目の位置のクエリとn番目の位置のキーをそれぞれ回してから内積をとった実数を返します。√dでは割りません。

dot(apply_rope(q, m, base), apply_rope(k, n, base))の1行です。片方だけ回すと、絶対位置がそのまま残り、後ろのステップの性質が崩れます。mとnが同じなら、回していないdot(q, k)と同じ値になる必要があります。同じ向きにいっしょに回したので、間の角がそのままだからです。

足す方式を並べて置く

sin_pos(pos, d, base=THETA_BASE)とadd_score(q, k, m, n, base=THETA_BASE)を作成してください。sin_posは長さdの位置ベクトルで、j番目の値は、angle = pos / (base ** ((2 * (j // 2)) / d))としたときに、jが偶数ならmath.sin(angle)、奇数ならmath.cos(angle)です。add_scoreは、qとkにそれぞれ位置ベクトルを足してから内積をとります。

こちらは、値を作って足す方式です。片方にだけ足すと、比較が成り立ちません。内積を展開すると、q·kのほかにq·P[n]とP[m]·kが出てきますが、この2つの項には片方の位置しか入っておらず、差としてまとめられません。次のステップで、その結果を数字で見ることになります。

同じ間隔なら同じスコアになる

offset_scan(q, k, offset, starts)を作成してください。startsの位置sごとに、クエリをs + offsetに、キーをsに置き、2つの方式でスコアを測ります。返す値は、(s, 회전 점수, 더하는 점수)(プレースホルダーは回転のスコアと足す方式のスコアです)のペアのリストで、順序はstartsと同じです。作成したら、offset_scan(q, k, 2, [3, 10, 100, 4000])を直接出力して、回転のほうの4つの値が同じで、足すほうは揺らぐことを、目で確認してください。

前に作ったrope_scoreとadd_scoreをそのまま呼び出せば、5行です。クエリとキーを逆に置くと、間隔の符号が反転して、別の値が出ます。回転のほうの4つの値は、小数点以下12桁まで同じなら正常です。完全に同じにはならず、ごく小さな差が残りますが、それが浮動小数点の誤差です。足すほうのばらつきと、桁数を比べてみてください。

離れるとどのペアが残るかを調べる

turns(d, delta, base=THETA_BASE)とslow_pairs(d, delta, base=THETA_BASE)を作成してください。turnsは、長さd // 2のリストで、i番目の値はdelta * theta_i / (2 * math.pi)、つまりそのペアが回る周回数です。slow_pairsは、その値が1.0より小さいペアの個数(整数)です。

1周は2πラジアンです。割るのを忘れると、周回数ではなく角度になり、個数がまったく変わります。1周を超えたペアは、その距離と、1周分を引いた距離を同じ角度で表すので、両者を見分けられません。間隔1・16・256・4096で呼び出して、残るペアの数がどのように減っていくかを見てください。

測った結果を記録に残す

demo_vectors()のqとkで測り、/root/work/tf-rope/rope_report.jsonにdim・base・offset・starts・rope_scores・rope_spread・add_scores・add_spread・norm_before・norm_max_gap・deltas・slow_pairs・turns_first・turns_lastを書き、/root/work/tf-rope/rope_report.mdに## 무엇을 쟀나 ## 같은 간격이면 같은 점수다 ## 더하는 방식은 왜 다른가 ## 멀어지면 어느 성분이 남는가の4つの節で書いてください(韓国語の見出しは順に、「何を測ったか」「同じ間隔なら同じスコアになる」「足す方式はなぜ違うのか」「離れるとどの成分が残るのか」という意味です)。間隔は2、開始位置は3・10・100・4000、ノルムを測る位置は0・1・7・100・4096、周回数を数える間隔は1・16・256・4096です。

数字は手で書かず、自分のコードを実際に動かして得た値で埋めてください。rope_scores・add_scoresは、offset_scanが返したペアから取り出します。rope_spread・add_spreadは、それぞれ最大値から最小値を引いた値で、norm_max_gapは、5つの位置で測った差の絶対値のうち最も大きいものです。turns_firstは、間隔ごとのペア0の周回数で、turns_lastは、最後のペアの周回数です。レポートの本文には、回転のほうの最初のスコア・足すほうのばらつき・qのノルムを数字で書いてください。採点ツールが、その3つの値が文章の中にあるかを見ます。