rq_vae.py

1"""对应正文第 9 章 §9.3.1「RQ-VAE」与 §9.3.2「RQ-VAE 的工程实践问题」。23实现内容:4  * Encoder f_theta: x -> z_0;5  * 逐层残差量化(RQ): s^l = argmin_{k in C^l} ||z_{l-1} - e_k||^2,6    z_l = z_{l-1} - e_{s^l}, 于是 z_0 ~= sum_{l=1}^{L} e_{s^l};7  * Decoder g_phi: x_hat = g_phi(sum_l e_{s^l});8  * 正文的整体训练目标 L = L_recon + lambda * sum_l ||z_{l-1} - e_{s^l}||^2,9    工程上把该量化项按 VQ-VAE 惯例拆成 codebook loss 与 commitment loss:10      ||sg[z_{l-1}] - e_{s^l}||^2 + beta * ||z_{l-1} - sg[e_{s^l}]||^2;11  * STE 梯度近似: z_STE = z_e + sg(z_q - z_e), 前向等于 z_q, 反向 dz_STE/dz_e = 1;12  * 码本坍缩诊断: p_k = n_k / sum_j n_j, Perplexity = exp(-sum_k p_k log p_k),13    均匀使用时 Perplexity = K, 远小于 K 则说明出现 Codebook Collapse。1415本脚本在合成的「带簇结构」item embedding 上训练一个小规模 RQ-VAE, 打印16每层码本的使用率 / Perplexity / 逐层累计重构误差数值表。17"""1819from __future__ import annotations2021import random2223import numpy as np24import torch25import torch.nn as nn26import torch.nn.functional as F272829def make_items(n: int, d: int, n_cluster: int, gen: torch.Generator) -> torch.Tensor:30    """合成带簇结构的 item embedding(模拟内容大模型输出)。"""31    centers = torch.randn(n_cluster, d, generator=gen) * 2.032    idx = torch.randint(0, n_cluster, (n,), generator=gen)33    return centers[idx] + 0.35 * torch.randn(n, d, generator=gen)343536class ResidualQuantizer(nn.Module):37    """L 层残差量化器, 每层码本大小 K。"""3839    def __init__(self, n_layer: int, codebook_size: int, dim: int, beta: float = 0.25):40        super().__init__()41        self.n_layer, self.K, self.beta = n_layer, codebook_size, beta42        self.codebooks = nn.ParameterList(43            [nn.Parameter(torch.randn(codebook_size, dim) * 0.1) for _ in range(n_layer)]44        )4546    def forward(self, z0: torch.Tensor):47        residual = z048        quantized = torch.zeros_like(z0)49        codes, cb_loss, cm_loss = [], z0.new_zeros(()), z0.new_zeros(())50        layer_err = []51        for layer in range(self.n_layer):52            e = self.codebooks[layer]                       # (K, d)53            # s^l = argmin_k ||z_{l-1} - e_k||^254            dist = (residual.pow(2).sum(1, keepdim=True)55                    - 2 * residual @ e.t()56                    + e.pow(2).sum(1).unsqueeze(0))57            s = dist.argmin(1)58            codes.append(s)59            e_s = e[s]60            # 正文量化项的两个工程拆分: codebook loss + commitment loss61            cb_loss = cb_loss + F.mse_loss(e_s, residual.detach())62            cm_loss = cm_loss + F.mse_loss(residual, e_s.detach())63            quantized = quantized + e_s64            residual = residual - e_s65            layer_err.append(residual.detach().pow(2).sum(1).mean().item())66        # STE: z_q 前向用量化值, 反向把梯度直通给 encoder 输出67        z_ste = z0 + (quantized - z0).detach()68        quant_loss = cb_loss + self.beta * cm_loss69        return z_ste, torch.stack(codes, 1), quant_loss, layer_err707172class RQVAE(nn.Module):73    def __init__(self, in_dim: int, lat_dim: int, n_layer: int, K: int):74        super().__init__()75        self.encoder = nn.Sequential(76            nn.Linear(in_dim, 64), nn.GELU(), nn.Linear(64, lat_dim))77        self.decoder = nn.Sequential(78            nn.Linear(lat_dim, 64), nn.GELU(), nn.Linear(64, in_dim))79        self.rq = ResidualQuantizer(n_layer, K, lat_dim)8081    def forward(self, x: torch.Tensor, lam: float = 1.0):82        z0 = self.encoder(x)83        z_ste, codes, quant_loss, layer_err = self.rq(z0)84        x_hat = self.decoder(z_ste)85        recon = F.mse_loss(x_hat, x)86        return recon + lam * quant_loss, recon, codes, layer_err878889def perplexity(codes_layer: torch.Tensor, K: int) -> tuple[float, float]:90    """返回 (Perplexity, 码本使用率)。"""91    n_k = torch.bincount(codes_layer, minlength=K).double()92    p = n_k / n_k.sum()93    nz = p[p > 0]94    return float(torch.exp(-(nz * nz.log()).sum())), float((n_k > 0).float().mean())959697def main() -> None:98    random.seed(42)99    np.random.seed(42)100    torch.manual_seed(42)101    gen = torch.Generator().manual_seed(42)102103    N, D, LAT, L, K = 3000, 48, 16, 3, 64104    x = make_items(N, D, n_cluster=40, gen=gen)105    model = RQVAE(D, LAT, L, K)106    opt = torch.optim.Adam(model.parameters(), lr=3e-3)107108    print("=" * 72)109    print(f"RQ-VAE 训练: N={N} items, dim={D}, latent={LAT}, L={L} 层, K={K}/层")110    print(f"语义 ID 空间 = K^L = {K ** L:,}")111    print("-" * 72)112    print(f"{'step':>6}{'total':>11}{'recon':>11}{'层1 PPL':>10}{'层2 PPL':>10}{'层3 PPL':>10}")113    hist = []114    for step in range(301):115        perm = torch.randperm(N, generator=gen)[:512]116        loss, recon, codes, layer_err = model(x[perm])117        opt.zero_grad()118        loss.backward()119        opt.step()120        if step % 60 == 0:121            ppls = [perplexity(codes[:, l], K)[0] for l in range(L)]122            print(f"{step:>6}{loss.item():>11.4f}{recon.item():>11.4f}"123                  + "".join(f"{p:>10.1f}" for p in ppls))124            hist.append((step, recon.item(), ppls))125126    with torch.no_grad():127        loss, recon, codes, layer_err = model(x)128        z0 = model.encoder(x)129130    print()131    print("=" * 72)132    print("表 1  各层码本使用情况(Perplexity 上限 = K = %d)" % K)133    print("-" * 72)134    print(f"{'层 l':>5}{'被用码字数':>12}{'使用率':>10}{'Perplexity':>13}{'PPL/K':>9}")135    for l in range(L):136        ppl, usage = perplexity(codes[:, l], K)137        print(f"{l + 1:>5}{int(usage * K):>12}{usage:>10.1%}{ppl:>13.2f}{ppl / K:>9.1%}")138139    print()140    print("表 2  逐层残差量化的累计逼近效果 z_0 ~= sum_l e_{s^l}")141    print("-" * 72)142    base = z0.pow(2).sum(1).mean().item()143    print(f"{'量化层数':>10}{'残差能量 E||z_l||^2':>22}{'相对初始':>12}")144    print(f"{0:>10}{base:>22.4f}{1.0:>12.1%}")145    for l, err in enumerate(layer_err, 1):146        print(f"{l:>10}{err:>22.4f}{err / base:>12.1%}")147148    print()149    print("表 3  最终重构质量")150    print("-" * 72)151    with torch.no_grad():152        x_hat = model.decoder(model.rq(z0)[0])153        mse = F.mse_loss(x_hat, x).item()154        var = x.var().item()155    print(f"输入方差            = {var:.4f}")156    print(f"重构 MSE L_recon    = {mse:.4f}")157    print(f"相对重构误差 MSE/Var = {mse / var:.2%}")158    uniq = len({tuple(r.tolist()) for r in codes})159    print(f"唯一语义 ID 数       = {uniq} / {N} items (冲突率 {1 - uniq / N:.2%})")160    print("-" * 72)161    print("注: 本实现严格按正文 L = L_recon + lambda*sum_l||z_{l-1}-e_{s^l}||^2,")162    print("    并用 STE(z_e + sg(z_q - z_e)) 把 Decoder 梯度直通到 Encoder。")163164165if __name__ == "__main__":166    main()