本記事について
RoPE での量子化精度への影響を簡単なプログラムで確認するという実験をしています。
RoPE はベクトルを回転させる演算ですが、量子化によって生じた誤差も同様に回転されるため、計算結果へ影響する可能性があるという点を見ています。今回は RoPE の前後で量子化を行った場合に、Attention score が変化するかどうかを見てみたいと思います。
前回の記事は以下です。
RoPE の前後で量子化した場合の Attention score の比較
Attention score は、Query ($ Q $) と Key ($ K $) がどれだけ似ているかを表す値です。Transformer の Attention では最も重要な計算の一つであり、RoPE や量子化の影響もまずこの Attention score ($ S $) に現れます。
\begin{align}
S = \frac {QK^{T}}{\sqrt{d}}
\end{align}
Query 行列 ($ Q $) と Key 行列 ($ K $) は $ (n, d)$ 次元のベクトルです。（$ n $ はシーケンス長、$ d $ は埋め込みベクトルの次元）
また、ここで、$ Q $, $ K $ の行ベクトル、つまりトークンごとの埋め込みベクトルの集合を $q_i$, ${k_i}$ と置くと、$ { x, y } $ をベクトル $ x $ と $ y $ の内積として、以下のように書くことができます。
\begin{align}
QK^{T} =
\begin{pmatrix}
\{q_{1},k_{1}\} & \dots & \{q_{1},k_{m}\} \\
\vdots & \ddots & \vdots \\
\{q_{n},k_{1}\} & \dots & \{q_{n},k_{m}\}
\end{pmatrix}
\end{align}
Self-Attention の場合は、この行列は入力シーケンス中の各トークンペアの間でアテンションスコアを総当たりで計算することになります。
この Attention score について、以下の2パターンでリファレンス （FP32 での結果）との平均二乗誤差(MSE)がどの様になるかを見てみます。
- RoPE の前に量子化した場合の Attention score
- RoPE の後に量子化した場合の Attention score
評価用コード
import numpy as np
import matplotlib.pyplot as plt
def quantize_symmetric(x, bits):
qmin = -(2 ** (bits - 1))
qmax = 2 ** (bits - 1) - 1
scale = np.max(np.abs(x)) / qmax
if scale == 0:
return x.copy()
q = np.round(x / scale)
q = np.clip(q, qmin, qmax)
return q * scale
def apply_rope(x):
"""
x: [seq_len, dim]
"""
seq_len, dim = x.shape
assert dim % 2 == 0
half = dim // 2
positions = np.arange(seq_len)[:, None]
freqs = 1.0 / (10000 ** (np.arange(0, half) / half))
angles = positions * freqs[None, :]
cos = np.cos(angles)
sin = np.sin(angles)
x1 = x[:, 0::2]
x2 = x[:, 1::2]
y = np.empty_like(x)
y[:, 0::2] = x1 * cos - x2 * sin
y[:, 1::2] = x1 * sin + x2 * cos
return y
def mse(a, b):
return np.mean((a - b) ** 2)
def attention_score(q, k):
"""
Attention score = QK^T / sqrt(d)
"""
d = q.shape[-1]
return q @ k.T / np.sqrt(d)
def main():
np.random.seed(0)
seq_len = 128
dim = 64
Q = np.random.randn(seq_len, dim).astype(np.float32)
K = np.random.randn(seq_len, dim).astype(np.float32)
# Reference
Q_rope = apply_rope(Q)
K_rope = apply_rope(K)
ref_score = attention_score(Q_rope, K_rope)
results = []
for bits in [8, 6, 4, 3, 2]:
# Case A: Quantize -> RoPE
Q_q_before = quantize_symmetric(Q, bits)
K_q_before = quantize_symmetric(K, bits)
Q_q_before_rope = apply_rope(Q_q_before)
K_q_before_rope = apply_rope(K_q_before)
score_a = attention_score(Q_q_before_rope, K_q_before_rope)
# Case B: RoPE -> Quantize
Q_rope_q = quantize_symmetric(Q_rope, bits)
K_rope_q = quantize_symmetric(K_rope, bits)
score_b = attention_score(Q_rope_q, K_rope_q)
results.append({
"bits": bits,
"quantize_before_rope": mse(ref_score, score_a),
"quantize_after_rope": mse(ref_score, score_b),
})
print(f"{bits}-bit")
print(f" Case A: quantize before RoPE MSE = {mse(ref_score, score_a):.8f}")
print(f" Case B: quantize after RoPE MSE = {mse(ref_score, score_b):.8f}")
print()
bits = [r["bits"] for r in results]
mse_before = [r["quantize_before_rope"] for r in results]
mse_after = [r["quantize_after_rope"] for r in results]
plt.figure()
plt.plot(bits, mse_before, marker="o", label="Quantize before RoPE")
plt.plot(bits, mse_after, marker="o", label="Quantize after RoPE")
plt.xlabel("Quantization bits")
plt.ylabel("MSE")
plt.title("Quantization Error of Attention Score Before/After RoPE")
plt.gca().invert_xaxis()
plt.grid(True)
plt.legend()
plt.show()
if __name__ == "__main__":
main()
結果は以下のようになりました。
ビット数が減るほど量子化の刻み幅（step size）が大きくなるため、Q と K の各要素の誤差が増えています。Attention score は多数の積和演算なので、各要素の誤差が積み重なり、MSE が急激に増加しています。
今回の結果では、すべてのビット幅で Q,K -> Quantize -> RoPE の方が Q,K -> RoPE -> Quantize よりわずかに誤差が小さくなりました。ただ、これは乱数初期値によっては変わる可能性があり、必ずしもこの結果が再現されるわけではありませんが、RoPE による回転は量子化誤差を変化させる可能性があり、結果的にAttention scoreも変動するという結果が見て取れました。