メインコンテンツまでスキップ

MRoPEとは?

概要 (画像は、Geminiで作成されたものです)

MRoPEの概要

MRoPE(Multi-dimensional Rotary Position Embedding) とは、 トークンの特徴量ベクトルを「回転」させることで、時間や空間といった多次元の相対位置情報を表現する技術です。 主にQwen2-VLなどのマルチモーダル大規模モデル(LMM)で採用されており、以下の特徴を持ちます。

  • 多次元対応: 画像の縦・横(2D)や、動画の時間(3D)の位置を同時に扱える
  • 高い外挿性: 学習時より高解像度な画像や、長い動画が来ても破綻せずに理解できる
  • 相対位置の学習: トークン同士の「距離感」をモデルが自然に捉えられる

動画理解における位置表現の進化:2D RoPEから「MRoPE」へ

比較 (画像は、Geminiで作成されたものです)

2次元の画像認識において、空間を縦(yy軸)と横(xx軸)に分解して処理する「2D RoPE」は、解像度やアスペクト比の変化に対して極めて高い柔軟性を発揮しました。 しかし、このアプローチをさらに「動画(Vision-Languageタスク)」へと拡張しようとした際、もう一つの決定的な壁にぶつかることになります。 それが「時間軸(tt軸)」という3つ目の次元の存在です。

従来の動画理解モデル(Video Transformerなど)では、動画データを処理する際、大きく分けて2つのアプローチが取られてきました。

1つ目は、「1Dへの平坦化(Flattening)」です。 これは、動画の各フレームをパッチに分割し、さらにそれらを時間順に並べた上で、全てのパッチを強引に1列の長いシーケンスへと平坦化する方法です。 ここに1次元の絶対位置埋め込みを足し算します。 しかし、この方法では「同じフレーム内の隣のパッチ(空間的な近さ)」と、「次のフレームの同じ位置(時間的な近さ)」の区別がモデルにとって等価になってしまいます。 結果として、動画が持つ本質的な「3次元的な時空間構造」を効率的に学習できないという致命的なデメリットがありました。

2つ目は、「空間・時間分離型(Factorized / Separable)位置埋め込み」です。 2D RoPEの発想に近く、空間(2D)の位置ベクトルと、時間(1D)の位置ベクトルを別々に用意し、後から「足し算」または「結合」して3次元的な位置を表現します。 一見すると理にかなっていますが、このアプローチの本質は「固定された絶対位置の足し算」です。 あらかじめ「16フレーム、224x224ピクセル」のように最大サイズを固定して位置ベクトルを学習するため、推論時にそれを超える「長い動画」や「高解像度な映像」が入力されると、 対応する未知の座標ラベルが存在しないため、モデルの位置感覚が完全に崩壊してしまう(外挿性の低さ)という弱点がありました。

この「時空間の構造維持」と「サイズ固定の壁」という2つの課題を同時に解決したのが、Qwen2-VLなどに採用されている「MRoPE(Multi-dimensional Rotary Position Embedding:多次元回転位置埋め込み)」です。

MRoPEは、2D RoPEの「隠れ層を分割して独立した回転を施す」というエレガントな特性を、さらに「時間軸(tt軸)」、「縦方向(yy軸:Height)」、「横方向(xx軸:Width)」の3つの独立した回転へと拡張しました。 具体的には、モデルの隠れ層のベクトル(特徴量)を綺麗に3つのブロックに分割します。

  • 1つ目のブロックには、そのトークンが「何フレーム目(何秒目)か」に応じた時間軸の回転を施します。
  • 2つ目のブロックには、フレーム内の縦方向の位置に応じた回転を施します。
  • 3つ目のブロックには、フレーム内の横方向の位置に応じた回転を施します。

これらを1つのベクトルとして結合(Concat)することで、QueryとKeyの内積(アテンション・スコア)を計算する際、数学的に「時間的な距離」「縦方向の距離」「横方向の距離」のすべてがアテンションの重みに直接、かつ独立して反映されるようになります。

MRoPEの導入によって、モデルは「特定の時間(いつ)、特定の場所(どこ)」で起きた事象かを完璧にトラッキングできるようになりました。 さらに、位置情報が無限に遠くへ飛んでいかない「回転角の差分(周波数)」として表現されているため、学習時よりはるかに高解像度なカメラ映像や、数分に及ぶ長い動画が入力されても、空間や時間の歪みを滑らかに補正しながら正確にコンテキストを理解することが可能になったのです。

