rq_kmeans.py

1"""对应正文第 9 章 §9.3.3「RQ-KMeans 算法」(快手 OneRec 使用的语义 ID 方案)。23正文给出的算法流程(逐层残差聚类):4    c_1 = argmin_j ||x_0 - mu_j||^25    r_1 = x_0 - mu_{c_1}6    c_2 = argmin_j ||r_1 - mu_j^{(2)}||^27    ...  最终得到语义 ID (c_1, c_2, ..., c_L), 且 x ~= sum_{l=1}^{L} mu_{c_l}^{(l)}8与 RQ-VAE 相比不需要训练 Encoder/Decoder, 也不存在 STE 带来的梯度近似问题。910正文还提到改进版 RQ-KMeans++: 继承 KMeans++ 思想, 每层残差聚类时新的聚类中心11不再随机初始化, 而是按样本到已有中心距离平方成比例的概率采样:12    P(选中 x) ∝ min_{mu in 已选} ||x - mu||^213使码本更均匀覆盖向量空间, 降低陷入局部最优的风险。1415本脚本用 numpy 实现 RQ-KMeans 与 RQ-KMeans++, 在同一合成 embedding 上对比16逐层残差能量、码本使用率、语义 ID 冲突率。17"""1819from __future__ import annotations2021import random2223import numpy as np242526def make_items(n: int, d: int, n_cluster: int, rng: np.random.Generator) -> np.ndarray:27    centers = rng.normal(0, 2.0, size=(n_cluster, d))28    idx = rng.integers(0, n_cluster, size=n)29    return centers[idx] + 0.3 * rng.normal(size=(n, d))303132def _assign(x: np.ndarray, mu: np.ndarray) -> np.ndarray:33    d2 = (x ** 2).sum(1)[:, None] - 2 * x @ mu.T + (mu ** 2).sum(1)[None, :]34    return d2.argmin(1)353637def init_random(x: np.ndarray, k: int, rng: np.random.Generator) -> np.ndarray:38    """KMeans 的随机初始化: 直接随机抽 k 个样本。"""39    return x[rng.choice(len(x), size=k, replace=False)].copy()404142def init_pp(x: np.ndarray, k: int, rng: np.random.Generator) -> np.ndarray:43    """KMeans++ 初始化: 按到已有中心的距离平方成比例采样。"""44    centers = [x[rng.integers(len(x))]]45    d2 = ((x - centers[0]) ** 2).sum(1)46    for _ in range(k - 1):47        p = d2 / max(d2.sum(), 1e-12)48        centers.append(x[rng.choice(len(x), p=p)])49        nd = ((x - centers[-1]) ** 2).sum(1)50        d2 = np.minimum(d2, nd)51    return np.stack(centers)525354def kmeans(x: np.ndarray, k: int, rng: np.random.Generator, *,55           plus_plus: bool, iters: int = 15) -> tuple[np.ndarray, np.ndarray]:56    mu = init_pp(x, k, rng) if plus_plus else init_random(x, k, rng)57    assign = _assign(x, mu)58    for _ in range(iters):59        for j in range(k):60            m = assign == j61            if m.any():62                mu[j] = x[m].mean(0)63        new = _assign(x, mu)64        if (new == assign).all():65            assign = new66            break67        assign = new68    return mu, assign697071def rq_kmeans(x: np.ndarray, n_layer: int, k: int, rng: np.random.Generator,72              plus_plus: bool) -> dict:73    residual = x.copy()74    codebooks, codes, layer_energy = [], [], []75    base = (x ** 2).sum(1).mean()76    for _ in range(n_layer):77        mu, assign = kmeans(residual, k, rng, plus_plus=plus_plus)78        codebooks.append(mu)79        codes.append(assign)80        residual = residual - mu[assign]81        layer_energy.append(float((residual ** 2).sum(1).mean()))82    codes = np.stack(codes, 1)83    recon = sum(cb[codes[:, l]] for l, cb in enumerate(codebooks))84    usage = [float((np.bincount(codes[:, l], minlength=k) > 0).mean())85             for l in range(n_layer)]86    ppl = []87    for l in range(n_layer):88        p = np.bincount(codes[:, l], minlength=k) / len(x)89        nz = p[p > 0]90        ppl.append(float(np.exp(-(nz * np.log(nz)).sum())))91    uniq = len({tuple(r) for r in codes})92    return {"base": float(base), "layer_energy": layer_energy,93            "mse": float(((x - recon) ** 2).mean()),94            "usage": usage, "ppl": ppl, "uniq": uniq,95            "codes": codes}969798def main() -> None:99    random.seed(42)100    np.random.seed(42)101102    N, D, L, K = 3000, 32, 3, 64103    rng = np.random.default_rng(42)104    x = make_items(N, D, n_cluster=50, rng=rng)105106    res_rand = rq_kmeans(x, L, K, np.random.default_rng(42), plus_plus=False)107    res_pp = rq_kmeans(x, L, K, np.random.default_rng(42), plus_plus=True)108109    print("=" * 80)110    print(f"RQ-KMeans / RQ-KMeans++ (N={N}, dim={D}, L={L} 层, K={K}/层, "111          f"ID 空间 K^L={K ** L:,})")112    print("-" * 80)113    print("表 1  逐层残差能量 E||r_l||^2 (x ~= sum_l mu_{c_l}^{(l)})")114    print(f"{'层 l':>6}{'RQ-KMeans':>16}{'占初始':>10}{'RQ-KMeans++':>16}{'占初始':>10}")115    print(f"{0:>6}{res_rand['base']:>16.4f}{1.0:>10.1%}"116          f"{res_pp['base']:>16.4f}{1.0:>10.1%}")117    for l in range(L):118        a, b = res_rand["layer_energy"][l], res_pp["layer_energy"][l]119        print(f"{l + 1:>6}{a:>16.4f}{a / res_rand['base']:>10.1%}"120              f"{b:>16.4f}{b / res_pp['base']:>10.1%}")121122    print()123    print("表 2  每层码本使用率与 Perplexity (上限 K=%d)" % K)124    print("-" * 80)125    print(f"{'层 l':>6}{'随机初始化使用率':>18}{'PPL':>9}{'KMeans++使用率':>18}{'PPL':>9}")126    for l in range(L):127        print(f"{l + 1:>6}{res_rand['usage'][l]:>17.1%}{res_rand['ppl'][l]:>9.1f}"128              f"{res_pp['usage'][l]:>17.1%}{res_pp['ppl'][l]:>9.1f}")129130    print()131    print("表 3  整体对比")132    print("-" * 80)133    print(f"{'指标':<28}{'RQ-KMeans(随机)':>20}{'RQ-KMeans++':>18}")134    print(f"{'重构 MSE':<28}{res_rand['mse']:>20.5f}{res_pp['mse']:>18.5f}")135    print(f"{'相对残差(末层/初始)':<28}"136          f"{res_rand['layer_energy'][-1] / res_rand['base']:>19.2%}"137          f"{res_pp['layer_energy'][-1] / res_pp['base']:>17.2%}")138    print(f"{'唯一语义 ID 数':<28}{res_rand['uniq']:>20}{res_pp['uniq']:>18}")139    print(f"{'语义 ID 冲突率':<28}{1 - res_rand['uniq'] / N:>19.2%}"140          f"{1 - res_pp['uniq'] / N:>17.2%}")141    print(f"{'平均 Perplexity':<28}"142          f"{sum(res_rand['ppl']) / L:>20.2f}{sum(res_pp['ppl']) / L:>18.2f}")143    print("-" * 80)144    print("注: RQ-KMeans 无需神经网络与 STE, 第一层即吸收大部分能量, 后续层做细化;")145    print("    KMeans++ 初始化让码本更均匀覆盖残差空间(Perplexity 更高)。")146147    print()148    print("表 4  前 5 个物品的语义 ID(RQ-KMeans++)")149    print("-" * 80)150    for i in range(5):151        print(f"  item_{i:<4} SID = ({', '.join(str(c) for c in res_pp['codes'][i])})")152153154if __name__ == "__main__":155    main()