RoPEの量子化誤差の影響を調べる

はじめに

最近の LLM では、Llama, Qwen, Gemma などをはじめ多くのモデルで RoPE が採用されています。
一方、推論高速化や省メモリ化を目的に、Weight Quantization, KV Cache Quantization, Activation Quantization などの量子化技術も研究・実用化されています。
RoPE はベクトルを回転させる演算ですが、量子化によって生じた誤差も同様に回転されるため、計算に影響するかは精度を考えると重要です。本記事では、RoPE の前後で量子化した場合の誤差を比較して、その点を学んでみたいと思います。

RoPE とは

RoPE (Rotaty Positional Embedding) は、Attention に相対位置情報を組み込むための位置エンコーディング手法です。
Transformer の Self-Attention では、各トークンから生成した Query(Q) と Key(K) の内積によって関連度を計算します。

\begin{align}
 Attention(Q,K,V) = softmax(\frac{QK^{T}}{\sqrt{d}})V
\end{align}

しかし、この計算だけではトークンの並び順（順序情報）は考慮されません。そのため、Q と K に位置情報を組み込む仕組みが必要になります。
RoPE は、この位置情報を表現する方法の一つです。従来の位置埋め込みのように位置ベクトルを加算するのではなく、各トークンの位置に応じた回転行列を Q と K に適用します。
位置 m, n の Query と Key は、それぞれ次のように回転されます。

\begin{align}
 Q'_{m} = R(\theta_{m}) Q_{m} \\
 K'_{n} = R(\theta_{n}) K_{n}
\end{align}

2次元の場合の回転行列は次のようになります。

\begin{align}
 R(\theta) = \begin{pmatrix} \cos \theta & - \sin \theta \\ \sin \theta & \cos \theta \end{pmatrix}
\end{align}

このとき、Attention で用いられる内積は

\begin{align}
 Q'^{T}_{m} K'_{n} = Q^{T}_{m} R(\theta_{n} - \theta_{m}) K_{n}
\end{align}

となり、回転角の差だけが内積に現れます。回転角はトークンの位置に対応しているため、Attention のスコアには自然に相対位置の情報が反映されます。
実際の Transformer では、2 次元だけでなく、隣接する 2 次元ごとに異なる周波数の回転を適用します。これにより、多様なスケールの位置情報を表現でき、長い系列に対しても相対位置を効果的に扱うことができます。

Implementation of Rotary Position Embedding (RoPE)¹

なぜ量子化が問題になるのか

RoPE では、Q を回転して Q' を得ていますが、これが量子化により Q + ε になると、回転後は RQ + Rε となり、誤差成分も RoPE によって回転されます。
さらに、Attention では、

Q'K'^{T}

を計算するため、Query と Key の両方に量子化誤差が含まれると、それらの誤差項が Attention score に影響を与える可能性があります。

RoPE 前後で量子化誤差を比較

以下の誤差を比較します。

• RoPE前に量子化する場合
• RoPE後に量子化する場合

量子化は対称量子化(symmetric uniform quantization)でシミュレートします。
例えばINT4の場合、bits=4となり、qmin=-8, qmax=7、つまり[-8, 7]の範囲にxを写像します。
仮に、x=[-1.8, -0.5, 0.4, 1.6]であれば、絶対値の最大であるmax(abs(x))=1.8がqmax=7に対応し、scale=0.257となります。
写像した値を最も近い整数へ丸めたものが量子化値となります。例えば、x=[-1.8, -0.5, 0.4, 1.6]に対しては、-1.8 -> -7, -0.5 -> -2, 0.4 -> 2, 1.6 -> 6となります。

