""" ============================================================================= CS236 · 第 5 讲模块：lvm.py —— 隐变量模型的小工具：EM、ELBO、一个有 2^m 种隐状态的世界 ============================================================================= 【这个模块是干什么的】 第 5 讲的主角是隐变量模型 p(x) = Σ_z p(x, z)。这里提供本讲各脚本共用的几样东西： fit_gmm(X, K, generator, ...) 用 EM 算法拟合高斯混合（toy2d.GaussianMixture），可记录每一步的 ELBO MaskedBernoulliMixture(K) 伯努利混合，EM 训练**允许训练图有缺失像素**（缺失像素被精确地边缘化） elbo_discrete(log_joint, q) 离散隐变量的 ELBO = Σ_z q(z) [log p(x, z) − log q(z)] kl_discrete(q, log_p) KL(q ‖ p) = Σ_z q(z) [log q(z) − log p(z)] log_mean_exp(a, dim) log( (1/k) Σ exp(a) )：重要性采样估计 log p(x) 时用 StrokeWorld(m) "笔画世界"：m 个二值隐特征（m ≤ 16），隐状态有 2^m 种， log p(x) 只能把 2^m 项全部加起来——用来感受"求和的项数随隐变量个数指数增长" 和第 1 讲 baselines.BernoulliMixture 的区别：那个类只处理完整的图；这里的版本支持缺失像素， 并能记录 EM 每一步之后的 ELBO（用来验证"EM 是在 ELBO 上做坐标上升"）。 【运行方式】 python3 lvm.py # 自检：EM 的单调性、ELBO 恒等式、笔画世界的 Σ_x p(x) = 1（约 3 秒） 【Python 小知识】 - (torch.arange(2**m).unsqueeze(1) >> torch.arange(m)) & 1：把 0 … 2^m − 1 的每个整数拆成 m 个二进制位， 得到全部 2^m 种 0/1 组合（>> 是右移，& 1 取最低位）。 """ import math import os import sys HERE = os.path.dirname(os.path.abspath(__file__)) sys.path.insert(0, os.path.join(HERE, "..", "lecture01")) import torch from toy2d import GaussianMixture # ============================================================================= # 小工具：log-mean-exp、离散 ELBO、离散 KL # ============================================================================= def log_mean_exp(a, dim=-1): """ log((1/k) Σ_j exp(a_j))，数值稳定（logsumexp 先减最大值）。 【DGM 联系】重要性采样估计 p(x) ≈ (1/k) Σ_j w_j，其中 w_j = p(x, z_j)/q(z_j)； 权重常常小到 1e-30 以下，只能在对数空间里平均：log p̂ = log_mean_exp(log w)。 """ return torch.logsumexp(a, dim=dim) - math.log(a.shape[dim]) def _xlogx(q): """q log q，约定 0 log 0 = 0（避免 0 × (−∞) = nan）。""" return torch.where(q > 0, q * torch.log(torch.where(q > 0, q, torch.ones_like(q))), torch.zeros_like(q)) def elbo_discrete(log_joint, q): """ 离散隐变量的证据下界 ELBO(q) = Σ_z q(z) log p(x, z) + H(q)。 【参数】log_joint：(N, K) 的 log p(x, z)；q：(N, K) 的分布（每行和为 1，可以有 0） 【返回】(N,) 张量（单位 nat） 【边界】q(z) = 0 的项不计入（即使 log p(x, z) = −∞ 也不会产生 nan）。 """ lj = torch.where(q > 0, log_joint, torch.zeros_like(log_joint)) return (q * lj).sum(dim=1) - _xlogx(q).sum(dim=1) def kl_discrete(q, log_p): """ KL(q ‖ p) = Σ_z q(z) [log q(z) − log p(z)]。 【参数】q：(N, K) 分布；log_p：(N, K) 或 (K,) 的 log p(z) 【返回】(N,) 张量（nat） """ log_p = log_p.expand_as(q) lp = torch.where(q > 0, log_p, torch.zeros_like(q)) return _xlogx(q).sum(dim=1) - (q * lp).sum(dim=1) # ============================================================================= # 高斯混合的 EM # ============================================================================= def fit_gmm(X, K, generator, n_iter=300, tol=1e-8, reg=1e-6, means_init=None): """ 用 EM 算法拟合 K 个成分的高斯混合（全协方差）。 【思路】EM = 在 ELBO(q, θ) = Σ_n Σ_k q_nk [log p(x_n, k; θ) − log q_nk] 上做坐标上升： E 步：固定 θ，令 q_nk = p(k | x_n; θ)（后验，"责任"）。这时 ELBO = log p(x; θ)，下界是紧的。 M 步：固定 q，最大化 ELBO 里和 θ 有关的部分 Σ q_nk log p(x_n, k; θ)——它有闭式解： π_k = N_k / N，μ_k = Σ_n q_nk x_n / N_k，Σ_k = Σ_n q_nk (x_n − μ_k)(x_n − μ_k)ᵀ / N_k， 其中 N_k = Σ_n q_nk（第 k 个成分"认领"了多少个点）。 于是 log p(x; θ_t) = ELBO(q_t, θ_t) ≤ ELBO(q_t, θ_{t+1}) ≤ log p(x; θ_{t+1})：对数似然单调不降。 【参数】 X：(N, D) 数据；K：成分数；generator：初始化用的随机数生成器 n_iter：最多迭代几轮；tol：平均对数似然的提升小于 tol（nat/点）时停止 reg：加到协方差对角线上的小数，防止某个成分缩成一个点（协方差奇异、似然 → +∞） means_init：None（随机挑 K 个数据点当初始均值）或 (K, D) 张量 【返回】(model, history) model：toy2d.GaussianMixture history：列表，第 t 个元素是字典 {"ll": log p(x; θ_t), "elbo_m": ELBO(q_t, θ_{t+1})}（都是每个点的平均，nat） 【边界】 - EM 只保证收敛到**局部**最优；不同的初始化可能停在不同的解（第 5 讲 04 号脚本的第 2 部分）。 - reg > 0 时 M 步不再是精确的最大化，单调性可能在 1e-6 量级上被破坏；验证单调性时请用 reg = 0。 - 某个成分一个点都没认领到（N_k ≈ 0）时，它的均值会退化，之后一直"死"着——不报错。 【PyTorch 小知识】torch.einsum("nk,nd,ne->kde", r, d, d) 一次算出 K 个加权的外积和： 对每个 k，Σ_n r_nk d_n d_nᵀ（这里 d 实际上按成分分别计算，见代码）。 """ X = torch.as_tensor(X, dtype=torch.float64) N, D = X.shape if means_init is None: means = X[torch.randperm(N, generator=generator)[:K]].clone() else: means = torch.as_tensor(means_init, dtype=torch.float64).clone() var = X.var(dim=0).mean().item() model = GaussianMixture(torch.full((K,), 1.0 / K, dtype=torch.float64), means, var * torch.eye(D, dtype=torch.float64).repeat(K, 1, 1)) eye = torch.eye(D, dtype=torch.float64) history = [] for _ in range(n_iter): lj = model.log_joint(X) # E 步 ll = torch.logsumexp(lj, dim=1) r = torch.exp(lj - ll.unsqueeze(1)) # 责任 q_nk = p(k | x_n) Nk = r.sum(dim=0) + 1e-12 # M 步 w = Nk / Nk.sum() mu = (r.T @ X) / Nk.unsqueeze(1) diff = X.unsqueeze(0) - mu.unsqueeze(1) # (K, N, D) cov = torch.einsum("nk,knd,kne->kde", r, diff, diff) / Nk.reshape(-1, 1, 1) + reg * eye new = GaussianMixture(w, mu, cov) elbo_m = elbo_discrete(new.log_joint(X), r).mean().item() # 用旧的 q、新的 θ 算 ELBO history.append({"ll": ll.mean().item(), "elbo_m": elbo_m}) model = new if len(history) >= 2 and history[-1]["ll"] - history[-2]["ll"] < tol: break history.append({"ll": model.log_prob(X).mean().item(), "elbo_m": float("nan")}) return model, history # ============================================================================= # 允许缺失像素的伯努利混合 # ============================================================================= class MaskedBernoulliMixture: """ 伯努利混合 p(x) = Σ_k π_k Π_i μ_ki^{x_i} (1 − μ_ki)^{1 − x_i}，EM 训练时允许每张训练图缺一部分像素。 【参数】K：原型个数；alpha：μ 的平滑伪计数（防止 μ 变成严格的 0 或 1）；n_iter、tol：EM 的停止条件 【DGM 联系】原课第 5 讲 slide 15："假设训练时有些像素缺失（例如上半张）"。缺失的像素也是隐变量： p(x_观测) = Σ_k π_k Π_{i 观测} p(x_i | k)。因为**给定 k 时像素相互独立**，把缺失像素加掉只是把它们从乘积里拿掉； 于是 E 步只用观测到的像素算责任，M 步只用观测到的像素估计 μ。 对一般的模型（例如 PixelCNN），对缺失像素求和要遍历 2^(缺失个数) 种补全，这就难了。 """ def __init__(self, K, alpha=0.5, n_iter=200, tol=1e-7): self.K, self.alpha, self.n_iter, self.tol = K, alpha, n_iter, tol self.log_pi = None self.mu = None self.history = [] def _log_lik(self, X, M): """log p(x_观测 | k)，形状 (N, K)：只对 M = 1 的像素求和。""" lm, l1m = self.mu.log(), (1 - self.mu).log() return (X * M) @ lm.T + ((1 - X) * M) @ l1m.T def fit(self, X, generator, mask=None, init_mu=None): """ EM 训练。 【参数】X：(N, D) 0/1 张量（缺失位置的值会被忽略）；generator：初始化用； mask：None（全部观测）或 (N, D) 的 0/1 张量（1 = 观测到）；init_mu：None 或 (K, D) 的初始原型 【返回】self；self.history 是每一轮的 {"ll": 平均 log p(x_观测)（nat/图）} 【思路】 初始化：随机挑 K 张训练图，观测到的像素软化成 0.25 / 0.75，缺失的设成 0.5。 E 步：r_nk ∝ π_k p(x_n,观测 | k)。 M 步：π_k ∝ Σ_n r_nk；μ_ki = (Σ_n r_nk m_ni x_ni + α) / (Σ_n r_nk m_ni + 2α) ——只在"观测到第 i 个像素"的图里统计。 【边界】某个像素在所有训练图里都缺失时，μ_ki 只由平滑决定（= 0.5）。 """ X = torch.as_tensor(X, dtype=torch.float64) N, D = X.shape M = torch.ones_like(X) if mask is None else torch.as_tensor(mask, dtype=torch.float64) if init_mu is None: idx = torch.randint(0, N, (self.K,), generator=generator) self.mu = M[idx] * (0.25 + 0.5 * X[idx]) + (1 - M[idx]) * 0.5 else: self.mu = torch.as_tensor(init_mu, dtype=torch.float64).clone() self.log_pi = torch.full((self.K,), -math.log(self.K), dtype=torch.float64) self.history = [] for _ in range(self.n_iter): lj = self._log_lik(X, M) + self.log_pi ll = torch.logsumexp(lj, dim=1) r = torch.exp(lj - ll.unsqueeze(1)) Nk = r.sum(dim=0) self.log_pi = ((Nk + 1e-3) / (N + 1e-3 * self.K)).log() self.mu = (r.T @ (X * M) + self.alpha) / (r.T @ M + 2 * self.alpha) self.history.append({"ll": ll.mean().item()}) if len(self.history) >= 2 and self.history[-1]["ll"] - self.history[-2]["ll"] < self.tol: break return self def log_joint(self, X, mask=None): """log p(x_观测, k)，形状 (N, K)。""" X = torch.as_tensor(X, dtype=torch.float64) M = torch.ones_like(X) if mask is None else torch.as_tensor(mask, dtype=torch.float64).expand_as(X) return self._log_lik(X, M) + self.log_pi def log_prob(self, X, mask=None): """log p(x)；mask 给出时是 log p(x_观测)。""" return torch.logsumexp(self.log_joint(X, mask), dim=1) def posterior(self, X, mask=None): """p(k | x_观测)，形状 (N, K)。""" return torch.softmax(self.log_joint(X, mask), dim=1) def sample(self, n, generator): k = torch.multinomial(self.log_pi.exp(), n, replacement=True, generator=generator) u = torch.rand(n, self.mu.shape[1], generator=generator, dtype=torch.float64) return (u < self.mu[k]).double() # ============================================================================= # 笔画世界：m 个二值隐特征 # ============================================================================= class StrokeWorld: """ 笔画世界：8×8 画布上有 m 条可能的笔画（整行或整列的直线），每条以概率 rho 独立出现； 出现的笔画叠在一起（"或"），再让每个像素以 ε 翻转。 z ∈ {0, 1}^m，p(z) = Π_j rho^{z_j} (1 − rho)^{1 − z_j} t_i(z) = 1 如果有某条出现的笔画经过像素 i，否则 0 p(x | z) = Π_i Bernoulli(x_i; (1 − ε) t_i + ε (1 − t_i)) 【参数】m：笔画条数，1–16（前 ⌈m/2⌉ 条是第 0、1、… 行，后 ⌊m/2⌋ 条是第 0、1、… 列）；rho；eps 【DGM 联系】和像素世界（177 个互斥的配置）不同，这里的隐变量是 m 个**可以同时出现**的特征， 隐状态有 2^m 种。原课 slide 17："假设有 30 个二值隐特征，求 Σ_z p(x, z) 就要加 2^30 项"—— 这里 m = 16 时是 65,536 项，还能算；每多一个特征，计算量翻一倍。 """ def __init__(self, m=16, rho=0.15, eps=0.05): if not 1 <= m <= 16: raise ValueError("m 必须在 1 到 16 之间") self.m, self.rho, self.eps = m, rho, eps n_rows, n_cols = (m + 1) // 2, m // 2 S = torch.zeros(m, 8, 8, dtype=torch.float64) for j in range(n_rows): S[j, j, :] = 1 for j in range(n_cols): S[n_rows + j, :, j] = 1 self.S = S.reshape(m, 64) def enumerate_z(self): """全部 2^m 种隐状态，形状 (2^m, m) 的 0/1 float64 张量。""" codes = torch.arange(2 ** self.m).unsqueeze(1) return ((codes >> torch.arange(self.m)) & 1).double() def _Q(self, Z): """P(x_i = 1 | z)，形状 (n_z, 64)。""" T = (Z @ self.S).clamp(max=1.0) return T * (1 - self.eps) + (1 - T) * self.eps def log_prior(self, Z): k = Z.sum(dim=1) return k * math.log(self.rho) + (self.m - k) * math.log(1 - self.rho) def log_lik(self, X, Z): """log p(x | z)，形状 (N, n_z)。""" Q = self._Q(Z) return X @ Q.log().T + (1 - X) @ (1 - Q).log().T def log_joint(self, X, Z): return self.log_lik(X, Z) + self.log_prior(Z) def log_prob(self, X, chunk=8192): """ 精确的 log p(x) = log Σ_{全部 2^m 个 z} p(x, z)，分块计算（每块 chunk 个 z）以免占用太多内存。 【返回】(N,) 张量 【边界】计算量 ∝ N × 2^m × 64：m = 16 时一张图约 400 万次乘加；m 再大就不现实了。 """ X = torch.as_tensor(X, dtype=torch.float64) Z = self.enumerate_z() out = torch.full((X.shape[0],), -float("inf"), dtype=torch.float64) for s in range(0, Z.shape[0], chunk): out = torch.logaddexp(out, torch.logsumexp(self.log_joint(X, Z[s:s + chunk]), dim=1)) return out def sample(self, n, generator, return_z=False): Z = (torch.rand(n, self.m, generator=generator, dtype=torch.float64) < self.rho).double() Q = self._Q(Z) X = (torch.rand(n, 64, generator=generator, dtype=torch.float64) < Q).double() return (X, Z) if return_z else X # ============================================================================= # 自检 # ============================================================================= def _selftest(): torch.set_num_threads(1) from toy2d import blobs3 g = torch.Generator().manual_seed(0) X = blobs3().sample(500, g) model, hist = fit_gmm(X, 3, torch.Generator().manual_seed(1), reg=0.0) chain = [] for t in range(len(hist) - 1): chain += [hist[t]["ll"], hist[t]["elbo_m"]] chain.append(hist[-1]["ll"]) worst = min(b - a for a, b in zip(chain[:-1], chain[1:])) print(f"EM 单调性：log p(θ_t) ≤ ELBO(q_t, θ_t+1) ≤ log p(θ_t+1) 链条上最小的增量 = {worst:.2e}（应 ≥ −1e−12）") lj = model.log_joint(X) q = torch.softmax(torch.randn(500, 3, generator=g, dtype=torch.float64), dim=1) lhs = torch.logsumexp(lj, dim=1) rhs = elbo_discrete(lj, q) + kl_discrete(q, lj - lhs.unsqueeze(1)) print(f"恒等式 log p(x) = ELBO(q) + KL(q ‖ p(z|x)) 的最大误差 = {(lhs - rhs).abs().max().item():.2e}") sw = StrokeWorld(m=4) Z = sw.enumerate_z() # 2^64 张图无法枚举，所以不直接检查 Σ_x p(x) = 1；给定 z 时像素独立，Σ_x p(x | z) = 1 自动成立， # 只需检查先验 Σ_z p(z) = 1： print(f"笔画世界（m = 4）：Σ_z p(z) = {sw.log_prior(Z).exp().sum().item():.12f}") Xs = sw.sample(5, g) lp_enum = sw.log_prob(Xs) lp_chunk = sw.log_prob(Xs, chunk=3) print(f"分块求和 vs 一次求和的最大差 = {(lp_enum - lp_chunk).abs().max().item():.2e}") if __name__ == "__main__": _selftest()