この「絶対位置の足し算」から「多次元の回転(MRoPE)」へのパラダイムシフトこそが、Qwen2-VLが任意の解像度や動画の長さをそのままネイティブに高精度処理できる(Naive Dynamic Resolution)技術的コアとなっています。

1D RoPE, 2D RoPE, MRoPEの比較

以下では数式を用いて1D RoPE, 2D RoPE, MRoPEの比較を行います。
まずは行列による回転と複素数表現のつながりについて解説し、その後、それぞれの次元のRoPEについて解説していきます。

行列による回転と複素数表現のつながり

RoPEにおける位置エンコーディングの基本は、sin\sincos\cos を用いた2次元の回転行列を対角線上に並べたブロック対角行列をベクトルに掛けることです 。 ある特徴量ベクトルの隣り合う2つの要素 (x1,x2)(x_1, x_2) に対して、位置 mm と周波数 θ\theta に基づく回転行列を掛けると、以下のようになります。

(cosmθsinmθsinmθcosmθ)(x1x2)=(x1cosmθx2sinmθx1sinmθ+x2cosmθ)\begin{pmatrix} \cos m\theta & -\sin m\theta \\ \sin m\theta & \cos m\theta \end{pmatrix} \begin{pmatrix} x_1 \\ x_2 \end{pmatrix} = \begin{pmatrix} x_1 \cos m\theta - x_2 \sin m\theta \\ x_1 \sin m\theta + x_2 \cos m\theta \end{pmatrix}

ここで、この2要素を複素数 c=x1+jx2c = x_1 + j x_2jj は虚数単位)と見立ててみます。 オイラーの公式(ejθ=cosθ+jsinθe^{j\theta} = \cos\theta + j\sin\theta)を用いると、先ほどの行列による回転計算は、複素数の掛け算として次のように全く同じ結果として美しく表現できます。

cejmθ=(x1+jx2)(cosmθ+jsinmθ)=(x1cosmθx2sinmθ)+j(x1sinmθ+x2cosmθ)\begin{aligned} c \cdot e^{j m\theta} &= (x_1 + j x_2)(\cos m\theta + j\sin m\theta) \\ &= (x_1 \cos m\theta - x_2 \sin m\theta) + j(x_1 \sin m\theta + x_2 \cos m\theta) \end{aligned}

つまり、「ブロック対角行列を掛けること」と「複素数にして ejθe^{j\theta} を掛けること」は数学的に等価です。次元数が増えてベクトルを分割する処理を数式で書く際、巨大な行列を書くよりも複素数表記を用いた方が非常にスッキリするため、以下ではこの複素数表現 ckc_k を用いて各次元のRoPEを比較します。

1. 1次元 RoPE (1D RoPE)

  • 主な用途: テキスト(自然言語処理)などのシーケンスデータ
  • 位置情報: 1つの座標 mm (シーケンス内のトークン位置)

1次元RoPEでは、隠れ層の次元 dd 全体に対して、1つの位置座標 mm を用いて回転行列を適用します 。 周波数パラメータを θk\theta_k とすると、変換後の複素ベクトル ckc'_k は以下のようになります。

ck=ckejmθk(k=0,1,,d21)c'_k = c_k \cdot e^{j m \theta_k} \quad \left(k = 0, 1, \dots, \frac{d}{2}-1\right)

意味: 特徴量ベクトルのすべての次元ペアが、シーケンス上の位置 mm に応じた角度で一斉に回転します。

2. 2次元 RoPE (2D RoPE)

  • 主な用途: 画像(Vision Transformer等)やグリッドデータ
  • 位置情報: 2つの座標 (x,y)(x, y) (横と縦)

2次元RoPEでは、画像のように縦横の空間情報を持つため、隠れ層の次元 dd を前半の d2\frac{d}{2} 次元(縦方向 yy)と後半の d2\frac{d}{2} 次元(横方向 xx)に半分ずつ分割します 。そして、それぞれの軸に対して独立した回転を適用します 。