def quantize_dequantize(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

量子化後の整数値を逆量子化した実数 (q * scale) と元の x との差が量子化誤差です。
今回は以下の Case A, B の2パターンについての誤差を INT8, INT6, INT4, INT3, INT2 で比較します。

• Case A: 量子化 -> 逆量子化 -> RoPE
• Case B: RoPE -> 量子化 -> 逆量子化

Case A では、xを量子化・逆量子化して RoPE を適用します。これは、RoPEに入力される前の Q/K がすでに量子化誤差を含んでいる場合に相当します。
Case B では、xに RoPE を適用したあとに量子化・逆量子化を行います。これは、RoPE を計算したあとの Q/K を 量子化する場合に相当します。

# Case A: Quantize -> Dequantize -> RoPE
x_qdq = quantize_dequantize(x, bits)
rope_after_qdq = apply_rope(x_qdq, position)

# Case B: RoPE -> Quantize -> Dequantize
rope_qdq = quantize_dequantize(rope_fp32, bits)

INT8, INT6, INT4, INT3, INT2 について、Case A, B の 2パターンの差を平均二乗誤差(MSE)で確認すると以下のようになりました。

bits | MSE(before RoPE) | MSE(after RoPE) |
   8 |       0.00003259 |      0.00003406 |
   6 |       0.00055602 |      0.00047698 |
   4 |       0.01014037 |      0.00938462 |
   3 |       0.06220563 |      0.05733350 |
   2 |       0.53640521 |      0.47904366 |

評価コード

import numpy as np
import matplotlib.pyplot as plt


def quantize_dequantize(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, position, base=10000.0):
    d = x.shape[0]
    assert d % 2 == 0

    half = d // 2
    freq = 1.0 / (base ** (np.arange(0, half) / half))
    theta = position * freq

    x1 = x[0::2]
    x2 = x[1::2]

    cos = np.cos(theta)
    sin = np.sin(theta)

    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 cosine_similarity(a, b):
    return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))


def main():
    np.random.seed(0)

    d = 128
    position = 1024
    bits_list = [8, 6, 4, 3, 2]

    x = np.random.randn(d).astype(np.float32)

    rope_fp32 = apply_rope(x, position)

    results = []

    for bits in bits_list:
        # Case A: Quantize -> Dequantize -> RoPE
        x_qdq = quantize_dequantize(x, bits)
        rope_after_qdq = apply_rope(x_qdq, position)

        # Case B: RoPE -> Quantize -> Dequantize
        rope_qdq = quantize_dequantize(rope_fp32, bits)

        results.append({
            "bits": bits,
            "mse_quant_before_rope": mse(rope_fp32, rope_after_qdq),
            "mse_quant_after_rope": mse(rope_fp32, rope_qdq),
            "cos_quant_before_rope": cosine_similarity(rope_fp32, rope_after_qdq),
            "cos_quant_after_rope": cosine_similarity(rope_fp32, rope_qdq),
        })

    print("bits | MSE(before RoPE) | MSE(after RoPE) | Cos(before) | Cos(after)")
    for r in results:
        print(
            f"{r['bits']:>4} | "
            f"{r['mse_quant_before_rope']:.8f} | "
            f"{r['mse_quant_after_rope']:.8f} | "
            f"{r['cos_quant_before_rope']:.8f} | "
            f"{r['cos_quant_after_rope']:.8f}"
        )

    bits = [r["bits"] for r in results]
    mse_before = [r["mse_quant_before_rope"] for r in results]
    mse_after = [r["mse_quant_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 Before/After RoPE")
    plt.gca().invert_xaxis()
    plt.grid(True)
    plt.legend()
    plt.show()

if __name__ == "__main__":
    main()

結果として、　8bit ではほぼ誤差が見られませんが、2bit では MSE が大きくなり、量子化誤差が急激に増加することが分かります。これは、INT2 では INT8 よりも量子化の刻み幅が粗くなり、実数値をより大きく丸める必要があるためです。
また、Case A と Case B を比較すると、今回の実験条件では RoPE 後に量子化した方が MSE がわずかに小さいという結果になりました。
RoPE は回転変換であり、量子化誤差を含む Q + ε に適用すると RQ + Rε となります。このとき回転行列はベクトルのノルム（大きさ）を保存するため、RoPE 自体が誤差の大きさを増減させるわけではありません。
一方で、今回の実装では Per-Tensor Quantization を採用しており、テンソル全体の最大値から量子化スケールを決定しています。そのため、RoPE 適用後の値分布が今回の乱数ではわずかに量子化しやすくなり、Case B の方が小さな MSE になったと考えられます。
ただ、この実験で評価しているのは RoPE 適用後の単一ベクトルの誤差です。実際の LLM では、このベクトルは最終的に Query と Key の内積である Attention score の計算に利用されます。量子化誤差が Attention score に与える影響を低減するための量子化手法は現在も数多く研究されています。次は、RoPE の量子化誤差が Attention score にどのような影響を与えるのかも調べてみたいと思います。

1. RoFormer: Enhanced Transformer with Rotary Position Embedding (Su et al.,2021) ↩