""" ============================================================================= CS236 · 第 7 讲共享模块：flows.py —— 归一化流的积木 ============================================================================= 【这个模块是干什么的】 归一化流（normalizing flow）= 一个简单的先验 z ~ N(0, I) + 一串**可逆**的变换 x = f(z) = f_M ∘ … ∘ f_1(z)。 由变量替换公式（原课第 7 讲 slide 11、14）： log p_X(x) = log p_Z(f⁻¹(x)) + Σ_m log |det ∂(f_m)⁻¹/∂z_m| 每一层只需要做三件事：正向 z → x、反向 x → z、给出雅可比行列式的对数。这里的约定是： layer.forward(z) → (x, log|det ∂x/∂z|) 采样方向 layer.inverse(x) → (z, log|det ∂z/∂x|) 求密度方向（两个 log det 互为相反数） 积木（第 7 讲）： ElementwiseAffine(D) x = z ⊙ e^s + t：每一维单独缩放、平移（对角雅可比） LinearLU(D, seed) x = W z + b，W = P L U 用 LU 分解参数化：可逆、log|det W| = Σ log|U_ii| Planar(D) x = z + û h(wᵀz + b)（Rezende & Mohamed 2015）：行列式用矩阵行列式引理； 反向没有闭式解，用二分法逐个样本求解 AffineAutoregressive2D() 二维的三角变换：x₁ = z₁ e^{a₁} + m₁，x₂ = z₂ e^{a₂(x₁)} + m₂(x₁)（下三角雅可比） Swap() 交换两个坐标（|det| = 1） Sequential(layers) 把若干层串起来：正向依次、反向倒序，log det 相加 以及 FlowModel(flow, D) 先验 N(0, I) + 流：log_prob(x)、sample(n, g)、sample_with_log_prob(n, g) fit_mle(model, X, ...) 最大似然训练（Adam，余弦学习率） logdet_autograd(fn, z) 用自动求导算出完整的雅可比矩阵再求 log|det|——检查各层的 log det 公式 第 8 讲（code/lecture08/）会在同一个接口上加入耦合层（NICE、RealNVP）、MAF、IAF。 【运行方式】 python3 flows.py # 自检：每种层的 forward/inverse 互逆、log det 和自动求导一致、FlowModel 的密度积分为 1（约 3 秒） 【PyTorch 小知识】 - 本模块不指定数据类型：参数是什么精度，计算就是什么精度。第 7 讲的脚本用 torch.set_default_dtype(torch.float64)， 让二维的数值检查精确到 1e−12 左右。 - torch.linalg.solve_triangular(A, B, upper=…, unitriangular=…)：解三角方程组，O(D²)，比求逆矩阵快又稳。 - nn.Module 的 forward 方法通常通过 model(x) 调用；这里为了把"正向 / 反向"写得一目了然，直接调用 .forward / .inverse。 """ import math import torch import torch.nn as nn import torch.nn.functional as F LOG2PI = math.log(2 * math.pi) def std_normal_log_prob(z): """log N(z; 0, I)，对最后一维求和。""" return -0.5 * (z ** 2 + LOG2PI).sum(-1) class Flow(nn.Module): """所有流层的基类。子类实现 forward(z) → (x, logdet) 和 inverse(x) → (z, logdet)。""" def forward(self, z): raise NotImplementedError def inverse(self, x): raise NotImplementedError class Sequential(Flow): """ 把若干层串起来：x = f_M(…f_1(z))。 【思路】正向依次调用各层，log det 相加（复合函数的雅可比 = 各层雅可比之积，行列式之积 = 行列式之和的 exp）； 反向倒序调用各层的 inverse。原课 slide 14："determinant of product equals product of determinants"。 """ def __init__(self, layers): super().__init__() self.layers = nn.ModuleList(layers) def forward(self, z): total = torch.zeros(z.shape[0], dtype=z.dtype) for layer in self.layers: z, ld = layer.forward(z) total = total + ld return z, total def inverse(self, x): total = torch.zeros(x.shape[0], dtype=x.dtype) for layer in reversed(self.layers): x, ld = layer.inverse(x) total = total + ld return x, total class ElementwiseAffine(Flow): """ x = z ⊙ exp(s) + t（每一维单独缩放、平移）。 【参数】D：维数。参数 s、t 初始化为 0（恒等变换）。 【思路】雅可比是对角矩阵 diag(e^s)，log|det| = Σ s——不依赖 z。NICE 的"缩放层"就是它（第 8 讲）。 """ def __init__(self, D): super().__init__() self.s = nn.Parameter(torch.zeros(D)) self.t = nn.Parameter(torch.zeros(D)) def forward(self, z): return z * self.s.exp() + self.t, self.s.sum().expand(z.shape[0]) def inverse(self, x): return (x - self.t) * (-self.s).exp(), (-self.s.sum()).expand(x.shape[0]) class LinearLU(Flow): """ 可逆线性变换 x = W z + b，W = P L U：P 是固定的置换，L 是对角线为 1 的下三角，U 是上三角，对角线 = sign · e^{log|u|}。 【思路】直接学一个任意矩阵 W，既不保证可逆，log|det W| 也要 O(D³)。LU 参数化让三件事都便宜： log|det W| = Σ log|U_ii|（三角矩阵的行列式 = 对角线之积，slide 17）；求逆 = 解两个三角方程组； 对角线写成 e^{...} 保证非零，W 永远可逆。Glow（Kingma & Dhariwal 2018）的 1×1 可逆卷积就是这样参数化的。 【参数】D：维数；seed：生成随机置换 P 和初始权重的种子 """ def __init__(self, D, seed=0): super().__init__() g = torch.Generator().manual_seed(seed) dtype = torch.get_default_dtype() Q, _ = torch.linalg.qr(torch.randn(D, D, generator=g, dtype=dtype)) # 随机正交矩阵作初值 P, L, U = torch.linalg.lu(Q) self.register_buffer("P", P) self.register_buffer("sign", torch.sign(torch.diagonal(U))) self.L_raw = nn.Parameter(torch.tril(L, -1)) self.U_raw = nn.Parameter(torch.triu(U, 1)) self.log_u = nn.Parameter(torch.log(torch.diagonal(U).abs())) self.b = nn.Parameter(torch.zeros(D, dtype=dtype)) self.register_buffer("eye", torch.eye(D, dtype=dtype)) def _LU(self): L = torch.tril(self.L_raw, -1) + self.eye U = torch.triu(self.U_raw, 1) + torch.diag(self.sign * self.log_u.exp()) return L, U def weight(self): L, U = self._LU() return self.P @ L @ U def forward(self, z): return z @ self.weight().T + self.b, self.log_u.sum().expand(z.shape[0]) def inverse(self, x): L, U = self._LU() y = (x - self.b) @ self.P # P⁻¹ = Pᵀ；按行向量写：y = Pᵀ (x − b) y = torch.linalg.solve_triangular(L, y.T, upper=False, unitriangular=True) z = torch.linalg.solve_triangular(U, y, upper=True).T return z, (-self.log_u.sum()).expand(x.shape[0]) class Planar(Flow): """ 平面流（Rezende & Mohamed 2015）：x = z + û · tanh(wᵀz + b)。 【思路】沿方向 w 看，变换只改变了"沿 û 方向"的位置，改变量取决于 z 在 w 上的投影——像在一个超平面两侧拉伸 / 挤压。 雅可比 J = I + h'(wᵀz + b) · û wᵀ 是"单位阵 + 秩 1 矩阵"，矩阵行列式引理给出 det J = 1 + h'(wᵀz + b) · ûᵀw （O(D)，不用 O(D³)，原课 slide 19） 可逆的条件：h = tanh 时要求 ûᵀw ≥ −1。为此把自由参数 u 修正成 û = u + (m(wᵀu) − wᵀu) · w / ‖w‖²，m(a) = −1 + softplus(a) → 保证 ûᵀw = m(wᵀu) > −1 【反向】没有闭式解。记 s = wᵀz，由 x = z + û tanh(s + b) 得 wᵀx = s + (wᵀû) tanh(s + b)——关于 s 单调递增， 用二分法解出 s，再 z = x − û tanh(s + b)。每个样本都要迭代几十次：这就是平面流只用在"只需要正向"的场合 （例如做变分推断时的 q，第 6 讲）而不适合做密度估计的原因。 【参数】D：维数；enforce_invertible=False 时不做修正（用来演示"不可逆的平面流会出错"） """ def __init__(self, D, enforce_invertible=True): super().__init__() self.u = nn.Parameter(torch.randn(D) * 0.1) self.w = nn.Parameter(torch.randn(D) * 0.1) self.b = nn.Parameter(torch.zeros(1)) self.enforce = enforce_invertible def u_hat(self): if not self.enforce: return self.u wu = self.w @ self.u return self.u + (-1 + F.softplus(wu) - wu) * self.w / (self.w @ self.w) def forward(self, z): u = self.u_hat() a = z @ self.w + self.b x = z + u * torch.tanh(a).unsqueeze(-1) det = 1 + (1 - torch.tanh(a) ** 2) * (u @ self.w) return x, torch.log(det.abs()) def inverse(self, x, iters=80): u = self.u_hat() wx = x @ self.w wu = u @ self.w lo = wx - wu.abs() - 1.0 # |tanh| ≤ 1，所以解一定在 [wx − |wu|, wx + |wu|] 里 hi = wx + wu.abs() + 1.0 for _ in range(iters): # 二分法：g(s) = s + wu·tanh(s + b) − wx 单调递增 mid = 0.5 * (lo + hi) gm = mid + wu * torch.tanh(mid + self.b) - wx lo = torch.where(gm < 0, mid, lo) hi = torch.where(gm < 0, hi, mid) s = 0.5 * (lo + hi) z = x - u * torch.tanh(s + self.b).unsqueeze(-1) _, ld = self.forward(z) return z, -ld class Swap(Flow): """把坐标顺序反过来（二维时就是交换 x₁、x₂）。置换矩阵的 |det| = 1，log det = 0。""" def forward(self, z): return z.flip(-1), torch.zeros(z.shape[0], dtype=z.dtype) def inverse(self, x): return x.flip(-1), torch.zeros(x.shape[0], dtype=x.dtype) class AffineAutoregressive2D(Flow): """ 二维的三角（自回归）仿射变换： x₁ = z₁ · e^{a₁} + m₁ （a₁、m₁ 是常数参数） x₂ = z₂ · e^{a₂(x₁)} + m₂(x₁) （a₂、m₂ 是以 x₁ 为输入的小 MLP） 【思路】x₁ 只依赖 z₁，x₂ 依赖 (z₁, z₂)：雅可比是下三角的，log|det| = a₁ + a₂(x₁)——O(D)（slide 18）。 反向同样容易：先由 x₁ 得 z₁，再算 a₂(x₁)、m₂(x₁) 得 z₂。注意 a₂、m₂ 的输入是 x₁ 而不是 z₁： 这样求密度时（反向）一步算完；采样时（正向）要先算出 x₁ 才能算 x₂——这就是第 8 讲的 MAF。 【参数】hidden：MLP 的宽度 """ def __init__(self, hidden=32): super().__init__() self.a1 = nn.Parameter(torch.zeros(1)) self.m1 = nn.Parameter(torch.zeros(1)) self.net = nn.Sequential(nn.Linear(1, hidden), nn.Tanh(), nn.Linear(hidden, hidden), nn.Tanh(), nn.Linear(hidden, 2)) nn.init.zeros_(self.net[-1].weight) # 初始化成恒等变换附近 nn.init.zeros_(self.net[-1].bias) def _cond(self, x1): out = self.net(x1.unsqueeze(-1)) return out[:, 0], out[:, 1] # m₂(x₁), a₂(x₁) def forward(self, z): x1 = z[:, 0] * self.a1.exp() + self.m1 m2, a2 = self._cond(x1) x2 = z[:, 1] * a2.exp() + m2 return torch.stack([x1, x2], dim=1), self.a1 + a2 def inverse(self, x): z1 = (x[:, 0] - self.m1) * (-self.a1).exp() m2, a2 = self._cond(x[:, 0]) z2 = (x[:, 1] - m2) * (-a2).exp() return torch.stack([z1, z2], dim=1), -(self.a1 + a2) class FlowModel(nn.Module): """ 归一化流模型：z ~ N(0, I_D)，x = flow.forward(z)。 【方法】 log_prob(x) log p(x) = log N(f⁻¹(x)) + log|det ∂f⁻¹/∂x|（slide 13）——精确，不是下界 sample(n, g) z ~ N(0, I)，x = f(z)（slide 16：采样走正向） sample_with_log_prob(n, g) 同时返回样本自己的 log q(x) = log N(z) − log|det ∂f/∂z|：只用正向，不需要求逆 （变分推断、第 8 讲的概率密度蒸馏都靠这一点） encode(x) z = f⁻¹(x)：隐表示，"不需要推断网络"（slide 16） """ def __init__(self, flow, D): super().__init__() self.flow, self.D = flow, D def log_prob(self, x): z, ld = self.flow.inverse(x) return std_normal_log_prob(z) + ld def sample(self, n, generator): z = torch.randn(n, self.D, generator=generator, dtype=torch.get_default_dtype()) return self.flow.forward(z)[0] def sample_with_log_prob(self, n, generator): z = torch.randn(n, self.D, generator=generator, dtype=torch.get_default_dtype()) x, ld = self.flow.forward(z) return x, std_normal_log_prob(z) - ld def encode(self, x): return self.flow.inverse(x)[0] def fit_mle(model, X, steps=3000, batch=256, lr=5e-3, seed=0, X_val=None, log_every=0): """ 最大似然训练：最小化 −(1/B) Σ log p(x)（Adam + 余弦学习率），返回每 log_every 步的 (步数, 训练损失, 验证集平均 −log p)。 【参数】X：(N, D) 训练数据；X_val：可选的验证集；log_every = 0 时只在最后记录一次 【边界】学习率太大时，缩放参数可能一步跳得很远（e^{a} 爆炸），损失出现 nan；遇到这种情况把 lr 调小。 """ g = torch.Generator().manual_seed(seed) opt = torch.optim.Adam(model.parameters(), lr=lr) sched = torch.optim.lr_scheduler.LambdaLR(opt, lambda s: 0.5 * (1 + math.cos(math.pi * s / steps))) hist = [] for step in range(1, steps + 1): x = X[torch.randint(0, X.shape[0], (batch,), generator=g)] loss = -model.log_prob(x).mean() opt.zero_grad() loss.backward() opt.step() sched.step() if (log_every and step % log_every == 0) or step == steps: with torch.no_grad(): v = -model.log_prob(X_val).mean().item() if X_val is not None else float("nan") hist.append((step, loss.item(), v)) return hist def logdet_autograd(fn, z): """ 对每个样本用自动求导算出完整的雅可比矩阵 ∂fn(z)/∂z，返回 log|det|（形状 (N,)）。 【用途】检查各层手写的 log det 公式；代价是 O(N · D³)，只适合小维度。 """ out = [] for i in range(z.shape[0]): J = torch.autograd.functional.jacobian(lambda t: fn(t.unsqueeze(0))[0][0], z[i]) out.append(torch.linalg.slogdet(J)[1]) return torch.stack(out) def _selftest(): torch.set_num_threads(1) torch.set_default_dtype(torch.float64) torch.manual_seed(0) g = torch.Generator().manual_seed(1) layers = {"ElementwiseAffine": ElementwiseAffine(2), "LinearLU": LinearLU(2, seed=3), "Planar": Planar(2), "AffineAutoregressive2D": AffineAutoregressive2D(), "Swap": Swap()} for p in layers["ElementwiseAffine"].parameters(): nn.init.normal_(p, std=0.3) for p in layers["AffineAutoregressive2D"].net.parameters(): # 不要太大：否则 e^{a₂} 会让分布的尾巴极长， nn.init.normal_(p, std=0.2) # 有限的积分网格装不下 with torch.no_grad(): layers["Planar"].u.copy_(torch.tensor([1.5, -0.5])) layers["Planar"].w.copy_(torch.tensor([2.0, 1.0])) z = torch.randn(20, 2, generator=g) print("每种层：反向∘正向 的最大误差、log det（公式 vs 自动求导）的最大误差、正反两个 log det 之和") for name, layer in layers.items(): x, ld = layer.forward(z) z2, ld_inv = layer.inverse(x) ld_auto = logdet_autograd(layer.forward, z) print(f" {name:24s} 互逆 {(z - z2).abs().max().item():.1e} log det {(ld - ld_auto).abs().max().item():.1e} " f"正 + 反 {(ld + ld_inv).abs().max().item():.1e}") model = FlowModel(Sequential(list(layers.values())), 2) t = torch.linspace(-12, 12, 801) pts = torch.cartesian_prod(t, t) with torch.no_grad(): total = (model.log_prob(pts).exp().sum() * (t[1] - t[0]) ** 2).item() with torch.no_grad(): outside = (model.sample(100000, torch.Generator().manual_seed(5)).abs() > 12).any(1).double().mean().item() print(f"五层串起来的 FlowModel：在 [−12, 12]² 上数值积分 ∫p(x)dx = {total:.6f}（样本落在这个范围外的比例 {outside:.5f}）") if __name__ == "__main__": _selftest()