| """トレリス符号化量子化（QTIP方式の中核）。 |
| |
| k-means辞書の限界: k個をまとめると候補が 2^(k×bit) で爆発し、k=32が壁。 |
| トレリスは重みを1個ずつ符号化しつつ「直近Lビットの履歴」を状態として持ち、 |
| 系列全体の誤差が最小になる経路をViterbiで厳密に探す。 |
| 再生値は状態番号から計算で生成するので**コードブックを保存しない**。 |
| |
| 高速化の鍵（ビットシフトトレリスの構造）: |
| 遷移は next = ((state << K) | code) & (2^L - 1)。 |
| K=1 のとき、ある次状態 j の前状態は必ず 2つだけ: j>>1 と (j>>1) + 2^(L-1)。 |
| この「バタフライ構造」のおかげで、全状態・全行を一度に配列演算で処理できる。 |
| """ |
| import numpy as np |
|
|
|
|
| def gen_values(L, seed=0): |
| """状態ごとの再生値。保存不要＝シードだけで再現できる。""" |
| rng = np.random.default_rng(seed) |
| v = rng.standard_normal(1 << L) |
| return ((v - v.mean()) / v.std()).astype(np.float32) |
|
|
|
|
| def quantize_rows(X, L=10, seed=0): |
| """X: (N, T) の各行をトレリス符号化（K=1bit/要素）。全行を同時に処理する。""" |
| N, T = X.shape |
| S = 1 << L |
| H = S >> 1 |
| vals = gen_values(L, seed) |
| j = np.arange(S) |
| p0 = (j >> 1) |
| p1 = p0 + H |
|
|
| cost = np.zeros((N, S), np.float32) |
| back = np.empty((T, N, S), np.bool_) |
| for t in range(T): |
| c0 = cost[:, p0] |
| c1 = cost[:, p1] |
| take1 = c1 < c0 |
| back[t] = take1 |
| base = np.where(take1, c1, c0) |
| d = vals[None, :] - X[:, t:t + 1] |
| cost = base + d * d |
|
|
| |
| s = np.argmin(cost, axis=1).astype(np.int64) |
| out = np.empty((N, T), np.float32) |
| rows = np.arange(N) |
| for t in range(T - 1, -1, -1): |
| out[:, t] = vals[s] |
| take1 = back[t, rows, s] |
| s = (s >> 1) + np.where(take1, H, 0) |
| return out |
|
|
|
|
| def quantize_matrix(W, G=256, L=10, seed=0, ls_scale=True): |
| """行列全体。グループ(=行の一区間)ごとにスケール、中身を1bitトレリスで符号化。 |
| 実効ビット = 1 + スケールのビット/G。""" |
| X = W.reshape(-1, G) |
| s = np.abs(X).max(1, keepdims=True) |
| s[s == 0] = 1 |
| R = quantize_rows((X / s).astype(np.float32), L=L, seed=seed) |
| if ls_scale: |
| num = (X * R).sum(1, keepdims=True) |
| den = (R * R).sum(1, keepdims=True) |
| den[den == 0] = 1 |
| s = num / den |
| return (R * s).reshape(W.shape).astype(np.float32), 1.0 + 8.0 / G |
|
|
