誤差を配る量子化(GPTQ)
実装:
llm/gptq// 実行:go test ./llm/gptq/
量子化の章では最寄りの格子点に丸めて、出た誤差はそのまま捨てた。だが層の出力は内積なので、効くのは重み1つずつの誤差ではなく、入力を掛けて足し合わせたあとの誤差になる。誤差を捨てずに、まだ丸めていない重みへ肩代わりさせる。配る先と量は入力どうしの相関が決める。重みそのもののずれは増えるが、出力のずれは8分の1になる。
この章で作るもの
量子化の章で、重みを整数の格子点に丸めた。1つずつ独立に丸めて、出た誤差はそのまま捨てる。誤差は格子間隔の半分以内に収まるので、それ以上できることは無いように見える。
だが、何を小さくしたいのかを問い直すと話が変わる。
層の出力は内積になる。y = Σ wᵢxᵢ なので、効いてくるのは重み1つずつのずれではなく、入力を掛けて足し合わせたあとのずれだ。個々の丸めが半格子以内でも、同じ向きに揃えば積み上がる。
だったら、丸めて出た誤差を捨てずに、まだ丸めていない重みへ肩代わりさせればよい。前から順に1つずつ確定させ、確定するたびに残りを動かす。
肩代わりが効くのは、入力どうしが似た動きをするときになる。入力 3 と入力 7 がいつも一緒に上下しているなら、w₃ を少し減らして w₇ を少し増やしても、出力はほとんど変わらない。逆に、まったく無関係に動く入力どうしでは肩代わりできない。
素朴な丸め(誤差を捨てる)
w₀ ──丸め──▶ ŵ₀ 誤差 +0.03 捨てる
w₁ ──丸め──▶ ŵ₁ 誤差 +0.02 捨てる
w₂ ──丸め──▶ ŵ₂ 誤差 +0.04 捨てる
────────────
出力には 0.09 ぶん積み上がる
誤差を配る
w₀ ──丸め──▶ ŵ₀ 誤差 +0.03 ─┐
├─▶ w₁, w₂ を動かしてから丸める
w₁'──丸め──▶ ŵ₁ 誤差 -0.01 ─┤
w₂'──丸め──▶ ŵ₂ 誤差 -0.01 ─┘
────────────
出力には 0.01 ぶんしか残らない
重みそのものは最寄りの格子点から離れる。それでよい。
小さくしたいのは重みのずれではなく、出力のずれだから。順に見ていく。
- 小さくしたいのは出力のずれ: 重みのずれは増やしてよい。増やして出力を守る
- 配る先と量は入力の相関が決める: 相関を無視すると配る先が消え、素朴な丸めに戻る
- 分解1回で全段ぶんが手に入る: 順に確定させても、なぞるのは1回で済む
① 小さくしたいのは出力のずれ
まず、何を測るのかを分ける:
// Outputs は入力 x に対する層の出力 Xᵀw を返す。
func Outputs(w []float64, x [][]float64) []float64 {
n := len(x[0])
out := make([]float64, n)
for k := 0; k < n; k++ {
s := 0.0
for i := range w {
s += w[i] * x[i][k]
}
out[k] = s
}
return out
}
// OutputError は元の重みと丸めた重みで、出力がどれだけずれたかを返す(二乗平均平方根)。
//
// 重みそのもののずれではなく、出力のずれを見るのがこの章の要点になる。
func OutputError(w, q []float64, x [][]float64) float64 {
a := Outputs(w, x)
b := Outputs(q, x)
s := 0.0
for k := range a {
d := a[k] - b[k]
s += d * d
}
return math.Sqrt(s / float64(len(a)))
}
// WeightError は重みそのもののずれ(二乗平均平方根)。
func WeightError(w, q []float64) float64 {
s := 0.0
for i := range w {
d := w[i] - q[i]
s += d * d
}
return math.Sqrt(s / float64(len(w)))
}重みそのもののずれと、出力のずれ。素朴な丸めが小さくしているのは前者で、本当に守りたいのは後者になる。
実測がこうなった(64次元、共通成分の強い入力を256回ぶん、3bit):
| 重みのずれ | 出力のずれ | |
|---|---|---|
| 素朴な丸め | 0.08682 | 0.32736 |
| 誤差を配る | 0.09272 | 0.04069 |
重みのずれは増えているのに、出力のずれは8分の1になっている。最寄りに丸めるのをやめたのだから、重みが離れるのは当然だ。それを承知で離す。テストで、この2つが逆向きに動くことを固定した。
② 配る先と量は入力の相関が決める
誰にどれだけ肩代わりさせられるかは、入力どうしの相関に入っている:
// Hessian は入力からヘッセ行列 H = X Xᵀ / n を作る。
//
// x は列が1回ぶんの入力。H[i][j] は入力 i と入力 j がどれくらい一緒に動くかで、
// ここが大きいほど、重み i の誤差を重み j に肩代わりさせやすい。
//
// damp は対角に足す下駄。相関が強すぎると逆行列が立たなくなるので、
// 対角の平均に対する割合で少しだけ持ち上げる(実物も同じことをする)。
func Hessian(x [][]float64, damp float64) [][]float64 {
d := len(x)
n := len(x[0])
h := make([][]float64, d)
for i := range h {
h[i] = make([]float64, d)
for j := 0; j < d; j++ {
s := 0.0
for k := 0; k < n; k++ {
s += x[i][k] * x[j][k]
}
h[i][j] = s / float64(n)
}
}
mean := 0.0
for i := 0; i < d; i++ {
mean += h[i][i]
}
mean /= float64(d)
for i := 0; i < d; i++ {
h[i][i] += damp * mean
}
return h
}
// Diagonal は非対角を落として対角だけ残す。
//
// 相関を無視したことになる。逆行列も対角だけになるので、配る先が消える。
func Diagonal(h [][]float64) [][]float64 {
out := make([][]float64, len(h))
for i := range h {
out[i] = make([]float64, len(h))
out[i][i] = h[i][i]
}
return out
}Hessian が H = X Xᵀ / n を作る。H[i][j] は入力 i と入力 j がどれくらい一緒に動くかで、ここが大きいほど肩代わりが効く。
対角に少し下駄を履かせているのは、相関が強すぎると逆行列が立たなくなるからだ。実物も同じことをする。
配る手順そのものは短い:
// RoundToNearest は素朴な丸め。1つずつ独立に最寄りの格子点へ寄せ、誤差は捨てる。
func RoundToNearest(w []float64, bits int) []float64 {
scale := scaleOf(w, bits)
qmax := (1 << (bits - 1)) - 1
out := make([]float64, len(w))
for i, v := range w {
out[i] = clamp(math.Round(v/scale), qmax) * scale
}
return out
}
// Quantize は前から順に1つずつ確定させ、出た誤差を残りへ配る。
//
// plan は Plan が作った配り方(H⁻¹ のコレスキー分解)。手順は3行で書ける。
//
// ① 重み i を丸めて確定させる
// ② 出た誤差を plan[i][i] で割る(その重みが出力に効く度合いで割り戻す)
// ③ まだ確定していない j へ、plan[i][j] に比例して配る
//
// 配る先は相関が決める。相関を無視した plan なら ③ の中身が 0 になり、
// 素朴な丸めと同じ結果になる。
//
// scale は元の重みから先に決めて動かさない。配ったせいで格子まで変わると、
// 何を比べているのか分からなくなる。
func Quantize(w []float64, plan [][]float64, bits int) []float64 {
scale := scaleOf(w, bits)
qmax := (1 << (bits - 1)) - 1
work := append([]float64(nil), w...)
out := make([]float64, len(w))
for i := range work {
q := clamp(math.Round(work[i]/scale), qmax) * scale
out[i] = q
err := (work[i] - q) / plan[i][i]
for j := i + 1; j < len(work); j++ {
work[j] -= err * plan[i][j]
}
}
return out
}丸めて、出た誤差を対角で割り戻して、まだ確定していない先へ比例配分する。それだけだ。
相関の強さを変えて測ると、得の大きさがそのまま動く:
| 入力の共通成分 | 素朴な丸め | 誤差を配る | 比 |
|---|---|---|---|
| 0.0 | 0.37310 | 0.37350 | 1.00 |
| 0.5 | 0.25517 | 0.18892 | 1.35 |
| 0.7 | 0.27364 | 0.11339 | 2.41 |
| 0.9 | 0.32736 | 0.04069 | 8.05 |
| 0.95 | 0.34461 | 0.02488 | 13.85 |
共通成分が 0、つまり入力がまったく無関係に動く場合は、配っても得が無い(1.00倍)。テストで、この場合に結果がほとんど変わらないこと、そして共通成分を増やすほど得が大きくなることを固定した。
同じことを別の角度からも確かめられる。ヘッセ行列の非対角を落として対角だけにすると、相関を無視したことになる。すると配る先が消えて、素朴な丸めと1ビットも違わない結果になる。テストで、全要素が一致することを固定した。
LLM の重み量子化でこの手が効くのは、活性の次元どうしに強い相関があるからになる。相関の無いところでは、この工夫は空回りする。
動かす
下のデモは、入力の相関の強さを変えながら、素朴な丸めと誤差を配る場合の出力のずれを比べる。相関を 0 にすると差が消える。
64 個の重み。色が付いているのが、素朴な丸めとは別の格子点へ行ったもの(4 個)
共通成分は、入力どうしがどれくらい一緒に動くかを表す。一緒に動く入力どうしなら、片方の重みを 減らして他方を増やしても出力はほとんど変わるので、丸めた誤差を肩代わりさせられる。0 にすると 肩代わりできる先が消え、素朴な丸めと同じところに落ち着く。緑と赤が列ごとに入れ替わっているのは、 重みのずれを増やして出力のずれを減らしているからになる。
③ 分解1回で全段ぶんが手に入る
ここが GPTQ が元の手法(OBQ)から進んだところになる。
理屈だけを追うと、1つ確定させるたびに「残った重みについてのヘッセ行列」を作り直す必要がある。行と列を1つ削った行列の逆行列は、元の逆行列から行と列を削ったものとは違うからだ。素直にやると、重み1つごとに逆行列を取り直すことになる。
// CholeskyUpper は対称正定値な行列 A を A = Uᵀ U と分解して、上三角 U を返す。
//
// これが GPTQ の肝になる。順に確定させるたびに、残りの重みについての
// ヘッセ行列を作り直す必要があるのだが、素直にやると毎段で逆行列が要る。
// H⁻¹ をコレスキー分解しておくと、その i 行目がそのまま i 段目の配り方になる。
// 分解1回で全段ぶんが手に入るので、なぞるのは1回で済む。
func CholeskyUpper(a [][]float64) [][]float64 {
d := len(a)
l := make([][]float64, d)
for i := range l {
l[i] = make([]float64, d)
}
for i := 0; i < d; i++ {
for j := 0; j <= i; j++ {
s := a[i][j]
for k := 0; k < j; k++ {
s -= l[i][k] * l[j][k]
}
if i == j {
if s <= 0 {
s = 1e-12 // 数値のゆらぎで負に落ちたときの下限
}
l[i][i] = math.Sqrt(s)
continue
}
l[i][j] = s / l[j][j]
}
}
u := make([][]float64, d)
for i := range u {
u[i] = make([]float64, d)
for j := 0; j < d; j++ {
u[i][j] = l[j][i]
}
}
return u
}
// Plan は入力から「各段でどう配るか」を作る。H⁻¹ のコレスキー分解になる。
func Plan(x [][]float64, damp float64) [][]float64 {
return CholeskyUpper(Inverse(Hessian(x, damp)))
}CholeskyUpper がこれを1回で片づける。H⁻¹ を Uᵀ U に分解しておくと、上三角 U の i 行目が、そのまま i 段目の配り方になる。削っていく操作が、分解の中にすでに入っている。
だから、重みを前から1回なぞるだけで済む。逆行列も分解も最初の1回だけで、あとは掛け算と引き算しか出てこない。数万次元の層を現実的な時間で量子化できるのは、これがあるからだ。
順に確定させる形なので、端では性質が変わる。先頭の重みには誰も配っていないので、素朴な丸めとまったく同じところへ行く。テストで固定した。逆に最後の重みは、それまでの全部を受け取ったうえで、配る先が無い。
なお、格子間隔は元の重みから先に決めて動かさない。配ったせいで最大値が変わり格子まで動くと、何と何を比べているのか分からなくなる。
設計の観点
- 何を小さくしたいのかを先に決める: 中間の量を小さくしても、最後に効く量が小さくなるとは限らない
- 中間の精度は譲ってよい: 重みのずれを増やして出力のずれを減らす。譲れる場所を見つける
- 効く条件を言えるようにする: 相関が無ければ得が無い。効かない場合を測って示す
- 同じ計算を繰り返さない: 分解1回に畳み込めるなら、段ごとに解き直さない
- 順に確定させる形は端が特殊になる: 先頭と末尾の扱いを把握しておく
- 他の工夫と組み合わせる: 区画ごとの格子や外れ値の分離とは目的が違うので、同時に使える
対照と実例
| 手法 | 何を動かすか | 効く条件 | 要るもの |
|---|---|---|---|
| 素朴な丸め | 何も動かさない | 常に | 重みだけ |
| 区画ごとの格子 | 格子の細かさ | 区画の中でレンジが揃う | 重みだけ |
| 外れ値の分離 | 持ち方 | 桁違いの値がごく一部 | 重みだけ |
| 列の引き伸ばし(AWQ) | 格子の配り方 | 効きの偏りがある | 活性の大きさ |
| 誤差を配る(GPTQ) | 丸め先 | 入力に相関がある | 入力の相関(ヘッセ行列) |
裏どり:
- GPTQ(2022): Frantar et al.。OBQ を大規模モデルに載せられるところまで速くした。列の順を固定してよいこと、コレスキー分解で段ごとの更新をまとめられることが要点
- OBQ / OBS: 元をたどると LeCun らの Optimal Brain Damage、Hassibi らの Optimal Brain Surgeon。枝刈りのための道具だったものを、丸め先を選ぶ問題に持ち込んだ形になる
- キャリブレーションデータ: ヘッセ行列は実際の入力から作るので、代表的な入力を数百件ほど流す必要がある。重みだけを見る手法との大きな違い
- damping: 対角に下駄を履かせないと逆行列が立たない。実装では対角の平均に対する割合(1%程度)を足すのが定番
- 組み合わせ: 実物の GPTQ は区画ごとの格子と併用する。この章では配ることだけを取り出している
簡略化したこと
- 1行だけ: 実物は行列の全行を同時に処理する。ここでは1本の重みベクトルだけ
- 区画ごとの格子なし: 格子は全体で1つ。実物は 64〜128 要素ごとに持つ
- 列の並べ替えなし: 実物は効きの大きい列から先に確定させる選択肢がある
- 段階的な更新なし: 実物はブロックごとに区切って、まとめて更新する(計算の局所性のため)
- キャリブレーションが人工的: 入力は共通成分と揺れから作る。実物は実際の文章を流して集める
- 枝刈りなし: 同じ枠組みで重みを 0 にする(枝刈り)話は扱わない
参考資料
- Frantar et al., GPTQ: Accurate Post-Training Quantization for Generative Pre-trained Transformers(2022)
- Frantar, Alistarh, Optimal Brain Compression(2022) — OBQ。GPTQ の土台
- Hassibi, Stork, Second Order Derivatives for Network Pruning: Optimal Brain Surgeon(1992)
- 実装: llm/gptq