ck={ckejyθk(k=0,1,,d41) 前半: 縦(y)の適用ckejxθkd/4(k=d4,,d21) 後半: 横(x)の適用c'_k = \begin{cases} c_k \cdot e^{j y \theta_k} & \left(k = 0, 1, \dots, \frac{d}{4}-1\right) \quad \text{… 前半: 縦(y)の適用} \\c_k \cdot e^{j x \theta_{k - d/4}} & \left(k = \frac{d}{4}, \dots, \frac{d}{2}-1\right) \quad \text{… 後半: 横(x)の適用}\end{cases}

意味: 前半成分は縦の位置 yy に応じて回転し、後半成分は横の位置 xx に応じて回転します。これにより、2D空間の幾何学的な位置関係を別々に捉えることができます。

3. 3次元 RoPE (3D RoPE)

  • 主な用途: 動画(時間・縦・横)や3D空間データ(深さ・縦・横)
  • 位置情報: 3つの座標 (t,y,x)(t, y, x) (時間・縦・横など)

動画など3つの座標軸を持つデータを扱う場合、次元 dd を3等分(1/3ずつ)に分割します。それぞれに「時間(t)」「縦(y)」「横(x)」の位置情報を割り当てて回転させます。

ck={ckejtθk(k=0,,d61) 1/3部分: 時間(t)の適用ckejyθkd/6(k=d6,,2d61) 2/3部分: 縦(y)の適用ckejxθk2d/6(k=2d6,,d21) 3/3部分: 横(x)の適用c'_k = \begin{cases} c_k \cdot e^{j t \theta_k} & \left(k = 0, \dots, \frac{d}{6}-1\right) \quad \text{… 1/3部分: 時間(t)の適用} \\c_k \cdot e^{j y \theta_{k - d/6}} & \left(k = \frac{d}{6}, \dots, \frac{2d}{6}-1\right) \quad \text{… 2/3部分: 縦(y)の適用} \\c_k \cdot e^{j x \theta_{k - 2d/6}} & \left(k = \frac{2d}{6}, \dots, \frac{d}{2}-1\right) \quad \text{… 3/3部分: 横(x)の適用}\end{cases}

意味: ベクトルを3つの空間に分け、それぞれ「フレーム間の距離」「縦のピクセル距離」「横のピクセル距離」に対応する独立した回転を適用します。

コードによるMRoPEの実装

コードからQuery V0H0T0に対するAttentionを可視化することで、上記の数式で説明した1D, 2D, MRoPEの挙動の違いを直感的に理解します。

d=144d=144の場合(時間48次元(24ペア)、縦48次元(24ペア)、横48次元(24ペア)の場合)

以下のコードでは、隠れ層の次元数 d が144の場合のMRoPEの挙動をシミュレートします。この理想的な設定では、時間(T)、縦(V)、横(H)の各軸に48次元(24ペア)ずつ均等に次元を割り当てます。numpy を用いて1D、2D、3D(MRoPE)の各方式でアテンションスコアを計算し、その結果を matplotlib で可視化することで、3次元の時空間情報を各モデルがどのように捉えるかの違いを比較します。

import numpy as np
import matplotlib.pyplot as plt

# --- 1. 基本設定 ---
# 時間(T), 縦(V), 横(H) のグリッドサイズ
T_size, V_size, H_size = 4, 4, 4

# 隠れ層の次元数 d と実際のRoPE周波数の計算
base = 10000
d = 144 # ★ 128から144に変更
k = np.arange(d // 2) # ペアのインデックス (0 ~ 71 の計72ペア)

# 実際の逆周波数 (Inverse Frequency) の数式
thetas = 1.0 / (base ** (2 * k / d))

# アテンションスコアを格納する3Dテンソル
attn_1d = np.zeros((T_size, V_size, H_size))
attn_2d = np.zeros((T_size, V_size, H_size))
attn_3d = np.zeros((T_size, V_size, H_size))

# 内積を計算するヘルパー関数
def calc_score(delta_pos, theta_list):
return sum(2 * np.cos(delta_pos * th) for th in theta_list)

# --- 2. Queryの位置設定 (V0, H0, T0) ---
q_T, q_V, q_H = 0, 0, 0
# 1Dシーケンスに平坦化した場合のインデックス
q_m = q_T * (V_size * H_size) + q_V * H_size + q_H

# --- 3. 全Key位置に対するAttention Scoreの計算 ---
# 周波数の均等配分(インターリーブ)
# 2D RoPE用: 72ペアを2分割 (36ペアずつ綺麗に分割)
thetas_2d_v = thetas[0::2]
thetas_2d_h = thetas[1::2]

# 3D RoPE用: 72ペアを3分割 (24ペアずつ綺麗に分割)
thetas_3d_t = thetas[0::3]
thetas_3d_v = thetas[1::3]
thetas_3d_h = thetas[2::3]

for t in range(T_size):
for v in range(V_size):
for h in range(H_size):
# 1D RoPE用: 平坦化された1次元の相対距離
m = t * (V_size * H_size) + v * H_size + h
delta_m = m - q_m

# 各軸の相対距離
delta_t = t - q_T
delta_v = v - q_V
delta_h = h - q_H

# ① 1D RoPE (全72ペアを使用)
attn_1d[t, v, h] = calc_score(delta_m, thetas)

# ② 2D RoPE (Tは無視してVとHのみ)
score_2d_v = calc_score(delta_v, thetas_2d_v)
score_2d_h = calc_score(delta_h, thetas_2d_h)
attn_2d[t, v, h] = score_2d_v + score_2d_h

# ③ 3D RoPE (T, V, Hのそれぞれに独立して適用)
score_3d_t = calc_score(delta_t, thetas_3d_t)
score_3d_v = calc_score(delta_v, thetas_3d_v)
score_3d_h = calc_score(delta_h, thetas_3d_h)
attn_3d[t, v, h] = score_3d_t + score_3d_v + score_3d_h

# --- 4. 可視化 (断面図のプロット) ---
fig, axes = plt.subplots(3, 3, figsize=(16, 16))
fig.suptitle("Attention Scores with Real RoPE Frequencies (d=144)", fontsize=20, y=0.95)

configs = [
{"title": "V-H Slice (T=0)", "get_data": lambda a: a[0, :, :], "xlabel": "Horizontal (H)", "ylabel": "Vertical (V)"},
{"title": "H-T Slice (V=0)", "get_data": lambda a: a[:, 0, :].T, "xlabel": "Time (T)", "ylabel": "Horizontal (H)"},
{"title": "V-T Slice (H=0)", "get_data": lambda a: a[:, :, 0].T, "xlabel": "Time (T)", "ylabel": "Vertical (V)"}
]

models = [
{"name": "1D RoPE", "data": attn_1d},
{"name": "2D RoPE", "data": attn_2d},
{"name": "3D RoPE", "data": attn_3d}
]

for row, model in enumerate(models):
for col, config in enumerate(configs):
ax = axes[row, col]
slice_data = config["get_data"](model["data"])

# ★ d=144のため、スコアの最大値(vmax)を144に設定
im = ax.imshow(slice_data, cmap='viridis', vmin=-10, vmax=144)

ax.set_title(f'{model["name"]} : {config["title"]}', fontsize=14, pad=10)
ax.set_xlabel(config["xlabel"])
ax.set_ylabel(config["ylabel"])

ax.set_xticks(range(4))
ax.set_yticks(range(4))

if col == 0:
ax.set_xticklabels([f"H{i}" for i in range(4)])
ax.set_yticklabels([f"V{i}" for i in range(4)])
elif col == 1:
ax.set_xticklabels([f"T{i}" for i in range(4)])
ax.set_yticklabels([f"H{i}" for i in range(4)])
elif col == 2:
ax.set_xticklabels([f"T{i}" for i in range(4)])
ax.set_yticklabels([f"V{i}" for i in range(4)])

# スコアのテキスト描画 (見やすく整数表記に)
for i in range(4):
for j in range(4):
val = slice_data[i, j]
# ★ 背景色に合わせて文字色を変える閾値を90に調整
color = "black" if val > 90 else "white"
ax.text(j, i, f"{val:.0f}", ha="center", va="center", color=color, fontweight='bold')

plt.tight_layout(rect=[0, 0, 1, 0.93])
plt.show()

実行結果 実行結果

このコードは、1D RoPE2D RoPE、そして 3D RoPE (MRoPE) が、4x4x4の時空間グリッド内で、原点 (T0, V0, H0) に置かれたQueryに対し、他のすべてのKeyトークンとの間でどのようなアテンションスコアを計算するかをシミュレーションしたものです。隠れ層の次元 d は144で、各軸に均等に48次元(24ペア)が割り当てられています。

  • 1D RoPE: 動画の3Dグリッドを1次元のシーケンスに平坦化(m = t * (V_size * H_size) + v * H_size + h)して相対距離 delta_m を計算しています。可視化結果を見ると、原点から1次元的に離れるにつれてスコアが減衰しているのがわかります。この方法では、例えば (T=0, V=0, H=1)(T=1, V=0, H=0) のような、本来は異なる時空間的な関係にあるトークンを区別できず、空間(縦横)の構造が壊れてしまいます。

  • 2D RoPE: 時間軸 t を完全に無視し、縦 delta_v と横 delta_h の相対距離のみでスコアを計算します。そのため、「V-H Slice (T=0)」のグラフでは、原点を中心に同心円状にスコアが分布しており、空間的な距離を完璧に捉えられています。しかし、「H-T Slice」や「V-T Slice」を見ると、時間 t が変化してもスコアが全く変わっていません。これは、2D RoPEが時間(フレーム)の経過を全く測定できていないことを意味します。

  • 3D RoPE (MRoPE): 隠れ層を3分割し、時間 delta_t、縦 delta_v、横 delta_h のそれぞれに独立した周波数セット(thetas_3d_t, thetas_3d_v, thetas_3d_h)を割り当ててスコアを計算し、それらを合算しています。結果、すべての断面図で原点を中心としたきれいなスコア分布が確認できます。これは、時間・縦・横にそれぞれ独立した「定規」を適用することで、3次元の時空間構造を崩すことなく、すべての軸の相対的な距離を同時に、かつ正確に測定できていることを示しています。

d=128d=128の場合(時間32次元(16ペア)、縦48次元(24ペア)、横48次元(24ペア)の場合)

続くコードは、より実用的なシナリオとして、隠れ層の次元数 d が128の場合を扱います。実際のAIモデル、特にQwen2-VLなどでは、GPU(特にNVIDIAのTensor Core)の計算効率を最大化するため、次元数に64や128といった「2のべき乗」が採用されるのが一般的です。この制約下で3次元情報を扱うため、次元を不均等に割り当てる工夫がなされます。ここでは、時間軸に16ペア(32次元)、縦横の空間軸にそれぞれ24ペア(48次元)を割り当てた場合のアテンションスコアを numpy で計算し、matplotlib を用いてその影響を可視化します。

import numpy as np
import matplotlib.pyplot as plt

# --- 1. 基本設定 ---
# 時間(T), 縦(V), 横(H) のグリッドサイズ
T_size, V_size, H_size = 4, 4, 4

# 隠れ層の次元数 d と実際のRoPE周波数の計算
base = 10000
d = 128
k = np.arange(d // 2) # ペアのインデックス (0 ~ 63 の計64ペア)

# 実際の逆周波数 (Inverse Frequency) の数式
thetas = 1.0 / (base ** (2 * k / d))

# アテンションスコアを格納する3Dテンソル
attn_1d = np.zeros((T_size, V_size, H_size))
attn_2d = np.zeros((T_size, V_size, H_size))
attn_3d = np.zeros((T_size, V_size, H_size))

# 内積を計算するヘルパー関数
def calc_score(delta_pos, theta_list):
return sum(2 * np.cos(delta_pos * th) for th in theta_list)

# --- 2. Queryの位置設定 (V0, H0, T0) ---
q_T, q_V, q_H = 0, 0, 0
q_m = q_T * (V_size * H_size) + q_V * H_size + q_H

# --- 3. 全Key位置に対するAttention Scoreの計算 ---

# 2D RoPE用: 64ペアを2等分 (32ペアずつ)
thetas_2d_v = thetas[0::2]
thetas_2d_h = thetas[1::2]

# ★ 3D RoPE用: 64ペアを 不均等 (16 : 24 : 24) にインターリーブ配分
all_idx = np.arange(64)

# 時間(T)には全体の1/4(4つに1つ)を割り当て = 16ペア
t_idx = all_idx[::4]

# 残りの48ペアのインデックスを取得
rem_idx = np.setdiff1d(all_idx, t_idx)

# 残りを縦(V)と横(H)で交互に半分ずつ割り当て = 24ペアずつ
v_idx = rem_idx[0::2]
h_idx = rem_idx[1::2]

thetas_3d_t = thetas[t_idx]
thetas_3d_v = thetas[v_idx]
thetas_3d_h = thetas[h_idx]

for t in range(T_size):
for v in range(V_size):
for h in range(H_size):
# 1Dの相対距離
m = t * (V_size * H_size) + v * H_size + h
delta_m = m - q_m

# 各軸の相対距離
delta_t = t - q_T
delta_v = v - q_V
delta_h = h - q_H

# ① 1D RoPE
attn_1d[t, v, h] = calc_score(delta_m, thetas)

# ② 2D RoPE
attn_2d[t, v, h] = calc_score(delta_v, thetas_2d_v) + calc_score(delta_h, thetas_2d_h)

# ③ 3D RoPE (16, 24, 24のペアで計算)
attn_3d[t, v, h] = calc_score(delta_t, thetas_3d_t) + \
calc_score(delta_v, thetas_3d_v) + \
calc_score(delta_h, thetas_3d_h)

# --- 4. 可視化 (断面図のプロット) ---
fig, axes = plt.subplots(3, 3, figsize=(16, 16))
fig.suptitle("Attention Scores (d=128) | 3D RoPE uses T:16, V:24, H:24", fontsize=20, y=0.95)

configs = [
{"title": "V-H Slice (T=0)", "get_data": lambda a: a[0, :, :], "xlabel": "Horizontal (H)", "ylabel": "Vertical (V)"},
{"title": "H-T Slice (V=0)", "get_data": lambda a: a[:, 0, :].T, "xlabel": "Time (T)", "ylabel": "Horizontal (H)"},
{"title": "V-T Slice (H=0)", "get_data": lambda a: a[:, :, 0].T, "xlabel": "Time (T)", "ylabel": "Vertical (V)"}
]

models = [
{"name": "1D RoPE", "data": attn_1d},
{"name": "2D RoPE", "data": attn_2d},
{"name": "3D RoPE", "data": attn_3d}
]

for row, model in enumerate(models):
for col, config in enumerate(configs):
ax = axes[row, col]
slice_data = config["get_data"](model["data"])

# d=128のためスコアの最大値は128
im = ax.imshow(slice_data, cmap='viridis', vmin=-10, vmax=128)

ax.set_title(f'{model["name"]} : {config["title"]}', fontsize=14, pad=10)
ax.set_xlabel(config["xlabel"])
ax.set_ylabel(config["ylabel"])

ax.set_xticks(range(4))
ax.set_yticks(range(4))

if col == 0:
ax.set_xticklabels([f"H{i}" for i in range(4)])
ax.set_yticklabels([f"V{i}" for i in range(4)])
elif col == 1:
ax.set_xticklabels([f"T{i}" for i in range(4)])
ax.set_yticklabels([f"H{i}" for i in range(4)])
elif col == 2:
ax.set_xticklabels([f"T{i}" for i in range(4)])
ax.set_yticklabels([f"V{i}" for i in range(4)])

# スコアのテキスト描画
for i in range(4):
for j in range(4):
val = slice_data[i, j]
color = "black" if val > 80 else "white"
ax.text(j, i, f"{val:.0f}", ha="center", va="center", color=color, fontweight='bold')

plt.tight_layout(rect=[0, 0, 1, 0.93])
plt.show()

実行結果 実行結果 このシミュレーションでは、隠れ層の次元 d がGPUフレンドリーな128に設定されています。合計64ペアの次元を、時間(T)に16ペア、縦(V)に24ペア、横(H)に24ペアという不均等な配分で割り当てています。これは、動画においては空間的な解像度(V, H)の方が時間的な解像度(T)よりも重要度が高いという判断に基づく、実用的な設定です。

  • 1D RoPE2D RoPE の挙動は、d=144のケースと本質的に同じです。1Dは空間構造を無視し、2Dは時間構造を無視しています。スコアの最大値が128になっている点のみが異なります。

  • 3D RoPE (MRoPE): ここでは、周波数ペアをインターリーブ方式で不均等に分割しています。まず全体の1/4にあたるインデックス(all_idx[::4])を時間 t_idx に割り当て、残りを縦 v_idx と横 h_idx で交互に分け合っています。 可視化結果を見ると、d=144の均等配分の場合と同様に、すべての断面で原点を中心とした距離減衰が正しく表現できています。これは、MRoPEが各軸に割り当てる次元数が異なっていても、それぞれの軸の「定規」として独立して機能し、時空間構造を正確に捉えられる柔軟なアーキテクチャであることを示しています。次元配分を調整することで、タスクに応じて特定の軸の解像度を優先するチューニングが可能になります。

動画データへのMRoPEの適用

これまでの理論とシミュレーションを踏まえ、実際の動画データに対してMRoPE(3D RoPE)を適用する具体的なプロセスをコードで示します。torch を用いて、ダミーの動画テンソル(時間、チャネル、縦、横の4次元)を作成するところから始めます。次に、この動画をパッチに分割し、線形層で特徴量ベクトルに変換後、各パッチの時空間座標(t, y, x)に応じてMRoPEによる回転を適用し、最終的にTransformerに入力するための1次元シーケンスを生成するまでの一連の流れを解説します。

import torch
import numpy as np

# --- 1. 入力動画の設定 (5フレーム x 10x10画像 x 3チャネル) ---
# Batch=1, Time=5, Channel=3, Height=10, Width=10
torch.manual_seed(42)
video = torch.randn(1, 5, 3, 10, 10)

# 空間方向のパッチサイズ (2x2ピクセル)
patch_size = 2
# 隠れ層の次元数 d=6 (時間:2, 高さ:2, 幅:2 に3分割するため6に設定)
hidden_dim = 6

# --- 2. 動画をパッチに分割してベクトル化 (Linear Projection) ---
# 5フレームの各10x10画像から、縦5個 × 横5個 = 25パッチ 作成されます
# 全体で 5フレーム × 25パッチ = 計125パッチ になります。
# 各パッチの元の画素数は 2 * 2 * 3(RGB) = 12次元 です

# ① まず縦方向(Dim 3)をサイズ2、ストライド2で切り出す
patches = video.unfold(3, patch_size, patch_size)
# 形状: [1, 5, 3, 5, 10, 2]

# ② 次に横方向(Dim 4)をサイズ2、ストライド2で切り出す
patches = patches.unfold(4, patch_size, patch_size)
# 形状: [1, 5, 3, 5, 5, 2, 2]

# 扱いやすいように次元の順序を入れ替える
# [Batch, Time, Channel, Grid_H, Grid_W, Patch_H, Patch_W]
# -> [Batch, Time, Grid_H, Grid_W, Channel, Patch_H, Patch_W]
patches = patches.permute(0, 1, 3, 4, 2, 5, 6)

# パッチ内の画素情報を1つのベクトルに平坦化
# 形状を [Batch(1), Time(5), Grid_H(5), Grid_W(5), パッチ内画素数(12)] に整理
patches = patches.flatten(4)

# 線形投影(Linear Projection)により、画素数12からモデルの隠れ層次元数6へ変換
projection_layer = torch.nn.Linear(patch_size * patch_size * 3, hidden_dim)
patch_features = projection_layer(patches)
# 形状: [1, 5, 5, 5, 6]

# --- 3. 3D RoPE の回転行列の生成 ---
# 時間(5) × 縦(5) × 横(5) のグリッド座標に対する回転行列を構築します
grid_t, grid_h, grid_w = patch_features.shape[1], patch_features.shape[2], patch_features.shape[3]
theta = 0.5

# 各座標用の回転行列を格納するテンソル [5, 5, 5, 6, 6]
R_3d_grid = torch.zeros(grid_t, grid_h, grid_w, hidden_dim, hidden_dim)

for t in range(grid_t):
for y in range(grid_h):
for x in range(grid_w):
# 1/3部分 (インデックス 0, 1): 時間方向 t に応じた回転
R_3d_grid[t, y, x, 0, 0] = np.cos(t * theta)
R_3d_grid[t, y, x, 0, 1] = -np.sin(t * theta)
R_3d_grid[t, y, x, 1, 0] = np.sin(t * theta)
R_3d_grid[t, y, x, 1, 1] = np.cos(t * theta)

# 2/3部分 (インデックス 2, 3): 縦方向 y に応じた回転
R_3d_grid[t, y, x, 2, 2] = np.cos(y * theta)
R_3d_grid[t, y, x, 2, 3] = -np.sin(y * theta)
R_3d_grid[t, y, x, 3, 2] = np.sin(y * theta)
R_3d_grid[t, y, x, 3, 3] = np.cos(y * theta)

# 3/3部分 (インデックス 4, 5): 横方向 x に応じた回転
R_3d_grid[t, y, x, 4, 4] = np.cos(x * theta)
R_3d_grid[t, y, x, 4, 5] = -np.sin(x * theta)
R_3d_grid[t, y, x, 5, 4] = np.sin(x * theta)
R_3d_grid[t, y, x, 5, 5] = np.cos(x * theta)

# --- 4. 3D RoPE の適用 (ベクトルの回転) ---
embedded_features = torch.zeros_like(patch_features)

for t in range(grid_t):
for y in range(grid_h):
for x in range(grid_w):
# 特徴量ベクトル [Batch=1, 6] 取得
q_or_k = patch_features[0, t, y, x, :]
# 対応する6x6の回転行列を取得
R = R_3d_grid[t, y, x, :, :]
# 行列掛け算(回転の適用)
embedded_features[0, t, y, x, :] = R @ q_or_k

# 最終的にトランスフォーマーブロックに渡すために1列のシーケンスに平坦化
# [Batch=1, 全パッチ数=125, 次元数=6]
final_sequence = embedded_features.flatten(1, 3)

print("元の動画形状:", video.shape)
print("パッチ分割後のグリッド形状 (時間, 縦, 横, 次元):", patch_features.shape[1:])
print("3D RoPE適用後の最終入力シーケンス形状:", final_sequence.shape)

実行結果

元の動画形状: torch.Size([1, 5, 3, 10, 10])
パッチ分割後のグリッド形状 (時間,,, 次元): torch.Size([5, 5, 5, 6])
3D RoPE適用後の最終入力シーケンス形状: torch.Size([1, 125, 6])

このコードは、torch を用いて、実際の動画データにMRoPEを適用する一連のパイプラインを実装したものです。

  1. 入力データ: まず、torch.randn(1, 5, 3, 10, 10) で、(Batch, Time, Channel, Height, Width) の5次元を持つダミーの動画テンソルを生成します。これは「10x10ピクセルのRGB画像が5フレーム続く動画」を意味します。

  2. パッチ化と線形射影: Vision Transformerと同様に、動画をパッチに分割します。video.unfold() を縦・横方向に適用することで、2x2ピクセルのパッチを切り出します。 その後、patches.permute()patches.flatten(4) を使って次元を整理し、各パッチを12次元(2x2x3)のベクトルに変換します。 最後に projection_layer(線形層)を通して、モデルの隠れ層の次元数である6次元の patch_features に変換します。この時点でテンソルの形状は [1, 5, 5, 5, 6] となり、(Batch, Time, Grid_H, Grid_W, Hidden_Dim) という、時空間グリッドに対応した構造になります。

  3. MRoPEの適用: forループを用いて、5x5x5のグリッド上のすべてのパッチ特徴量 q_or_k に対して、その座標 (t, y, x) に応じた回転行列 R を適用(@は行列積)しています。 回転行列 R_3d_grid はブロック対角行列として実装されており、特徴量ベクトルの最初の2次元を時間 t で、次の2次元を縦 y で、最後の2次元を横 x で独立して回転させます。

  4. 最終シーケンス生成: 位置情報が付与されたすべてのパッチ特徴量を embedded_features.flatten(1, 3) で1次元のシーケンスに平坦化します。最終的な形状は [1, 125, 6] となり、(Batch, 全パッチ数, Hidden_Dim) の形式でTransformerブロックに入力できる状態になります。

まとめ

本記事では、動画などの多次元データを扱うための強力な位置エンコーディング技術である MRoPE(Multi-dimensional Rotary Position Embedding) について、その理論的背景から具体的な実装までを詳細に解説しました。

記事を通じて、以下のステップでMRoPEへの理解を深めました。

  • RoPEの進化の学習: 1D RoPE、2D RoPEの限界点を学び、なぜ時間・空間を同時に扱うためにMRoPEが必要とされるのかを理解しました。
  • 理論と実装による比較: numpyを用いたシミュレーションを通じて、1D/2D RoPEでは時空間情報が欠落する一方、MRoPEはこれらを正確に捉えられることを可視化して確認しました。また、GPU効率を考慮した不均等な次元割り当ても有効であることを学びました。
  • PyTorchによる適用パイプラインの実装: 実際の動画テンソルをパッチに分割し、各パッチの時空間座標に応じてMRoPEを適用し、最終的にTransformerへの入力シーケンスを生成するまでの一連の流れをコードで実践しました。

MRoPEは、Qwen2-VLのような先進的なマルチモーダルモデルが、学習時よりも高解像度な画像や長い動画を「そのまま」扱えるようにするための根幹技術です。この「多次元回転」というエレガントなアプローチを理解することは、今後のVision-Languageモデルの進化を読み解く上で非常に重要となります。

本記事の文章・構成の一部に生成AIを使用しています。