other_tokenizers.py

1"""对应正文第 9 章 §9.3.4「其他离散 Tokenizer 方法」。23正文把离散 Tokenizer 形式化为: 学习一个 Codebook 矩阵 C in R^{m x d}, 把 Embedding4矩阵 Z in R^{n x d} 中的每个向量映射到最近的码字:5    q: z -> argmin_{k=1..m} D(z, c_k)6并给出了三类范式:7  (1) Level-wise Quantization(层次化): RQ-VAE / RQ-KMeans, 逐层量化残差;8  (2) Group-wise Quantization(分组): Product Quantization(PQ), 把 d 维向量切成 S 段9      x = concat(x_1, ..., x_S), 每段独立量化, C(x) = (C_1(x_1), ..., C_S(x_S));10  (3) Lookup-Free Quantization(无查找表): FSQ 与 LFQ。11      FSQ: 先把表征映射到很少的维度, 再逐维四舍五入到有限标量集合, 隐式构成码本,12           z_hat_i = round(L_i/2 * tanh(z_i)) 这类形式, 隐式码本大小 = prod_i L_i;13      LFQ: 把码本维度直接降到 1, 每维只取 {-1, +1}, q(z)_i = sign(z_i), 隐式码本14           大小 2^d; 并引入熵正则提升码本利用率:15           L_entropy = E_x[H(q(z))] - H(E_x[q(z)])16           其中第一项降低单样本不确定性, 第二项鼓励样本在码本空间均匀分布。1718本脚本在**同一份合成 item embedding**上对比 RQ(层次化)、PQ(分组)、FSQ、LFQ19四种 Tokenizer 的重构误差、隐式/显式码本大小、码本利用率与语义 ID 冲突率。20"""2122from __future__ import annotations2324import math25import random2627import numpy as np282930def make_items(n: int, d: int, n_cluster: int, rng: np.random.Generator) -> np.ndarray:31    centers = rng.normal(0, 2.0, size=(n_cluster, d))32    idx = rng.integers(0, n_cluster, size=n)33    return centers[idx] + 0.3 * rng.normal(size=(n, d))343536def _kmeans(x: np.ndarray, k: int, rng: np.random.Generator, iters: int = 12):37    mu = x[rng.choice(len(x), size=min(k, len(x)), replace=False)].copy()38    assign = np.zeros(len(x), dtype=np.int64)39    for _ in range(iters):40        d2 = (x ** 2).sum(1)[:, None] - 2 * x @ mu.T + (mu ** 2).sum(1)[None, :]41        assign = d2.argmin(1)42        for j in range(len(mu)):43            m = assign == j44            if m.any():45                mu[j] = x[m].mean(0)46    return mu, assign474849# ---------------------------------------------------------------- (1) 层次化 RQ50def rq_tokenize(x: np.ndarray, L: int, K: int, rng: np.random.Generator):51    residual, codes, cbs = x.copy(), [], []52    for _ in range(L):53        mu, a = _kmeans(residual, K, rng)54        cbs.append(mu)55        codes.append(a)56        residual = residual - mu[a]57    codes = np.stack(codes, 1)58    recon = sum(cb[codes[:, i]] for i, cb in enumerate(cbs))59    used = sum(int((np.bincount(codes[:, i], minlength=K) > 0).sum()) for i in range(L))60    return codes, recon, L * K, used, K ** L616263# ---------------------------------------------------------------- (2) 分组 PQ64def pq_tokenize(x: np.ndarray, S: int, K: int, rng: np.random.Generator):65    """把 d 维切成 S 段, 每段独立 k-means 量化。"""66    d = x.shape[1]67    assert d % S == 068    seg = d // S69    codes, recon = [], np.zeros_like(x)70    used = 071    for s in range(S):72        sub = x[:, s * seg:(s + 1) * seg]73        mu, a = _kmeans(sub, K, rng)74        codes.append(a)75        recon[:, s * seg:(s + 1) * seg] = mu[a]76        used += int((np.bincount(a, minlength=K) > 0).sum())77    return np.stack(codes, 1), recon, S * K, used, K ** S787980# ---------------------------------------------------------------- (3) FSQ81def fsq_tokenize(x: np.ndarray, levels: list[int], rng: np.random.Generator):82    """FSQ: 先线性降维到 len(levels) 维, 再逐维 bounded + round。8384    每维 i 取 L_i 个标量取值: z_hat_i = round((L_i-1)/2 * tanh(z_i)),85    隐式码本大小 = prod_i L_i, 无需查找表。86    """87    d, m = x.shape[1], len(levels)88    W = rng.normal(0, 1.0 / math.sqrt(d), size=(d, m))   # 降维映射 f89    Wi = np.linalg.pinv(W)                               # 解码(最小二乘逆)90    z_raw = x @ W91    scale = z_raw.std(0, keepdims=True) + 1e-8           # 标准化使 tanh 落在有效区92    z = z_raw / scale93    half = np.array([(l - 1) / 2 for l in levels])94    q = np.round(half * np.tanh(z))                      # 逐维四舍五入到有限标量集合95    q = np.clip(q, -half, half)96    # 解码: 由离散标量反推 tanh 值 -> z -> 原空间97    z_hat = np.arctanh(np.clip(q / half, -0.999, 0.999)) * scale98    recon = z_hat @ Wi99    idx = np.rint(q + half).astype(np.int64)             # 每维编号 0..L_i-1100    codes = np.zeros(len(x), dtype=np.int64)101    radix = 1102    for i, l in enumerate(levels):103        codes = codes + idx[:, i] * radix104        radix *= l105    implicit = int(np.prod(levels))106    used = len(np.unique(codes))107    return idx, recon, 0, used, implicit108109110# ---------------------------------------------------------------- (4) LFQ111def lfq_tokenize(x: np.ndarray, m: int, rng: np.random.Generator,112                 tau: float = 1.0):113    """LFQ: 码本维度降为 1, 每维取值 {-1,+1}, q(z)_i = sign(z_i), 隐式码本 2^m。114115    同时按正文公式计算熵正则项:116        L_entropy = E_x[H(q(z))] - H(E_x[q(z)])117        H(q(z)) = sum_i H(q_hat(c_i | z_i)), q_hat 为 softmax(-tau*|c_i - z_i|) 的伯努利分布118    """119    d = x.shape[1]120    W = rng.normal(0, 1.0 / math.sqrt(d), size=(d, m))121    Wi = np.linalg.pinv(W)122    z_raw = x @ W123    scale = z_raw.std(0, keepdims=True) + 1e-8124    z = z_raw / scale125    q = np.where(z >= 0, 1.0, -1.0)126    recon = (q * scale) @ Wi127    # 伯努利概率: p_i = exp(-tau*|1 - z_i|) / (exp(-tau*|1-z_i|) + exp(-tau*|-1-z_i|))128    a = np.exp(-tau * np.abs(1.0 - z))129    b = np.exp(-tau * np.abs(-1.0 - z))130    p = a / (a + b)131132    def H(pp):133        pp = np.clip(pp, 1e-9, 1 - 1e-9)134        return -(pp * np.log(pp) + (1 - pp) * np.log1p(-pp))135136    per_sample_H = H(p).sum(1).mean()          # E_x[H(q(z))]137    marginal_H = H(p.mean(0)).sum()            # H(E_x[q(z)])138    l_entropy = per_sample_H - marginal_H139    bits = ((q > 0).astype(np.int64))140    codes = np.zeros(len(x), dtype=np.int64)141    for i in range(m):142        codes = codes | (bits[:, i] << i)143    return bits, recon, 0, len(np.unique(codes)), 2 ** m, l_entropy, per_sample_H, marginal_H144145146def rel_mse(x: np.ndarray, recon: np.ndarray) -> float:147    return float(((x - recon) ** 2).mean() / x.var())148149150def main() -> None:151    random.seed(42)152    np.random.seed(42)153    rng = np.random.default_rng(42)154155    N, D = 2000, 32156    x = make_items(N, D, n_cluster=50, rng=rng)157158    rq_codes, rq_rec, rq_tbl, rq_used, rq_space = rq_tokenize(159        x, L=3, K=32, rng=np.random.default_rng(42))160    pq_codes, pq_rec, pq_tbl, pq_used, pq_space = pq_tokenize(161        x, S=4, K=32, rng=np.random.default_rng(42))162    fsq_idx, fsq_rec, _, fsq_used, fsq_space = fsq_tokenize(163        x, levels=[8, 8, 5, 5], rng=np.random.default_rng(42))164    (lfq_bits, lfq_rec, _, lfq_used, lfq_space,165     l_ent, h_sample, h_marg) = lfq_tokenize(x, m=12, rng=np.random.default_rng(42))166167    print("=" * 94)168    print(f"同一合成 item embedding 上的离散 Tokenizer 对比 (N={N}, dim={D})")169    print("-" * 94)170    print(f"{'方法':<26}{'范式':<18}{'显式码字数':>11}{'隐式ID空间':>13}"171          f"{'实际用到ID':>12}{'相对MSE':>10}")172    rows = [173        ("RQ (L=3,K=32)", "Level-wise", rq_tbl, rq_space, len({tuple(r) for r in rq_codes}),174         rel_mse(x, rq_rec)),175        ("PQ (S=4,K=32)", "Group-wise", pq_tbl, pq_space, len({tuple(r) for r in pq_codes}),176         rel_mse(x, pq_rec)),177        ("FSQ (levels=8,8,5,5)", "Lookup-Free", 0, fsq_space, fsq_used,178         rel_mse(x, fsq_rec)),179        ("LFQ (m=12 bits)", "Lookup-Free", 0, lfq_space, lfq_used,180         rel_mse(x, lfq_rec)),181    ]182    for name, para, tbl, space, used, mse in rows:183        tbl_s = "-(无查找表)" if tbl == 0 else f"{tbl}"184        print(f"{name:<26}{para:<18}{tbl_s:>11}{space:>13,}{used:>12}{mse:>10.3f}")185186    print()187    print("表 2  码本/ID 空间利用率")188    print("-" * 94)189    print(f"{'方法':<26}{'实际ID数':>11}{'ID空间':>13}{'空间利用率':>13}{'样本覆盖度':>13}")190    for name, _, _, space, used, _ in rows:191        print(f"{name:<26}{used:>11}{space:>13,}{used / space:>13.2%}{used / N:>13.1%}")192193    print()194    print("表 3  LFQ 熵正则项各分量 (L_entropy = E_x[H(q(z))] - H(E_x[q(z)]))")195    print("-" * 94)196    print(f"{'分量':<40}{'数值(nats)':>14}{'说明':<30}")197    print(f"{'E_x[H(q(z))] 单样本平均熵':<40}{h_sample:>14.4f}  越小越确定")198    print(f"{'H(E_x[q(z)]) 边缘分布熵':<40}{h_marg:>14.4f}  越大越均匀(上限 m*ln2="199          f"{12 * math.log(2):.3f})")200    print(f"{'L_entropy':<40}{l_ent:>14.4f}  训练中最小化该项")201202    print()203    print("表 4  各方法逐层/逐组的 ID 分布(前 3 个物品的离散编码)")204    print("-" * 94)205    for i in range(3):206        print(f"  item_{i}: RQ={tuple(rq_codes[i])}  PQ={tuple(pq_codes[i])}  "207              f"FSQ={tuple(int(v) for v in fsq_idx[i])}  "208              f"LFQ={''.join(str(int(b)) for b in lfq_bits[i])}")209    print("-" * 94)210    print("注: 相对 MSE = MSE / Var(x)。RQ/PQ 依赖显式码本查找, 重构最准;")211    print("    FSQ/LFQ 无需查找表、天然不会出现'整块码字死亡', 但受降维瓶颈限制,")212    print("    重构误差更高——与正文'无查找表量化以表达能力换工程简洁'的结论一致。")213214215if __name__ == "__main__":216    main()