やぁ〜ん
こんにちは、データサイエンティストをしている服部です。
ABEJAアドベントカレンダー2025の10日目の記事です。
LLMといえばロングコンテキスト大事ですよね（唐突） そんなLLMのロングコンテキストを支える重要技術である「YaRN」を紹介したいと思います。
https://arxiv.org/abs/2309.00071
脳内再生されるあれ
実際になんと読むのが正解なのか知りませんが、私の頭の中ではこれしか流れません。
「やぁ〜ん」
どないやねん ヤドン
ということで、「やぁ〜ん」っとヤドンを題材にしながらヤドンでも分かるようにYaRNを紹介したいと思います。
YaRNの正式名称
論文によると正式名称は「Yet another RoPE extensioN method」....
また別のRoPE拡張方法ということみたいなのですが、かなり無理矢理なやぁ〜んという感じですね。 よほど、好きだったんでしょうか。
RoPE
さて、RoPEの拡張方法ということは、流石にRoPEを避けては説明できなさそうです。ということでRoPEから説明したいと思います。 ヤドンでも分かる感じにしようと思いましたがいきなり厳しいかもしれません....
そもそもRoPEというのは、LLMに入力するトークンの並びに、各トークンがどの位置にあるものか位置情報を付与する方法の一つです。 Positional EncodingとかPositional Embeddingという言い方をします。 この位置情報がないと、LLMは各トークンがまとめて入力された時に、それぞれが何番目のトークンか分かりません。 「ヤ」「ド」「ン」と順番に入力しても、「ヤドン」なのか「ドンヤ」なのか区別がつきません。 「シェルダーがヤドンを噛んだ」のか「ヤドンがシェルダーを噛んだ」といった主語の認識も出来ません。
この位置情報の付与のさせかたについては、様々な研究がされてきました。
RoPEは、「Rotary Position Embedding」という名前で、ベクトルを回転させることで、この位置情報を付与させるという考え方で、現在主流な方法です。
回転による位置情報付与
ただベクトル？回転？位置情報？とイメージしづらいので、まずざっくりしたイメージを図示してみます。
私はヤドンです。やぁ〜ん という文章で考えてみます。
まずこれをLLMにわたすようにトークンに分割して、それぞれをトークンEmbedding（ベクトル化）します。
そして、それらのベクトルを回転させるわけですが、位置が進むにつれて徐々に増えるように回転量を変えていきます。
なんで回転させるの？なんで回転させる量が変わるの？と疑問に思う方もいるかもしれません。 まずここでやりたいのは、「トークン同士が相対的にどれくらい離れているか」の情報付与です。
例えば、「私はヤドンです。」という文章に対して「私」と「ヤ」は「は」の1トークン分挟んで離れてるのに対して「私」と「ン」は「は」「ヤ」「ド」の3トークン分挟んで離れています。 このようなトークン同士の距離を測りたいわけですが、距離を計算するのに回転を使うのが便利だったりします。
まず離れてるほうが互いの回転量の差が大きくなるので、そこで区別できるようになります。 また、1番目と3番目、2番目と4番目といった同じ2個離れた関係性も同じ用に扱うこともできます。 そして、この回転という計算はコサインを使うと簡単に表現できるし、長さが伸びても周期的に扱えたりもします。 そんな理由で、採用されています。
同じトークン内のベクトルの中でも回転させる量を変えている
RoPEでは回転を使うわけですが、同じベクトルの中でも回転させる量をベクトルの位置によって変えます。 こちらの図が、とあるトークン位置のベクトルの中でやっている処理です。
トークンごとのベクトル、例えばV_私をd次元のベクトルとして、そこから2つずつをペアにします。 そしてペアになった2つを2次元ベクトルとして、二次元座標上で回転させます。 図の(a, b)から(a′, b′)に回転させてるわけですが、sin/cosを使った行列計算できます。
この回転させる量がペアの場所によって徐々に変わっていきます。
なぜ回転させる量をわざわざ変えてるのでしょうか。 これは回転量を変えることで、回転の変化がゆっくりな成分では長距離同士の関係性を表現し、回転の変化が速い成分では短距離同士の関係を表現するという、短距離と長距離どちらにも強い状態を作っているからです。
回転量が少ない部分
回転量が少ない、すなわちトークン位置によって回転量がゆっくり変化する成分では、短距離同士では角度差が小さいので、ほぼ変化が分かりません。 つまり短距離同士の位置表現としてはあまり向いていません。 一方で、距離が離れた時には角度差が大きくなってくるので、長距離同士の位置表現には向いています。
回転量が多い部分
回転量が多い、すなわちトークン位置によって回転量が速く変化する成分では、短距離同士でも角度差が生まれるので短距離同士の表現には向いています。 しかしある程度離れると何周もぐるぐる回った状態になり、角度が折り返して区別がつかなくなり、長距離を表すのには向きません。
この両方の性質があることで、近い関係性も遠い関係性も把握できるようになります。
RoPEの全体像
最初の図と2つ目の図をまとめると、「トークンの場所によって回転量が変わる」し「ベクトルの場所によっても回転量が変わる」ということになります。 回転のさせ方は二番目の図に描いてあるものです。
両方をまとめると、こんな感じになります。
これがRoPEです！ヤドンのしっぽみたいなものでしょうか(?)
YaRNとは
さぁついにやぁ〜んの登場です。
「Yet another RoPE extensioN method」という名前の通りRoPEをExtension(拡張)する手法です。
YaRNを使うことで、元々のモデルで学習していたコンテキスト長の最大値より伸ばすことができます。 また、学習無しでも一定程度効果がありつつ、少量の学習をすることでより高性能かつ大きく伸ばすことが出来るというものです。
すごく誤解を招く分かりやすい言い方をすると「そのままヤドンのしっぽ(RoPE)を伸ばすとしっぽがちぎれるけど、やぁ〜んってするとしっぽをうまく伸びるテクニック」みたいな感じです(?)
※ちなみにヤドンのしっぽは切れても再び生えてきます。
どうやって伸ばしているか？
大きく2つのことをしています。
(1) RoPEにおける近距離同士の関係性を保つ部分はそのままにしておき、長距離同士の部分は伸ばす (NTK-by-parts)
(2) 伸ばしたことで関係性の強弱が変わる部分を補正 (Attention 温度パラメータ導入)
さらにコレに加えて推論時の文章の長さによって、伸ばし度合いを変える Dynamic Scalingという推論時のテクニックも紹介しているため、こちらも合わせて紹介します。
(1) NTK-by-partsについて
RoPEの説明のところで、同じトークンのベクトル内でもペア同士での回転角度を位置によって変えているという話がありました。そしてそれは、短距離と長距離の関係性を両方表現するためと述べました。
YaRN では、この中でも 短距離に強く作用する成分はそのままにしておき、長距離を扱う成分は伸ばし、中距離くらいは適度に伸ばす ということをします。
短距離同士の関係性をそのままにするのは、ほんの少し伸ばすだけでも敏感に影響を受けてしまうためです。例えば2倍に伸ばした時に、元の角度の差なら「2トークン離れている」と学習していたところが、実際には「4トークン離れている」となると、繊細なトークン間の距離の把握がうまくできなくなります。
一方で、遠く離れた位置同士をざっくり捉える成分は、多少伸ばしても影響が比較的少なく済みます。感覚的には、1万トークン離れてようが2万トークン離れてようが、とにかくめっちゃ遠いということに変わりはありません。 「やぁ〜〜.....〜〜〜ん」の “〜” が1万個だろうが2万個だろうが、とにかく長いと感じる部分はそこまで変わらないようなイメージです。
そのため、長距離を扱う成分は比較的伸ばしても、学習で得た性質を壊しにくいという考えです。
ちなみに、「YaRN でのコンテキスト長を伸ばす」というのはヤドンのしっぽ自体を1m から 100m まで成長させて伸ばすというよりも、メジャーに細工をしてヤドンに 100mまで伸びたと錯覚させるイメージのほうが近いです。90度の範囲で1万トークンの距離を測っていたところを、回転量を小さくして90度の範囲で10万トークンの距離まで測れるようにしているからです。
(2) Attention 温度パラメータ導入について
NTK-by-partsによる伸ばす処理をした際に、LLM(Transformer)の中での他の計算にも影響がでてきます。 その一つとして、Attention機構における、QueryとKeyの類似度からAttentionのweightを作る部分があります。
Attentionのweightというのは要は、今のトークンに対して関連性の高い注目すべき他トークンはどれなのか・どれくらいの関連度なのかを表現するものです。
NTK-by-partsで長距離同士のところを伸ばすと、全体的な回転量はより小さめになります。 例えば、10000トークン離れたところで今まで20度回転してたところを、40000トークン離れたところで20度回転してたことにするので、10000トークン付近を含めて全体的に回転量は減ります。 回転量が減ると、QueryとKeyの類似度を測る計算をする中でのCosine類似度もあがるため、Attention Weightは全体的に高くなってしまいます。
そうなると注目すべき他トークンの関係性として使っていた値が変わってきてしまうため、計算全体にも悪影響を及ぼしてしまいます。
そこで、全体的にcosine類似度が上がる部分を抑えるように、伸ばした分に応じて変化する温度パラメータtによって、類似度を減らすように、元の状態に近いように調整します。
このように伸ばした影響を抑えるのが、Attention 温度パラメータ導入というわけです。
これまたヤドンに例えると、ヤドンのしっぽが伸びると遠くまで攻撃が届くようになるのですが、 ヤドンは元々そんなに広い範囲をしっぽで狙う訓練をしていないので、 急に射程が広がるとうまく当てられません。 そこで、中心だけくっきり見えて周辺が少しぼやける特製メガネをプレゼントします。 このメガネのおかげで、
- 本当に狙うべき中心（重要トークン）に意識が向きやすくなる
- 周囲（不要トークン）の情報に惑わされにくくなる
ので、ヤドンは これまでと同じ感覚のまま遠くの敵にも正確に攻撃できるようになるようなイメージです。
推論時のテクニック: Dynamic Scaling
NTK-by-partsで長距離同士の関係性の部分を伸ばしていましたが、何でもかんでも無理やり伸ばす必要はありません。 例えば32,000を128,000に伸ばした場合、4倍に伸ばせるわけですが、64,000の長さの文章をわざわざ4倍まで伸ばす必要はありません。 2倍程度伸ばせたら十分なわけです。 YaRNは伸ばすこと自体はできるのですが、伸ばせば伸ばすほど精度も少しずつ悪化していきます。 なので、文章ごとに伸ばしたい分だけ伸ばすようにしたほうが、無駄な精度悪化を招かないというわけです。
ヤドンのしっぽも無駄に伸ばしすぎると、ヤドンが痛いだけですからね....
やぁ〜んをYaRNしてみる
さて実際にやってみないと分からないところはありますよね！ ということで実際に、やぁ〜んをYaRNを試してみましょう！(?)
流れとしては、
- やぁ〜んをベクトルで表現
- RoPEを適用する
- RoPE+YaRNを適用する
という流れの中で可視化しながら変化を見ていきます。
実際のLLMでは一つのトークンを数千次元や数万次元で表現するのですが、今回は可視化のわかりやすさ重視でかなり次元を落として試してみます。
1. やぁ〜んをベクトルで表現
まずは、「やぁ〜ん」をベクトルで表現します。 またこの「やぁ〜ん」を少し音声的に捉えて「や」と「ぁ」と「ん」の3つのベクトルを用意し、「〜」は「ぁ」と「ん」の間を補完する形とします。
ベクトルはそれぞれ10次元でなんとなくのお気持ちで表現してみます。 少し音的に考えて、母音の共通する「や」と「あ」を少し近づけています。あと分かりやすいように、0~1の範囲で表現しています。
v_や = [1.0, 0.7, 0.2, 0.1, 0.0, 0.0, 1.0, 0.8, 0.0, 1.0]
v_あ = [0.1, 0.0, 1.0, 0.7, 0.0, 0.0, 1.0, 0.9, 0.0, 1.0]
v_ん = [0.0, 0.0, 0.0, 0.0, 1.0, 0.7, 0.0, 0.0, 1.0, 1.0]
そして、補完した「〜」も加えて、このようなベクトルとしました。
2. RoPEの適用
RoPEをnumpyで簡易的に実装してみます。
数式を飛ばしたので、ここでパラメータが初めて登場しますがbaseというのがどれくらい回転させるかのパラメータです。
import numpy as np def rope(x, position, base=10000.0): """ x : ベクトル (shape: [dim]) position : 位置 (int) base : 角度のスケール (10000が一般的) """ x = np.asarray(x, dtype=np.float64) d = x.shape[0] assert d % 2 == 0 # 偶数 / 奇数成分 x_even = x[0::2] x_odd = x[1::2] # 角度 idx = np.arange(d // 2, dtype=np.float64) theta = position * (1.0 / (base ** (2.0 * idx / d))) # = m * θ_d cos = np.cos(theta, dtype=np.float64) sin = np.sin(theta, dtype=np.float64) # ペアの回転 out = np.empty_like(x) out[0::2] = x_even * cos - x_odd * sin out[1::2] = x_even * sin + x_odd * cos return out
ここに先程のベクトルを適用させていきます。
# ベクトルの定義 v_yaan = np.array([ [1.00, 0.70, 0.20, 0.10, 0.00, 0.00, 1.00, 0.80, 0.00, 1.00], # や [0.10, 0.00, 1.00, 0.70, 0.00, 0.00, 1.00, 0.90, 0.00, 1.00], # ぁ (= あ) [0.05, 0.00, 0.50, 0.35, 0.50, 0.35, 0.50, 0.45, 0.50, 1.00], # 〜 (あ+ん)/2 [0.00, 0.00, 0.00, 0.00, 1.00, 0.70, 0.00, 0.00, 1.00, 1.00] # ん ]) # RoPE適用 base = 1000 v_yaan_rotated = [] for token_ix in range(v_yaan.shape[0]): token_vector = v_yaan[token_ix] # token_ix番目のトークンベクトル roped_vector = rope(token_vector, token_ix, base=base) # RoPE適用 v_yaan_rotated.append(roped_vector) v_yaan_rotated = np.stack(v_yaan_rotated) # shape: (N, D)
RoPEのbaseパラメータを変えながらどれくらいベクトルが変化するかを見てみます。
一番左のグラフがRoPE適用前のベクトル、中央のグラフがRoPE適用後のベクトル、右側のグラフが両者の差分です。
横方向の0, 1, 2, 3がそれぞれ「や」「ぁ」「〜」「ん」にあたります。 縦方向がトークンをベクトル化した10次元で上から順番に先程表で定義した数値が入っており、それを色で表現しています。
base = 10000
base = 1000
base = 100
考察
baseを小さくするほど、差分が見られるのが分かります。 base=10000とかだと、左2つのベクトル自体の可視化では差がほとんどわからないですね。 （ちなみにサンプルコードや論文での例だと10000を使うことが多い印象です。）
RoPEの説明で述べたように回転量がベクトル内の場所によって違います。
まず、同じ文字のベクトル内でも最初のペアは回転量が大きく、トークンベクトル内での位置(図で言うDimension index)が大きいほど、回転量が小さくなります。 また、何番目のトークンかによっても回転量が変わり、最初の文字は回転せず、後半の文字につれて徐々に回転する量が増えます。
1トークン目は回転量が0なのでRoPEによって全く変化しておらず、それがグラフからも読み取れます。
右上ほど回転量が大きいので、よりRoPEによる変化が大きい気もしますが、実際にはそうなっていません。 これは回転量だけでなく、元のベクトルの値の大きさにもよって変化量は変わってくるので、その影響かと思います。
ヤドンを回転させるより、大きなヤドランを回転させるほうが変化量も大きいわけです(?)
3. RoPE+YaRNを適用する
次にYaRNを適用しましょう。
YaRNの実装はこちらです。
import numpy as np import math def yarn( x, position, orig_context, # 元の学習時コンテキスト長 L scale, # スケール係数 s = L_new / L base=10000.0, # RoPE のベース b alpha=1.0, # YaRN 推奨: α=1 beta=32.0, # YaRN 推奨: β=32 apply_attn_scale=False, # True にすると 1/t による attention スケーリングを適用 ): """ YaRN (NTK-by-parts + optional attention scaling) を 1 本のベクトルに適用する。 Parameters ---------- x : np.ndarray shape: [dim] の 1D ベクトル（dim は偶数） position : int 位置 m orig_context : int or float 元モデルが学習された最大コンテキスト長 L (例: 4096) scale : float コンテキスト拡張比 s = L_new / L (例: 8.0) base : float RoPE のベース b （通常 10000） alpha, beta : float YaRN の境界パラメータ (推奨 α=1, β=32) apply_attn_scale : bool True のとき、1/t = 0.1 * ln(s) + 1 に基づく attention スケーリングも入れる （論文の式 (22) 部分） Returns ------- x_rot : np.ndarray YaRN を適用したベクトル（shape: [dim]） """ x = np.asarray(x, dtype=np.float64) d = x.shape[0] assert d % 2 == 0 # ----------------------------------------------------- # 1) RoPE 周波数 θ_d を 計算 # ----------------------------------------------------- idx = np.arange(d // 2, dtype=np.float64) theta_base = 1.0 / (base ** (2.0 * idx / d)) # θ_d # ----------------------------------------------------- # 2) r(d) と γ(r) # ----------------------------------------------------- L = float(orig_context) r = L * theta_base / (2.0 * math.pi) # float64 gamma = np.zeros_like(r, dtype=np.float64) gamma[r > beta] = 1.0 mid = (r >= alpha) & (r <= beta) gamma[mid] = (r[mid] - alpha) / (beta - alpha) # ----------------------------------------------------- # 3) θ'_d = ((1-γ)θ/s + γθ) # ----------------------------------------------------- s = float(scale) theta_yarn = (1.0 - gamma) * (theta_base / s) + gamma * theta_base # ----------------------------------------------------- # 4) Attention scaling (optional, no side effect if false) # ----------------------------------------------------- x_even = x[0::2].copy() x_odd = x[1::2].copy() if apply_attn_scale: inv_t = 0.1 * math.log(s) + 1.0 # 1/t length_scale = math.sqrt(inv_t) x_even *= length_scale x_odd *= length_scale # ----------------------------------------------------- # 5) 回転 # ----------------------------------------------------- angles = position * theta_yarn cos = np.cos(angles, dtype=np.float64) sin = np.sin(angles, dtype=np.float64) out = np.empty_like(x) out[0::2] = x_even * cos - x_odd * sin out[1::2] = x_even * sin + x_odd * cos return out
RoPEをしてからYaRNをするというよりは、RoPEの回転角度の計算にYaRNを使うため、同時に合わせて処理します。
こちらも同様に可視化してみます。
今度はRoPEとも比較したいので、左から順番に
- 元のトークンベクトル
- RoPE適用
- YaRN適用
- 元のとYaRNの差分
- RoPEとYaRNの差分
として可視化していきます。
baseは100、伸ばす長さは8倍と固定して、「元のモデルのコンテキスト長」を変えてみます。 異なる長さのヤドンの尻尾を伸ばす（伸びたように見せる）イメージです。
「Attention 温度パラメータ導入」については、推論時は計算上は入れるべきですが、ベクトルの違いを見るだけなので今回はスキップします。
コンテキスト長: 4096
初期の頃のLLMのコンテキスト長が4096くらいだったため採用してみました。 RoPEとYaRNを比較すると、このスケールではほとんど差がない事がわかります。（右下だけ若干差が見える）
コンテキスト長: 16
4096と比較すると、RoPEとYaRNで差が出てきているのが分かります。
考察
YaRNの説明の中で、短距離に強く作用する成分はそのままにしておき、長距離を扱う成分は伸ばすという話をしました。
コンテキスト長を変えると、そのモデルにとっての長距離とはどれくらいか？が変わります。 例えば8トークンの長さのものがあったとき、4096にとって8はかなり短距離なのに対して、16にとっては半分になるのでそこそこの長さという扱いになります。 体の長いハクリューにとっては、ヤドンの尻尾はそんな長くないけど、尻尾の短いイーブイにとってはヤドンの尻尾は長いとみえますよね（たぶん）。
つまり、長距離の基準がここによって変わります。 長距離を扱う成分を伸ばすYaRNですが、4096のほうが長距離扱いする対象が厳しく、16のほうが長距離扱いの基準がゆるくて伸ばす方向になる条件が増えます。 そして、今回の例でもコンテキスト長16のほうが反応したというふうに解釈できそうです。
最後に
というわけで、ロングコンテキストなLLMを作るのに必須ともいえるYaRNについて、やぁ〜んと説明してみました。
ヤドンでも分かるような説明をしようと、概念的に説明してみましたが、思ったより複雑になってしまいました。
また数式を避けて概念的な説明ゆえに正確性に欠けている部分があると思いますので、より正確に理解するためには論文を読むことを推奨します。
ところで、YaRNはやぁ〜んと読むのでしょうか
We Are Hiring!
ABEJAは、テクノロジーの社会実装に取り組んでいます。 技術はもちろん、技術をどのようにして社会やビジネスに組み込んでいくかを考えるのが好きな方は、下記採用ページからエントリーください！ （新卒の方やインターンシップのエントリーもお待ちしております！） careers.abejainc.com
特に下記ポジションの募集を強化しています！ぜひ御覧ください！
