模块 5: 其他值得了解的网络

模块 01 至 04 把输入映射到一个标签、一个数，或序列的下一个元素。本模块的各个模型族则压缩和生成数据、读取图、服从微分方程、在没有标签的情况下学习，或者在不增加计算量的情况下增加参数。每个模型族都由它优化的目标来定义，它的难点大多也源于这个目标。本节给出全景图，外加本模块其余部分所依赖的数学。

每个模型族优化什么

混合专家改变的是网络如何构建，而不是网络为什么而训练。神经算子（第 10 节）与 PINN 并列，但它从求解器的输出中学习，而不是从方程中学习。

构建生成器的三种方式

图 5.1 的上行展示了三种取舍。VAE 是一个显式的潜变量模型，在似然的一个下界上训练。GAN 是一个隐式模型：生成器把噪声变成样本，并且只能通过第二个网络来学习，这个网络试图把样本与数据区分开。扩散模型学习逆转对数据的逐步加噪，每次一小步。

第 3 至 6 节逐一解释表中的每一行。

KL 散度与詹森不等式

从分布 q 到分布 p 的 KL 散度（Kullback–Leibler divergence）为

它永不为负。詹森不等式（Jensen’s inequality）指出，对 \log 这样的凹函数，\E[\log Y] \le \log \E[Y]。在 q 下取 Y = p/q 应用它：

等号成立要求 p/q 在 q > 0 处为常数，对归一化的密度而言，这意味着 q = p。这个散度不对称，所以它不是距离。

对两个一元高斯分布，对数密度之差为

而在 q = \mathcal{N}(\mu_1, s_1^2) 下，\E_q[(z - \mu_1)^2] = s_1^2，\E_q[(z - \mu_2)^2] = s_1^2 + (\mu_1 - \mu_2)^2。因此

相对于标准正态分布（\mu_2 = 0，s_2 = 1），它变为 \tfrac12(\mu^2 + s^2 - \log s^2 - 1)，这正是第 3 节中每个 VAE 都要计算的那一项。

高斯代数

第 5 节用到三个事实。相互独立的 a \sim \mathcal{N}(0, s_a^2) 和 b \sim \mathcal{N}(0, s_b^2) 给出 a + b \sim \mathcal{N}(0, s_a^2 + s_b^2)：方差相加。对 \epsilon \sim \mathcal{N}(0, 1)，有 c\,\epsilon \sim \mathcal{N}(0, c^2)。任何高斯样本都可以写成 \mu + \sigma \epsilon。

蒙特卡洛，以及它求不了的梯度

期望用样本平均来估计，\E_p[f(\mathbf{x})] \approx \frac{1}{S}\sum_{s=1}^{S} f(\mathbf{x}_s)，其中 \mathbf{x}_s \sim p；这一估计是无偏的，标准误差按 1/\sqrt{S} 下降。当分布依赖于正在训练的参数时，麻烦就来了：

参数位于密度中，而不在 f 中。取 f(z) = z^2、q = \mathcal{N}(\mu, 1)，期望为 \mu^2 + 1，梯度为 2\mu；而把一个采样得到的 z^2 对 \mu 求导得 0，因为样本不记录自己从何而来。第 3 节解决这个问题。

有三个先前的工具直接使用，不再重新推导：极大似然及其损失（模块 01 第 5 节）、反向模式自动微分（模块 02 第 4 节），以及主成分分析（模块 01 第 11 节）。

自编码器是作为一个整体来训练的两个网络。编码器把输入映射为一个编码 \mathbf{z} = f_\phi(\mathbf{x}) \in \R^{d_z}，解码器再把编码映射回去，\hat{\mathbf{x}} = g_\theta(\mathbf{z}) \in \R^{d_x}。训练在数据上最小化重构误差：

不需要标签；输入就是它自己的目标。编码位于潜在空间（latent space）中，而让它有价值的是一个约束。欠完备（undercomplete）自编码器满足 d_z < d_x，即有一个瓶颈（bottleneck）：编码装不下一切，所以训练必须决定保留什么，它保留的是在整个数据集上最能降低误差的东西（图 5.2）。一个没有其他约束的过完备（overcomplete）自编码器（d_z \ge d_x）可以学到恒等映射，完美地重构，却对哪些输入是可能出现的一无所知。

线性自编码器就是 PCA

取中心化的数据、线性编码器 \mathbf{z} = \mathbf{W}_e\mathbf{x}（其中 \mathbf{W}_e \in \R^{d_z \times d_x}）、线性解码器 \hat{\mathbf{x}} = \mathbf{W}_d\mathbf{z}，以及平方误差。重构为 \mathbf{M}\mathbf{x}，其中 \mathbf{M} = \mathbf{W}_d\mathbf{W}_e 是秩至多为 d_z 的矩阵，所以每个重构都位于一个 d_z 维子空间中。对固定的子空间，离 \mathbf{x} 最近的点是它的正交投影，所以问题在于选哪个子空间。把到一组标准正交基 \mathbf{U} \in \R^{d_x \times d_z} 上的投影写成 \mathbf{U}\mathbf{U}^\top，协方差写成 \mathbf{C} = \frac{1}{N}\sum_i \mathbf{x}_i\mathbf{x}_i^\top。平均误差为

这是因为投影与残差正交。最小化它，就是最大化保留下来的方差 \operatorname{tr}(\mathbf{U}^\top\mathbf{C}\mathbf{U})，而 \mathbf{C} 的前 d_z 个特征向量做到了这一点（Eckart–Young 定理）。最小误差是 \operatorname{tr}(\mathbf{C}) 减去前 d_z 个特征值，即被舍弃的特征值之和。Baldi 和 Hornik（1989）证明，线性自编码器上的梯度下降没有其他局部极小值，所以它能找到这个子空间。

它找到的是子空间，而不是主成分本身。对任意可逆的 d_z \times d_z 矩阵 \mathbf{A}，(\mathbf{A}\mathbf{W}_e,\ \mathbf{W}_d\mathbf{A}^{-1}) 这一对给出相同的乘积 \mathbf{M} 和相同的误差，所以编码坐标可以是各主成分的任意混合。

非线性的编码器和解码器能沿着弯曲的流形走，而没有哪个平直的子空间能拟合这样的流形。在实验 1 中，在缩放到 [0, 1] 的 8×8 数字上，一个两侧各有一个 128 单元隐藏层的自编码器，在 d_z = 2 时达到每像素 0.037 的测试误差，PCA 为 0.053；在 d_z = 8 时为 0.011 对 0.025。PCA 仍是应当首先拟合的基线：它是精确的，瞬间就能算完，而且当数据接近一个子空间时，它几乎一样好。

去噪

去噪自编码器（denoising autoencoder，Vincent 等人 2008）的输入是一个被破坏的 \tilde{\mathbf{x}}（加上高斯噪声，或遮住部分像素），训练它输出干净的 \mathbf{x}。即使是过完备的网络，也无法靠复制输入来完成这个任务；它必须学到数据位于何处，并把点移回数据那里。

它学到的东西有精确的形式。在平方误差下，最好的去噪器是条件均值 r(\tilde{\mathbf{x}}) = \E[\mathbf{x} \mid \tilde{\mathbf{x}}]。取 \tilde{\mathbf{x}} = \mathbf{x} + \sigma\boldsymbol{\epsilon}，带噪声的密度为 p_\sigma(\tilde{\mathbf{x}}) = \int p(\mathbf{x})\,\mathcal{N}(\tilde{\mathbf{x}}; \mathbf{x}, \sigma^2\mathbf{I})\, d\mathbf{x}。在积分号下求导；高斯函数的梯度等于高斯函数乘以 (\mathbf{x} - \tilde{\mathbf{x}})/\sigma^2：

所以 r(\tilde{\mathbf{x}}) - \tilde{\mathbf{x}} = \sigma^2 \nabla \log p_\sigma(\tilde{\mathbf{x}})，在噪声较小时近似为 \sigma^2 \nabla \log p(\tilde{\mathbf{x}})（Vincent 2011）。去噪器的修正指向对数密度上升的方向。这个梯度，即分数（score），正是第 5 节中扩散模型要学习的东西。

异常检测

基于重构的异常检测（anomaly detection）只在正常运行的数据上训练自编码器，并用重构误差给新输入打分。一个与训练数据中任何东西都不相像的输入，例如一个以新方式表现的传感器通道，或一个超出设计空间的零件，应当重构得很差。真正要紧的决定是阈值。把它设在留出正常数据误差的一个高百分位上，比如第 95 百分位，误报率按构造就约为 5%。检出率无法设定，只能测量，而且只能在你已知的异常上测量，它取决于异常长什么样。

同样的配方也适用于预测残差，而不只是重构（模块 04 第 9 节）。它有三种失效方式。与正常数据相似的异常重构得很好，于是通过了检测，许多 9 就是这样。运行工况的改变，例如新的载荷工况或更换了传感器，会把正常误差推到阈值之上，用误报淹没操作人员。而藏在“正常”训练数据中的异常会被当作正常学进去：检测器的干净程度，取决于被告知是健康的那些数据有多干净。

自编码器的潜在空间有空洞。解码两个编码簇之间的一个点，输出不是任何特定的数字（实验 1 展示了这一点），所以解码器不能用来生成：没有任何东西规定编码位于何处。变分自编码器（VAE；Kingma 和 Welling 2014，Rezende、Mohamed 和 Wierstra 2014 年独立提出了它）把编码变成一个带先验的随机变量，并把编码器和解码器作为一个概率模型来训练。

模型，以及为什么它的似然无法求得

生成过程分两步：抽取一个编码 \mathbf{z} \sim p(\mathbf{z}) = \mathcal{N}(\mathbf{0}, \mathbf{I})，再抽取 \mathbf{x} \sim p_\theta(\mathbf{x} \mid \mathbf{z})，这是一个简单分布，它的参数由解码器网络从 \mathbf{z} 算出。一个数据点的似然为

极大似然（模块 01 第 5 节）需要每个训练点的 \log p_\theta(\mathbf{x}_i)，而这个积分里面有一个神经网络，还要覆盖 d_z 个维度。对从先验抽取的编码求 p_\theta(\mathbf{x} \mid \mathbf{z}) 的平均是无偏的，却毫无希望：几乎每个编码解码出来的都是与 \mathbf{x} 无关的东西。要紧的编码是后验 p_\theta(\mathbf{z} \mid \mathbf{x}) = p_\theta(\mathbf{x} \mid \mathbf{z})p(\mathbf{z}) / p_\theta(\mathbf{x}) 中的编码，而后验需要 p_\theta(\mathbf{x})，正是我们无法计算的那个量。

证据下界的两种推导

引入一个近似后验 q_\phi(\mathbf{z} \mid \mathbf{x})，它可以是任何我们能采样、能求值的密度，并把它用作重要性分布。在积分内乘上再除以它，然后应用詹森不等式（第 1 节）：

这就是证据下界（evidence lower bound，ELBO）。第一项奖励那些能让解码器复现 \mathbf{x} 的编码；第二项按每个编码的分布偏离先验的程度向它收费，以奈特（nats）计。

不等式掩盖了放弃的东西；一个恒等式能把它显示出来。由贝叶斯法则，对每个 \mathbf{z}，有 \log p_\theta(\mathbf{x}) = \log p_\theta(\mathbf{x} \mid \mathbf{z}) + \log p(\mathbf{z}) - \log p_\theta(\mathbf{z} \mid \mathbf{x})。左边不依赖于 \mathbf{z}，所以它等于自身在 q_\phi 下的期望。在里面加上再减去 \log q_\phi：

差距恰好是从近似后验到真实后验的 KL 散度。由于 \log p_\theta(\mathbf{x}) 不依赖于 \phi，在 \phi 上提高 ELBO 只能缩小差距：编码器学会逼近后验。在 \theta 上提高 ELBO，则提高似然，或提高下界的紧度，或两者兼有。训练同时做这两件事。

摊销推断与高斯编码器

经典的变分推断为每个数据点单独拟合一个 q，各自做一次优化。VAE 使用摊销推断（amortised inference）：一个编码器网络对每个 \mathbf{x} 输出 q_\phi(\mathbf{z} \mid \mathbf{x}) = \mathcal{N}\big(\boldsymbol{\mu}_\phi(\mathbf{x}), \operatorname{diag}\boldsymbol{\sigma}^2_\phi(\mathbf{x})\big) 的参数，所以一个新输入只需一次前向传播。代价是编码器未必能对每个点都输出最好的 q，这是在高斯形状之外又一个使下界变松的来源。编码器输出 \log \boldsymbol{\sigma}^2 而不是 \boldsymbol{\sigma}^2，因为网络输出是不受约束的实数，取指数能使它为正。

对角高斯 q 配以标准正态先验时，两个对数密度都是对各维求和，所以 KL 是第 1 节中一元公式之和：

让梯度穿过采样

重构项是 q_\phi 下的期望，而 q_\phi 依赖于编码器的参数：这正是第 1 节留下的问题。它有两种解法。

分数函数（score-function）估计量，即 REINFORCE，利用 \nabla_\phi q_\phi = q_\phi \nabla_\phi \log q_\phi：

它是无偏的，即使对离散的 \mathbf{z} 也适用，但方差很高，因为每个样本都把 f 的大小乘上一个随机方向。对 f(z) = z^2 和 q = \mathcal{N}(\mu, 1)，在 \mu = 1 处，单样本估计 z^2(z - 1) 的均值为 2，即真实梯度，方差为 30。

重参数化技巧（reparameterisation trick）把样本写成参数和一个噪声输入的确定性函数，而噪声的分布不依赖于参数：

梯度能移到期望里面，因为现在期望是对 \boldsymbol{\epsilon} 取的，而 \phi 不触及 \boldsymbol{\epsilon}。逐维看，\partial z_j/\partial \mu_j = 1，\partial z_j/\partial \sigma_j = \epsilon_j。在同一个例子中，估计为 2z，均值为 2，方差为 4，而分数函数估计量为 30。这个技巧要求 \mathbf{z} 连续、f 可微。

解码器的似然设定了兑换率

重构项是一个对数似然，所以必须给解码器一个噪声模型。对 [0, 1] 中的强度，伯努利（Bernoulli）解码器为每个像素输出一个 logit，-\log p_\theta(\mathbf{x} \mid \mathbf{z}) 就是对各像素求和的二元交叉熵。高斯解码器 \mathcal{N}(\hat{\mathbf{x}}, \sigma_x^2\mathbf{I}) 给出

求和的平方误差，就是这个式子取 \sigma_x^2 = 1/2 并丢掉常数。对 [0, 1] 中的像素，这相当于 0.71 的噪声标准差，等于说图像几乎不由它的编码决定。\sigma_x 的值是重构与 KL 之间的兑换率：小的 \sigma_x 让每一单位误差都很昂贵，于是值得为 \mathbf{z} 中的信息付费；大的 \sigma_x 则让 KL 占主导。beta-VAE（Higgins 等人 2017）把这一权衡显式写出，用 \beta 给 KL 加权；\beta = 1 就是 ELBO。

下面的类就是实验 1 中使用的紧凑 VAE，并写明了它的似然：

import torch
import torch.nn as nn
import torch.nn.functional as F


class VAE(nn.Module):
    def __init__(self, d_in, d_z=2, h=128):
        super().__init__()
        self.enc = nn.Sequential(nn.Linear(d_in, h), nn.GELU(), nn.Linear(h, 2 * d_z))
        self.dec = nn.Sequential(nn.Linear(d_z, h), nn.GELU(), nn.Linear(h, d_in))
        self.beta = 1.0  # KL weight; 1 gives the ELBO itself

    def forward(self, x):
        mu, logvar = self.enc(x).chunk(2, dim=-1)
        z = mu + torch.exp(0.5 * logvar) * torch.randn_like(mu)  # reparameterisation
        xhat = self.dec(z)
        # Gaussian decoder with sigma_x^2 = 1/2: -log p(x|z) = ||x - xhat||^2 + const
        recon = F.mse_loss(xhat, x, reduction="sum") / len(x)
        # closed-form KL to N(0, I), summed over latent dimensions, averaged over batch
        kl = -0.5 * (1 + logvar - mu**2 - logvar.exp()).sum(-1).mean()
        return recon + self.beta * kl, xhat


def bernoulli_recon(logits, x):
    """-log p(x|z) for pixels in [0, 1]: decoder outputs are Bernoulli logits."""
    return F.binary_cross_entropy_with_logits(logits, x, reduction="sum") / len(x)

把 recon 换成 bernoulli_recon(xhat, x)，就得到实验 1 的主要结果所用的伯努利解码器。

后验坍塌

后验坍塌（posterior collapse）是这样一种状态：对每个 \mathbf{x}，q_\phi(\mathbf{z} \mid \mathbf{x}) 在部分维度或全部维度上都接近先验，而解码器学会忽略这些维度。它有两种成因。一是 KL 的权重相对于重构项可能太大：\beta > 1，或者隐含的 \sigma_x^2 很大。二是解码器可能不需要 \mathbf{z} 就能很好地建模 \mathbf{x}，文本 VAE 的自回归解码器就是如此（Bowman 等人 2016）。

可以用每个维度的 KL 来诊断它，也可以用活跃单元（active units，Burda 等人 2016）：若 \operatorname{Var}_{\mathbf{x}}\big(\E_{q}[z_j]\big) > 0.01，即它的均值编码随输入变化而移动，则维度 j 是活跃的。常用的修正有：KL 预热（把 \beta 从 0 逐渐升高）、自由比特（free bits；Kingma 等人 2016：每个维度在一个奈特下限之下不施加 KL 惩罚）、\beta \le 1，以及缩放得更好的似然。

实验 1 在测试数字上用 d_z = 8 和伯努利解码器测量了这一点：

在 \beta = 4 时，坍塌是最优解，而不是训练中的偶然结果。\beta = 1 的解在 \beta = 4 的目标下计分，代价为 21.0 + 4 \times 3.6 = 35.4 奈特；坍塌的解代价为 27.2 + 4 \times 0 = 27.2。使用编码能省下 6.2 奈特的重构，却要付出 14.4 奈特的加权 KL。这就是预热救不了它的原因：在实验 1 的“动手试试”第 1 项中，预热到 \beta = 4 仍会坍塌（KL 0.01 奈特，没有活跃单元），而预热到 \beta = 1 则把活跃单元从 6 个提高到 8 个。预热修正的是由优化路径造成的坍塌，而不是目标中内置的坍塌。求和平方误差那一行就是 \sigma_x^2 = 1/2 的效应：放弃重构的代价很低，只有三个维度仍在使用。

模糊的样本，以及 VAE 的用途

VAE 的样本是模糊的。一个编码与许多略有不同的图像都相容，而在高斯或伯努利似然上训练的解码器输出它们的平均，在它们不一致的地方就是平滑的。VAE 的价值反而在于它的潜在空间：平滑，处处接近先验，所以每个点都能解码出看似合理的东西，两个编码之间的直线也是合理的插值。用工程的话说，它是一个潜在设计空间。而一个 VAE 式的自编码器，加上较轻的 KL 惩罚和一个保持清晰度的对抗项，就是潜在扩散（第 6 节）内部的压缩器。

生成对抗网络（GAN；Goodfellow 等人 2014）完全放弃了似然。生成器（generator）G 把噪声 \mathbf{z} \sim p(\mathbf{z}) 映射为样本 G(\mathbf{z})；这些样本服从某个分布 p_g，但它的密度没有公式可写，所以不能用极大似然训练。取而代之的是一个判别器（discriminator）D，它输出其输入为真实数据的概率，D(\mathbf{x}) = \sigma(a(\mathbf{x}))，其中 a 是一个 logit；两者进行一场博弈：

D 是使用交叉熵损失的二元分类器（模块 01 第 6 节），数据标签为 1，样本标签为 0。G 被训练来让这个分类器失败。训练交替进行：在 D 上走一步，再在 G 上走一步（图 5.5）。

这场博弈优化的是什么

把第二个期望写成对 \mathbf{x} = G(\mathbf{z}) 的期望，使两项都成为对 \mathbf{x} 的积分：

对固定的 G，D 可以在每个 \mathbf{x} 处分别选择自己的取值，所以逐点最大化被积函数。记 a = p_{\text{data}}(\mathbf{x})，b = p_g(\mathbf{x})，则 h(y) = a\log y + b\log(1 - y) 有 h'(y) = a/y - b/(1 - y)，它在 y = a/(a + b) 处为零；由于 h 是凹的，这是极大值点。因此

最优判别器估计的是一个密度比：D^*/(1 - D^*) = p_{\text{data}}/p_g。代入它，并记 p_{\text{data}} + p_g = 2m，m 为混合分布：

其中 Jensen–Shannon 散度（Jensen–Shannon divergence，JSD）为 \mathrm{JSD}(p \,\|\, q) = \tfrac12 D_{\KL}(p \,\|\, m) + \tfrac12 D_{\KL}(q \,\|\, m)。它仅当两个分布相等时为零，所以面对一个完美的判别器，生成器最小化的是 JSD，而博弈的值 -\log 4 恰好在 p_g = p_{\text{data}} 时达到。

饱和与非饱和损失

上面的分析假设 D 是最优的，但训练是一串梯度步，而原始的生成器损失给出的梯度步很差。训练早期样本很差，D 轻易地拒绝它们，D(G(\mathbf{z})) = \sigma(a) 接近 0。生成器最小化 \log(1 - \sigma(a))，它对 logit 的导数是 -\sigma(a)：接近 0，恰恰是在生成器最需要信号的时候。非饱和（non-saturating）损失改为让生成器最小化 -\log D(G(\mathbf{z}))，其导数为 -(1 - \sigma(a))，在同一处接近 -1。两种损失都靠骗过 D 来最小化，所以不动点相同；不同的是动力学。

模式坍塌

损失中没有任何东西奖励覆盖整个 p_{\text{data}}。一个把许多 \mathbf{z} 映射到当前 D 所接受的少数几个输出上的生成器，按损失来看表现良好。随后 D 适应过来，学会拒绝这些输出，G 又跳到其他模式上：两者互相追逐，而不是收敛。Metz 等人（2017）在由八个高斯分布组成的一个环上展示了这一点，标准 GAN 在那里一个接一个地访问各个模式。这就是模式坍塌（mode collapse，图 5.6）。损失揭示不了它，查看单个样本也不行，因为每个样本单独看都令人信服。要用多样性度量来诊断它：覆盖了多少已知模式、留出数据点到其最近样本的距离，以及样本相对于数据的精确率和召回率。

不相交的支撑集与 Wasserstein 距离

真实数据（例如图像）位于一个巨大空间中某些低维集合的附近，训练早期生成器的样本也是如此。两个这样的集合通常不重叠。此时 m 在两者各自的支撑集上都等于它的一半，两者的 D_{\KL}(p \,\|\, m) 都是 \log 2，无论两个分布相距多远，JSD 都是 \log 2。一个在每次微小移动下都保持不变的散度，不给生成器任何方向。

Wasserstein-1 距离，即推土机距离（earth mover’s distance），度量概率质量必须移动多远：

其中 \Pi(p, q) 是边缘分布为 p 和 q 的联合分布（耦合）的集合。对耦合取下确界无法直接计算，但 Kantorovich–Rubinstein 对偶把它变成一个对函数的优化：W(p, q) = \sup_{\|f\|_L \le 1} \E_p[f] - \E_q[f]，即对 1-Lipschitz 函数 f 取上确界。Wasserstein GAN（Arjovsky、Chintala 和 Bottou 2017）训练一个网络 f，即评判器（critic），去达到这个上确界，并训练生成器去减小它。难点在于保持 f 的 Lipschitz 性。WGAN 把评判器的权重裁剪到一个小方框内，这样做有效，但限制了评判器。WGAN-GP（Gulrajani 等人 2017）在数据与样本之间的随机插值点 \hat{\mathbf{x}} 处加入惩罚 \lambda\,\E\big[(\|\nabla f(\hat{\mathbf{x}})\| - 1)^2\big]。谱归一化（spectral normalisation，Miyato 等人 2018）把每层的权重除以其最大奇异值，把每层的 Lipschitz 常数限制在 1 以内，是另一种常用的稳定手段。

GAN 的现状

GAN 生成了第一批照片级逼真的图像，至今仍是最快的生成器，每个样本只需一次前向传播。不过，截至 2026 年，扩散模型（第 5 节）已在大多数新的图像生成任务中取代了它们；这一转变大约始于 2021 年，当时 Dhariwal 和 Nichol 报告扩散模型在图像合成上击败了 GAN。扩散模型的胜出有四个原因：它作为一个稳定的回归来训练，只有一个网络和一个损失；它能覆盖各个模式，因为它的目标是一个会惩罚遗漏数据的似然下界；它很容易以文本、类别或图像为条件；它能扩展到大模型和大数据集。对抗损失作为一个附加项保留了下来：在图像和音频编解码器中，在潜在扩散的自编码器中（Rombach 等人用一个基于图像块的对抗项训练它，以保持清晰），以及在把扩散模型蒸馏到只需几步时。

在工程中，GAN 曾被用作探索设计空间的快速采样器，用于数据增强，以及对仿真得到的场做超分辨率。在信任它们的样本之前，先检查覆盖度：只生成常见设计的生成器看起来会非常出色，却会漏掉那些真正要紧的罕见设计。

第 3 节的 VAE 一步到位地学习生成器，从简单的潜变量直接到数据，代价是样本模糊。第 4 节的 GAN 生成的样本清晰，但训练是一场不稳定的博弈。扩散模型（diffusion model）把生成拆成许多小步。破坏数据很容易：加一点高斯噪声，再加一点，直到只剩下噪声。每一小步破坏几乎都是可逆的，而撤销一步是一个去噪问题，也就是一个回归问题。训练一个网络去撤销每一步，然后从纯噪声出发，反复应用它。本节构建这个破坏过程并推导回归损失；第 6 节把训练好的网络变成采样器。

前向过程

固定步数 T 和一个噪声调度（noise schedule）\beta_1, \dots, \beta_T，即一组小的正数。前向过程（forward process）每次加一步噪声：

记 \alpha_t = 1 - \beta_t，一步就是 \mathbf{x}_t = \sqrt{\alpha_t}\,\mathbf{x}_{t-1} + \sqrt{1-\alpha_t}\,\boldsymbol{\epsilon}_t，其中 \boldsymbol{\epsilon}_t \sim \mathcal{N}(\mathbf{0}, \mathbf{I}) 每步重新抽取。收缩因子 \sqrt{1-\beta_t} 是有原因的。若 \mathbf{x}_{t-1} 的某个坐标方差为 1，则 \mathbf{x}_t 的方差为 (1-\beta_t) \cdot 1 + \beta_t = 1：标准化到单位方差的数据始终保持单位方差，过程收敛到 \mathcal{N}(\mathbf{0}, \mathbf{I})，而不像单纯叠加噪声那样漂向越来越大的值。

这里没有任何东西是学出来的。前向过程没有参数；调度是人工选定的。

闭式解。训练需要随机 t 下的 \mathbf{x}_t，为此跑 t 步会很浪费。其实也没有必要。利用第 1 节的高斯代数，把两步复合起来：

括号里是两个相互独立的零均值高斯变量之和，因此它也是高斯的，每个坐标的方差为 \alpha_2(1-\alpha_1) + (1-\alpha_2) = 1 - \alpha_1\alpha_2。记 \bar\alpha_t = \prod_{s=1}^{t}\alpha_s。两步给出 \mathbf{x}_2 = \sqrt{\bar\alpha_2}\,\mathbf{x}_0 + \sqrt{1-\bar\alpha_2}\,\boldsymbol{\epsilon}，其中只有一个 \boldsymbol{\epsilon} \sim \mathcal{N}(\mathbf{0}, \mathbf{I})。同一个恒等式 \alpha_t(1-\bar\alpha_{t-1}) + (1-\alpha_t) = 1 - \bar\alpha_t 把结果从 t-1 推到 t，由归纳法得

练习 6 要求你写出这个归纳。式 (5.4) 是训练廉价的关键：任意噪声水平下的训练输入只需一行代码，不必对 t 循环。信号按 \sqrt{\bar\alpha_t} 缩放，噪声按 \sqrt{1-\bar\alpha_t} 缩放，两个缩放因子的平方和为 1。

调度

调度决定信号消退的快慢。一个有用的单一指标是信噪比（signal-to-noise ratio）

即对单位方差数据而言 (5.4) 中信号方差与噪声方差之比，常以分贝表示，10\log_{10}\mathrm{SNR}。有一条要求没有商量余地：\bar\alpha_T 必须接近 0，使 \mathbf{x}_T 接近 \mathcal{N}(\mathbf{0}, \mathbf{I})。采样从 \mathcal{N}(\mathbf{0}, \mathbf{I}) 开始；如果训练中网络从未见过这样的输入，采样的第一步就是迈向未知。

Ho、Jain 和 Abbeel（2020）使用线性调度（linear schedule）：在 T = 1000 步内，\beta_t 从 10^{-4} 均匀增大到 0.02，最终 \bar\alpha_T = 4.0 \times 10^{-5}。Nichol 和 Dhariwal（2021）注意到它过早地破坏了信息，提出了余弦调度（cosine schedule），直接定义 \bar\alpha_t：

并取 \beta_t = 1 - \bar\alpha_t/\bar\alpha_{t-1}，在 0.999 处截断，否则 f(T) = 0 会使最后的 \beta_T 等于 1。小偏移量 s 使 \beta_1 不至于小到可以忽略。

如果已知 \mathbf{x}_0，反向的一步

生成需要反方向的 q(\mathbf{x}_{t-1} \mid \mathbf{x}_t)。这个分布是难以处理的：它依赖于整个数据分布，因为许多干净的点都可能导向 \mathbf{x}_t。若再以干净的点为条件，它就是高斯的。由贝叶斯公式，q(\mathbf{x}_{t-1} \mid \mathbf{x}_t, \mathbf{x}_0) \propto q(\mathbf{x}_t \mid \mathbf{x}_{t-1})\, q(\mathbf{x}_{t-1} \mid \mathbf{x}_0)，这是关于 \mathbf{x}_{t-1} 的两个高斯的乘积。对 \mathbf{x}_{t-1} 配方（代数推导是常规的，这里略过），两者的精度相加，\alpha_t/\beta_t + 1/(1-\bar\alpha_{t-1}) = (1-\bar\alpha_t)/\big(\beta_t(1-\bar\alpha_{t-1})\big)，并得到

均值把点的来处和它现在的位置混合在一起；缺失的 \mathbf{x}_0 必须由模型提供。

从证据下界到预测噪声

把 \mathbf{x}_1, \dots, \mathbf{x}_T 当作潜变量，把反向模型 p_\theta(\mathbf{x}_{t-1} \mid \mathbf{x}_t) = \mathcal{N}\big(\boldsymbol{\mu}_\theta(\mathbf{x}_t, t),\ \sigma_t^2 \mathbf{I}\big)（\sigma_t^2 固定）当作解码器。以 q(\mathbf{x}_{1:T} \mid \mathbf{x}_0) 为编码器，第 3 节的证据下界可以重新整理为每步一个 KL 项（这需要一页记账式的推导，对每个前向步应用贝叶斯公式，见 Ho 等人的扩展推导）：

L_T 不含参数。每个 L_{t-1} 都是两个高斯之间的 KL 散度；当 p_\theta 的方差固定时，它化为均值之间的平方距离，L_{t-1} = \frac{1}{2\sigma_t^2}\|\tilde{\boldsymbol{\mu}}_t - \boldsymbol{\mu}_\theta\|^2 + C，其中 C 与 \theta 无关。

现在从 (5.5) 中消去 \mathbf{x}_0。由 (5.4)，\mathbf{x}_0 = (\mathbf{x}_t - \sqrt{1-\bar\alpha_t}\,\boldsymbol{\epsilon})/\sqrt{\bar\alpha_t}。代入后，\mathbf{x}_t 的系数变为 \big(\beta_t + \alpha_t(1-\bar\alpha_{t-1})\big)/\big((1-\bar\alpha_t)\sqrt{\alpha_t}\big) = 1/\sqrt{\alpha_t}，于是

网络能看到 \mathbf{x}_t，所以唯一的未知量是 \boldsymbol{\epsilon}。用同样的方式参数化模型均值，用网络 \boldsymbol{\epsilon}_\theta(\mathbf{x}_t, t) 代替 \boldsymbol{\epsilon}。两者相减时 \mathbf{x}_t 项相互抵消，得到

界中的每一项都是加权的噪声预测误差。Ho 等人去掉了权重，并对 t 均匀采样，得到实践中使用的损失：

被去掉的权重并非细微差别。取 \sigma_t^2 = \beta_t 时权重为 \beta_t/\big(2\alpha_t(1-\bar\alpha_t)\big)；对线性调度，它在 t = 1 处为 0.50（此处 1-\bar\alpha_1 = \beta_1），在 t = 100 处为 0.010，在 t = 500 处为 0.0055，在 t = 1000 处为 0.010。界把最小噪声水平上的权重放大了五十到九十倍，而在那里去噪最容易，对样本的观感影响也最小。相对于证据下界，\mathcal{L}_{\text{simple}} 把重心移到了更难、噪声更大的步上；Ho 等人发现它给出的样本更好，代价是它不再是似然的一个界。图 5.8 展示了一个训练步。

写成代码，(5.6) 的一个 batch 只需六行：

def diffusion_loss(eps_model, x0, alpha_bar):
    """L_simple for one batch; alpha_bar has length T + 1 with alpha_bar[0] = 1."""
    T = len(alpha_bar) - 1
    t = torch.randint(1, T + 1, (x0.shape[0],))         # one noise level per example
    eps = torch.randn_like(x0)
    ab = alpha_bar[t].view(-1, *[1] * (x0.dim() - 1))   # broadcast over feature dims
    xt = ab.sqrt() * x0 + (1 - ab).sqrt() * eps         # equation (5.4): no loop over t
    return ((eps - eps_model(xt, t)) ** 2).mean()

同一个损失，从分数估计的角度看

通往 (5.6) 的第二条路径能说明网络学到了什么。对数密度关于其自变量的梯度 \nabla_{\mathbf{x}} \log p(\mathbf{x}) 称为分数函数（score function）：一个指向密度更高处的向量场。对一步加噪，

在加噪样本上把网络回归到这个条件分数，在最优处恢复的是加噪边缘分布 p_t(\mathbf{x}_t) 的分数；这就是去噪分数匹配（denoising score matching，Vincent，2011），Song 和 Ermon（2019）把它变成了生成器。预测 \boldsymbol{\epsilon} 是同一个回归，只差一个固定的比例，所以训练好的噪声预测器就是一个分数估计器：\mathbf{s}_\theta(\mathbf{x}_t, t) = -\boldsymbol{\epsilon}_\theta(\mathbf{x}_t, t)/\sqrt{1-\bar\alpha_t}。

分数还说明了网络对干净数据的看法。Tweedie 公式（Tweedie’s formula）由加噪边缘分布的分数给出 \mathbf{x}_0 的后验均值：

因此 \hat{\mathbf{x}}_0，即经 (5.4) 反解噪声预测得到的量，是一个后验均值：对所有可能产生 \mathbf{x}_t 的干净点取平均。在高噪声下，许多截然不同的点都可能产生它，它们的平均是一团模糊。这就是为什么从纯噪声出发的一步去噪生成不了任何东西，也是为什么第 6 节的采样器要走许多小步。

它为什么好训练，网络长什么样

式 (5.6) 是对一个分布固定的目标做普通回归：噪声由训练循环抽取，不会随网络的学习而移动。没有对手，也没有博弈，所以损失虽然在 batch 之间有波动，却稳定下降，并且在整个训练过程中含义不变。这比什么都更能解释扩散模型为何取代了 GAN（第 4 节）。

预测 \boldsymbol{\epsilon} 只是携带相同信息的三种选择之一。网络也可以直接预测 \mathbf{x}_0，或预测速度（velocity）\mathbf{v} = \sqrt{\bar\alpha_t}\,\boldsymbol{\epsilon} - \sqrt{1-\bar\alpha_t}\,\mathbf{x}_0（Salimans 和 Ho，2022）。三者都可以通过 (5.4) 互相转换；区别在于对它们取平方误差时各噪声水平的权重不同，因而网络拟合得最好的步也不同。

对图像而言，\boldsymbol{\epsilon}_\theta 是一个输出形状与输入相同的 U-Net（模块 03 第 12 节），或一个作用于图像块的 Transformer；在实验 2 中，它是一个作用于两个坐标的小 MLP。时间步通过 t 的正弦嵌入输入网络，构造方式与模块 06 第 6 节的位置编码相同，这样同一个网络可以在每个噪声水平上表现不同。

训练好的 \boldsymbol{\epsilon}_\theta 预测的是噪声。本节把它变成生成器，用条件引导它，让它负担得起，并计算它的代价。

祖先采样

第 5 节的反向模型均值为 \boldsymbol{\mu}_\theta = \frac{1}{\sqrt{\alpha_t}}\big(\mathbf{x}_t - \frac{\beta_t}{\sqrt{1-\bar\alpha_t}}\boldsymbol{\epsilon}_\theta(\mathbf{x}_t, t)\big)，方差 \sigma_t^2 固定。从 t = T 一直采样到 1，就是祖先采样（ancestral sampling，即 Ho 等人的采样算法）：

最后一步取 \mathbf{z} = \mathbf{0}，\sigma_t^2 = \beta_t 或 \tilde\beta_t（两者都可行）。每一步减去预测噪声的一部分，重新缩放，再加一点新的噪声。一个样本要做 T 次网络求值。

@torch.no_grad()
def ddpm_sample(eps_model, shape, alpha_bar):
    """Ancestral sampling with sigma_t^2 = beta_t."""
    T = len(alpha_bar) - 1
    x = torch.randn(shape)
    for t in range(T, 0, -1):
        alpha_t = alpha_bar[t] / alpha_bar[t - 1]
        beta_t = 1 - alpha_t
        eps = eps_model(x, torch.full((shape[0],), t))
        x = (x - beta_t / (1 - alpha_bar[t]).sqrt() * eps) / alpha_t.sqrt()
        if t > 1:                                        # no noise on the last step
            x = x + beta_t.sqrt() * torch.randn_like(x)
    return x

图 5.9 显示了功夫花在哪里。在实验 2 的 200 步中，前 100 步里距离几乎不动（从 0.246 到 0.212；纯高斯抽样给出 0.242）；双月在最后 50 步才成形，最终为 0.022，而新抽取的数据为 0.013。

第一步。更新要除以 \sqrt{\alpha_t}，这会放大 \boldsymbol{\epsilon}_\theta 中的任何误差；通常 \alpha_t 接近 1。余弦调度在 \beta_T = 0.999 处的截断使第一步的放大因子为 1/\sqrt{0.001} = 31.6。因此实验 2 在 T = 200 时把 \beta_t 限制在 0.5 以内：\bar\alpha_T 仍为 6.8 \times 10^{-5}，放大因子为 1/\sqrt{0.5} = 1.41。在实验 2 的数据上，0.999 的截断未必有害（用它重新训练的实验副本在 w = 7 时给出 0.019 的类精确率距离），但一次原型运行用它时发散了，而这个上限没有任何代价。另一种等价的修正保留原调度：计算 \hat{\mathbf{x}}_0，把它截断到数据范围内，然后走到 (5.5) 的后验均值 \tilde{\boldsymbol{\mu}}_t(\mathbf{x}_t, \hat{\mathbf{x}}_0)。

更少的步数：DDIM

Song、Meng 和 Ermon（2021）注意到，训练损失只约束边缘分布 q(\mathbf{x}_t \mid \mathbf{x}_0)，而不约束逐步的过程，因此同一个训练好的网络可以服务于其他采样器。DDIM 走一个确定性的步：先估计干净的点，再用预测的噪声而不是新噪声把它重新加噪到更低的水平 t' < t，

所用的时间是一个带步幅的子序列，比如 200、190、……、0。实验 2 的扫描测出了这能换来什么：类精确率距离在 200 步时为 0.022，20 步时为 0.025，5 步时为 0.046，一步时为 1.71。从纯噪声出发的单步返回的是 \hat{\mathbf{x}}_0，即给定噪声时后验均值的估计。即使网络完美，这也是数据均值附近的一团模糊，而训练所得网络的小误差还会被放大 \sqrt{1-\bar\alpha_T}/\sqrt{\bar\alpha_T} = 121 倍。

第 5 节的高斯例子在没有任何训练误差的情况下展示了这一机制。用它的精确噪声预测器（线性调度，T = 1000，20,000 个样本），祖先采样给出标准差 0.50，正如应有的那样；DDIM 在 50、10 和 3 步时分别给出 0.47、0.37 和 0.16，一步时为 0.002，每个样本都落在后验均值 2.00 上。步数少时，首先丢失的是分布的离散程度。在扩散探索器中，分别用 10 步和 3 步在环上运行 DDIM：余弦面板退化得更慢，而在 3 步时，线性面板把点落在了簇与簇之间。

条件与引导

要生成所要求的东西，就给网络一个条件 \mathbf{c}：\boldsymbol{\epsilon}_\theta(\mathbf{x}_t, t, \mathbf{c})，其中 \mathbf{c} 可以是一个类别、一个文本嵌入、一张低分辨率图像或一个实测的边界条件。除了多一个输入，训练不变。条件模型常常只是松散地遵循它的条件；引导会加强这一点。分类器引导（classifier guidance，Dhariwal 和 Nichol，2021）在每一步把一个在加噪输入上训练的分类器的梯度 \nabla_{\mathbf{x}_t} \log p(\mathbf{c} \mid \mathbf{x}_t) 加到分数上。

无分类器引导（classifier-free guidance，Ho 和 Salimans，2022）不需要分类器。只训练一个网络，以概率 p_{\text{uncond}}（0.1 到 0.2；实验 2 中为 0.2）把 \mathbf{c} 替换为空 token \varnothing，使它同时学会条件预测和无条件预测。采样时用

w = 0 是无条件模型，w = 1 是条件模型，w > 1 则越过它向外外推。Ho 和 Salimans 把缩放写作 (1 + w)，所以他们的 w 等于本模块的 w 减 1；跨论文比较数字之前，先核对约定。

w > 1 意味着什么，可以从分数的角度得出。由于 \mathbf{s} = -\boldsymbol{\epsilon}/\sqrt{1-\bar\alpha_t}，同样的组合对分数也成立；再由贝叶斯公式 \nabla \log p(\mathbf{x} \mid \mathbf{c}) = \nabla \log p(\mathbf{x}) + \nabla \log p(\mathbf{c} \mid \mathbf{x})，得到

引导从条件分布中采样，但把分布向一个隐式分类器有把握地标为 \mathbf{c} 的那些点倾斜。现在每一步要做两次网络求值。

实验 2 在双月数据上请求 1,000 个类别 0 的样本。w = 0 时，50% 落在所请求的那个月牙上；w = 1、3 和 7 时为 100%。类召回率距离（从类别 0 的数据到样本）在 w = 1、3、7 时依次增长为 0.021、0.029、0.043；类精确率距离在 w = 0 到 3 时为 0.023、0.014、0.013，在 w = 7 时因样本冲过头而变差到 0.020。保真度是用多样性换来的，而 w 超过某个值后，保真度也会丢失。

潜在扩散

像素空间的扩散要在每个像素上多次运行一个大网络。潜在扩散（latent diffusion，Rombach 等人，2022）先训练一个自编码器（像 VAE 那样带一个小的 KL 项，并用感知损失和对抗损失保持清晰），把图像压缩成小得多的潜变量。扩散在潜变量上运行；解码器只在最后运行一次。文本条件通过交叉注意力（cross-attention）进入去噪器（从潜变量的各个位置注意文本的 token 嵌入，即模块 06 的机制）。

采样的代价

一个样本的代价 = 步数 × 每步的网络求值次数 × 一次求值的代价。

蒸馏（distillation）训练一个快速的学生去匹配一个慢速的教师。渐进式蒸馏（progressive distillation，Salimans 和 Ho，2022）反复训练学生用一步完成教师的两步，每一轮把步数减半；一致性模型（consistency models，Song 等人，2023）训练一个网络，把采样路径上的任意一点直接映射到路径的终点。两者都能做到 1 到 4 步，在最少的步数下会损失一些质量。模块 10 会在部署服务的场景下再谈多次顺序求值的代价。

工程用途

有三类用途适合这一家族：以需求（一个载荷工况、一个包络）为条件生成候选几何形状；以已观测的传感器通道为条件补全缺失的通道，给出一组合理取值的分布；以及对粗糙的仿真场做超分辨率。样本是一个提议，而不是一个结果。每个生成的设计仍然要经过求解器和常规检查：模型复现的是其训练数据的统计特性，对它没见过的物理一无所知。

到目前为止，每个网络都假定了一种规则的结构：定长的向量、像素网格（模块 03）、序列（模块 04）。许多工程数据并没有这样的结构。分子是由化学键连接的原子，有限元网格是由单元连接的节点，电路是由网络连接的元件，故障树是与其所馈入的门相连的事件。结构是一张图，它因样本而异，而且承载着信息。

作为数据的图

图 \mathcal{G} = (\mathcal{V}, \mathcal{E}) 有 n 个节点和一组边。它的邻接矩阵（adjacency matrix）\mathbf{A} \in \{0, 1\}^{n \times n} 在节点 i 与 j 相连时 A_{ij} = 1；对无向图它是对称的。度矩阵 \mathbf{D} 是对角矩阵，D_{ii} = \sum_j A_{ij} 即 i 的邻居数，\mathcal{N}(v) 表示 v 的邻居集合。每个节点带有一个特征向量，按行堆叠成 \mathbf{X} \in \R^{n \times d}；边也可以带特征 \mathbf{e}_{uv}（化学键类型、相对位置）。工程领域提供了大量的图：分子、网格、电路、故障树、SysML 模块图，以及用目标结构化表示法（Goal Structuring Notation，GSN）写成的安全论证，其中的论点、策略和证据都是节点。

约束：编号是任意的

图本身并不规定哪个节点是 1 号节点。用置换矩阵 \mathbf{P} 给节点重新编号，会把 \mathbf{A} 变成 \mathbf{P}\mathbf{A}\mathbf{P}^\top，把 \mathbf{X} 变成 \mathbf{P}\mathbf{X}，描述的仍是同一张图。因此，一个为每个节点输出一个向量的层必须是置换等变的（permutation-equivariant），

这样给输入重新编号只会给输出重新编号，别的什么都不变。整张图的输出必须是置换不变的，g(\mathbf{P}\mathbf{A}\mathbf{P}^\top, \mathbf{P}\mathbf{X}) = g(\mathbf{A}, \mathbf{X})，对节点求和、求均值或取最大值都能做到这一点。作用于展平邻接矩阵的 MLP 在两方面都失败：同一张图的 n! 种编号对它来说是不同的输入，而且它根本无法接受不同大小的图。

消息传递

出路是：用同一个函数计算每个节点的更新，这个函数的输入是节点自身的状态和其邻居状态的无序集合。Gilmer 等人（2017）写出了消息传递（message passing）层的一般形式：

其中聚合函数是求和、均值或最大值，任何与顺序无关的函数都可以。最简单的实例是

权重由所有节点共享，就像卷积在各个位置上共享卷积核（模块 03 第 2 节）；图就像一个邻域大小不一、且没有顺序的网格。一层让节点看到它的邻居；L 层给它一个 L 跳的感受野（图 5.11）。

为什么要归一化

如果直接求和，度为 50 的节点收到的消息是叶节点的五十倍，于是节点特征的尺度取决于它连接得有多紧密，而不是它是什么。堆叠多层会让情况更糟：L 个未归一化的层把特征乘以 \mathbf{A}^L，而 \mathbf{A} 的最大特征值至少等于平均度。在消息传递探索器（第 8 节）的 12 节点冷却系统故障树上，从独热特征出发，带自环的未归一化传播在 8 步之后，特征中的最大值达到 1.22 \times 10^4。这么大的特征会使激活函数饱和，破坏优化。

从消息传递到 GCN

图卷积网络（graph convolutional network，GCN；Kipf 和 Welling，2017）做了三个选择。只用一个权重矩阵，\mathbf{W}_{\text{self}} = \mathbf{W}_{\text{nbr}} = \mathbf{W}，于是节点自身的状态只是又一条消息；加入自环，\tilde{\mathbf{A}} = \mathbf{A} + \mathbf{I}，相应的度 \tilde d_i = d_i + 1 放在 \tilde{\mathbf{D}} 中；并做对称归一化。以节点特征为行，这一层为

每条消息都除以两个端点的度的平方根，所以枢纽节点发出的消息在每个接收者那里分量更轻，枢纽节点自身的求和也被缩小。

为什么这能保持尺度不变：令 \mathbf{u} = \tilde{\mathbf{D}}^{1/2}\mathbf{1}，即由 \sqrt{\tilde d_i} 组成的向量。那么 \hat{\mathbf{A}}\mathbf{u} = \tilde{\mathbf{D}}^{-1/2}\tilde{\mathbf{A}}\mathbf{1} = \tilde{\mathbf{D}}^{-1/2}\tilde{\mathbf{d}} = \tilde{\mathbf{D}}^{1/2}\mathbf{1} = \mathbf{u}，因为 \tilde{\mathbf{A}} 的行和就是各节点的度。所以 1 是一个特征值。它也是最大的：\hat{\mathbf{A}} = \tilde{\mathbf{D}}^{1/2}(\tilde{\mathbf{D}}^{-1}\tilde{\mathbf{A}})\tilde{\mathbf{D}}^{-1/2} 与 \tilde{\mathbf{D}}^{-1}\tilde{\mathbf{A}} 有相同的特征值，后者每行非负且行和为 1，这样的矩阵没有模大于 1 的特征值。反复应用既不会爆炸，也不会消失。另一种选择是随机游走归一化（random-walk normalisation）\tilde{\mathbf{D}}^{-1}\tilde{\mathbf{A}}：每个节点取其邻域上的均值。它的特征值相同，但不对称。这一层是置换等变的，因为重新编号给出 \hat{\mathbf{A}} \to \mathbf{P}\hat{\mathbf{A}}\mathbf{P}^\top，且利用 \mathbf{P}^\top\mathbf{P} = \mathbf{I} 有 \mathbf{P}\hat{\mathbf{A}}\mathbf{P}^\top\mathbf{P}\mathbf{H}\mathbf{W} = \mathbf{P}\hat{\mathbf{A}}\mathbf{H}\mathbf{W}。

Kipf 和 Welling 是从另一个方向得到 (5.7) 的。谱图理论通过图拉普拉斯矩阵的特征向量定义图上的卷积，代价很高。他们把谱滤波器在拉普拉斯矩阵上做一阶近似，把它的两个系数合并为一个，得到传播矩阵 \mathbf{I} + \mathbf{D}^{-1/2}\mathbf{A}\mathbf{D}^{-1/2}，它的特征值可达 2，所以反复使用并不稳定。他们的“重归一化技巧”（renormalisation trick）把它换成 \tilde{\mathbf{D}}^{-1/2}\tilde{\mathbf{A}}\tilde{\mathbf{D}}^{-1/2}，也就是 (5.7)。本模块的内容都不需要谱的观点；消息传递的推导得到的是同一个层。

计算

在稠密矩阵上，这一层只需一行：

def gcn_layer(A_hat, H, W):
    """A_hat: normalised adjacency with self-loops (n x n); H: (n x d); W: (d x d')."""
    return torch.relu(A_hat @ H @ W)

稠密的 \hat{\mathbf{A}} 需要 O(n^2) 的内存和时间，大部分花在乘零上。真实的图是稀疏的，所以实验 3 把 \hat{\mathbf{A}} 存成三元组 (i, j, \hat A_{ij})，包含每条边的两个方向以及自环，并用一次 scatter-add 计算 \text{out}_i = \sum_j \hat A_{ij}\mathbf{H}_j：

def normalised_edges(edges, n):
    """A-hat as (i, j, weight) triples: both directions of every edge plus self-loops."""
    src = [a for a, b in edges] + [b for a, b in edges] + list(range(n))
    dst = [b for a, b in edges] + [a for a, b in edges] + list(range(n))
    i, j = torch.tensor(dst), torch.tensor(src)       # message travels j -> i
    deg = torch.zeros(n).index_add_(0, i, torch.ones(len(i)))  # d-tilde
    w = (deg[i] * deg[j]).rsqrt()                    # 1 / sqrt(d_i d_j)
    return i, j, w


def propagate(H, i, j, w):
    """out_i = sum over j of A-hat_ij H_j, one multiply-add per stored entry."""
    return torch.zeros_like(H).index_add_(0, i, w[:, None] * H[j])

于是一层就是 torch.relu(propagate(H, i, j, w) @ W)，代价为 O(|\mathcal{E}|\,d + n\,d\,d')：每条边、每个特征一次乘加，再加上与 \mathbf{W} 的稠密乘积。

作为图的故障树

故障树（fault tree）把一个不希望发生的顶事件分解为各种原因。它的节点是事件和门：或门的输出事件在任一输入发生时发生，与门的输出事件只在所有输入都发生时发生，叶节点是基本事件（泵失效、阀门卡死）。边把每个输入连到它的门上。节点特征取独热编码的类型：[OR, AND, basic]，即（或门，与门，基本事件）。单点故障（single point of failure）是这样一个基本事件：它单独失效就会导致顶事件，当且仅当它通往顶事件的路径上每个门都是或门时成立。判断这一点需要几跳之外的信息，这正是它被选为实验 3 任务的原因。

直推式学习与归纳式学习

GCN 论文在一张大图的节点上训练和测试：部分节点有标签，其余节点待预测。这是直推式（transductive）学习，测试节点的特征在训练时是可见的，只是没有标签。工程模型通常是归纳式（inductive）的：网络在一些故障树上训练，然后应用到新的故障树上。评估也要用同样的方式，按图而不是按节点划分，这是模块 01 第 10 节防止泄漏的规则在图上的版本；实验 3 在 200 棵树上训练，在另外 100 棵上测试。

GCN 只按度来给邻居加权，只能做浅层，忽略方向，而且无法区分某些图。

图注意力

图注意力网络（graph attention network，GAT；Veličković 等人，2018）让特征来决定每个邻居的分量：给每一对打分，在邻域（包括自身）上用 softmax 归一化，再取加权和：

其中 \| 表示拼接，\mathbf{a} 是一个学习得到的向量。多个头并行运行；隐藏层把各头的输出拼接起来，最后一层则取平均。Transformer 层（模块 06）就是在其 token 构成的完全图上做注意力，再加上位置信息；GAT 则是限制在边上的注意力。

过平滑

对邻域求平均的次数足够多，一切看起来就都一样了。\hat{\mathbf{A}} 是对称的，所以 \hat{\mathbf{A}} = \sum_i \lambda_i \mathbf{u}_i\mathbf{u}_i^\top，其中特征向量是标准正交的，于是

第 7 节证明了 \lambda_1 = 1，对应的特征向量为 \mathbf{u}_1 \propto \tilde{\mathbf{D}}^{1/2}\mathbf{1}。对带自环的连通图，其余每个特征值都满足 |\lambda_i| < 1（Perron–Frobenius 定理；自环排除了 -1）。所以除第一项外每一项都会衰减，最慢的按 |\lambda_2|^k 衰减，并且

极限的第 i 行是 (\mathbf{u}_1)_i\,(\mathbf{u}_1^\top\mathbf{X})，即同一个公共向量乘以 \sqrt{\tilde d_i}。所有节点都指向同一个方向，只有度能把它们区分开。这就是过平滑（over-smoothing，Li、Han 和 Wu，2018）；传播之间的权重和非线性改变的是细节，而不是这个趋势。

没有自环时，树是二部图（bipartite；相邻层交替构成两类，每条边都连接这两类），\mathbf{D}^{-1/2}\mathbf{A}\mathbf{D}^{-1/2} 有特征值 -1，它对应的项每一步都变号：特征来回振荡而不收敛。在消息传递探索器中，k = 0 到 4 时余弦读数为 0.455、0.384、0.762、0.661、0.832，最终稳定在 0.976 附近，永远到不了 1。

实践中的深度

两到四层是典型的。在实验 3 中，朴素的 GCN 在四层以内提升很快（测试准确率一层时为 0.810，三层时为 0.897，四层时为 0.938），在八层时达到峰值（0.959）；到 12 和 16 层时，它们对每个事件都预测多数类“不是单点故障”：0.781。深的朴素堆叠既会平滑其特征，又难以训练。残差更新 \mathbf{H} \leftarrow \mathbf{H} + \phi(\hat{\mathbf{A}}\mathbf{H}\mathbf{W}) 让 16 层恢复到 0.948。归一化层和跳跃知识连接（jumping-knowledge connection，读出能看到每一层）也有帮助。

第二个极限是过度挤压（over-squashing，Alon 和 Yahav，2021）：r 跳以内的节点数可以随 r 指数增长，它们的信息必须经过少数几条边挤进一个固定宽度的向量，所以即使深度已经够得着，长程依赖也会受损。

消息传递分不清什么

Weisfeiler–Lehman 检验（Weisfeiler–Lehman test，1-WL）通过颜色细化来比较图：从一种颜色开始，然后反复根据每个节点自身的颜色及其邻居颜色的多重集给它重新着色；只要颜色计数出现差异，两张图就不同。Xu 等人（2019）证明，从相同的特征出发，任何消息传递 GNN 都无法区分 1-WL 区分不了的图，而先求和再应用 MLP 的图同构网络（graph isomorphism network，GIN）能达到这个上限。均值或最大值丢掉了每条消息由多少个邻居发出的信息；求和则保留了它。

修正办法是用结构特征（度、经过某节点的环的数目）、由图计算出的位置编码或随机的节点标识符来打破对称性。

有向边与带类型的边

故障树是有向的，门也有类型。对称的 \hat{\mathbf{A}} 把来自门的消息与来自输入的消息同等对待，而在两跳之外还会混入兄弟节点。关系图卷积网络（relational GCN，Schlichtkrull 等人，2018）给每种边类型和方向各自的权重矩阵：\mathbf{h}_v' = \phi\big(\mathbf{W}_0\mathbf{h}_v + \sum_r \sum_{u \in \mathcal{N}_r(v)} \frac{1}{c_{v,r}}\mathbf{W}_r\mathbf{h}_u\big)，其中 c_{v,r} 是一个归一化常数，比如 r 类邻居的个数。在实验 3 中，一个感知方向的层对来自节点所属门的消息和来自其输入的消息使用不同的权重，在两层、三层和四层时分别达到 0.853、0.941 和 1.000；无向 GCN 的最好成绩是八层时的 0.959。

工程中的图，以及整图输出

分子性质预测把分子读作它的化学键图。网格上的学习型模拟器，如 MeshGraphNets（Pfaff 等人，2021），对节点和边的特征编码（网格边携带相对位置），用消息传递块处理，再解码出逐节点的物理量，以此把一个偏微分方程在时间上向前推进一步；它们在传统求解器的轨迹上训练。系统模型也是图：故障树（事件和门）、GSN 安全论证（论点、策略、证据）、SysML 模型（部件和连接器）。把它们当作图而不是展平的文本来读的网络，是检查或补全它们的自然工具。

对于整张图上的标签，比如一棵故障树的顶事件概率是否超过目标值，按求和或均值池化节点向量，再应用一个 MLP。测试集必须是另外的图。

到目前为止，每一类模型都是从样本中学习。工程师的处境却常常相反：微分方程是已知的，来自守恒定律和本构模型，而测量数据很少。物理信息神经网络（physics-informed neural network，PINN；Raissi、Perdikaris 和 Karniadakis，2019）把方程本身用作训练信号。一个网络表示解 u_\theta(\mathbf{x}, t)，它是空间和时间的光滑函数，参数为 \theta；损失衡量的是这个函数在多大程度上不满足方程、边界条件和初始条件，以及现有的数据。

复合损失

把方程写成 \mathcal{N}[u] = 0，其中 \mathcal{N} 是一个微分算子（对热方程，\mathcal{N}[u] = u_t - u_{xx}），把初始条件和边界条件写成 \mathcal{B}[u] = 0（对一端保持为零的情形，\mathcal{B}[u] = u(0, t)）。损失由三项组成：

数据项是在 N_d 个测量值 u_i 上的普通回归。残差项在 N_c 个配点（collocation points）上计算，即强制满足方程的那些位置；它们不需要测量值，所以算力允许多少就可以取多少，可以放在网格上，也可以每一步重新随机抽取。边界项在边界上和 t = 0 处的 N_b 个点上强制满足条件。权重 \lambda_d、\lambda_r 和 \lambda_b 设定各项之间的兑换比率，而选择它们就是难点的大部分，本节余下的内容会说明这一点。没有数据时，PINN 是一个求解器；有数据和一个未知系数时，它是一个估计器。

这里没有什么需要网格：配点就只是点。这正是它的吸引力所在。也正因如此，这种方法没有继承基于网格的求解器的收敛理论所附带的任何误差估计。图 5.14 展示了各部分如何组合在一起。

对输入求导

残差需要网络输出对其输入的导数，而不是对其权重的导数。反向模式自动微分（模块 02 第 4 节）能精确地算出它们，精度达到浮点精度，不需要有限差分：

import math, torch, torch.nn as nn

W0, ZETA, T_END = 2 * math.pi, 0.1, 2.0

def d(u, t):
    """du/dt at every point; create_graph keeps the result differentiable."""
    return torch.autograd.grad(u, t, grad_outputs=torch.ones_like(u),
                               create_graph=True)[0]

net = nn.Sequential(nn.Linear(1, 32), nn.Tanh(), nn.Linear(32, 32), nn.Tanh(),
                    nn.Linear(32, 32), nn.Tanh(), nn.Linear(32, 1))
t_c = torch.linspace(0, T_END, 200).reshape(-1, 1).requires_grad_(True)
u = net(t_c / T_END)                     # inputs scaled to [0, 1]
u_t = d(u, t_c)
u_tt = d(u_t, t_c)                       # the same call applied twice
residual = u_tt + 2 * ZETA * W0 * u_t + W0**2 * u
loss_res = residual.pow(2).mean()        # then add the condition terms and backward()

有两个细节很重要。grad_outputs=torch.ones_like(u) 求的是 \sum_j u_\theta(t_j) 的梯度；每个输出只依赖于它自己的输入，所以这个梯度就是由 200 个逐点导数组成的向量。create_graph=True 把求导计算本身记录到计算图中。没有它，一阶导数返回时是一个常量张量：它不能再求导得到 u''，残差也没有回到 \theta 的路径。有了它，backward() 会穿过这些导数求导，代价是每一步多出相当于几次前向传播的计算量。

激活函数必须有有用的二阶导数。ReLU 是分段线性的，所以它的二阶导数几乎处处为零：ReLU 网络的 u_\theta'' 在每个配点上都为零，u'' 项悄无声息地从残差中消失。PINN 使用光滑的激活函数，最常用的是 tanh。

阻尼振子

贯穿本节的例子是一个质量-弹簧-阻尼系统，从单位位移处由静止释放：

固有频率 \omega_0 = 2\pi rad/s（1 Hz），阻尼比 \zeta = 0.1，时间范围 t \in [0, 2] s。一个候选解的残差为 r(t) = u_\theta'' + 2\zeta\omega_0 u_\theta' + \omega_0^2 u_\theta，\zeta < 1 时的精确解为

所以 \omega_d = 6.252 rad/s，包络按 e^{-0.628 t} 衰减。条件成立：u(0) = 1，且 u'(0) = -\zeta\omega_0 + \omega_d \cdot \zeta\omega_0/\omega_d = 0。正是这个精确解让这个例子有用：实验 4 可以报告相对 L_2 误差，而真实问题不会提供这样的条件。

平凡解

这个方程是齐次的，所以 u \equiv 0 精确地满足它。只有初始条件能排除零，而优化器并不知道它的哪一项表达了建模者的意图：它只会减小最大的那一项。在实验 4 中，一个刚初始化的网络开始时残差损失为 52.5，初始条件损失为 1.36。把 52.5 缩小的最快办法是缩小输出，优化器正是这么做的。取 \lambda_{\text{ic}} = 1 训练 3,000 步之后，残差项约为 6 \times 10^{-3}，初始条件项约为 0.99，相对误差为 0.999：一条接近零的平线。取 \lambda_{\text{ic}} = 0 时误差为 1.000。损失很小，答案却是错的，这是关于 PINN 最需要知道的一点。

三种修复办法，逐一测量

给条件加权。 取 \lambda_{\text{ic}} = 100 时，5,000 步后相对误差约为 0.015，10,000 步后为 0.004。这管用，但权重是试出来的。

无量纲化。 以 1/\omega_0 为单位度量时间。这是工程师本来就会对数值求解器采用的修复办法。

在实验 4 中，无量纲的 PINN 开始时残差损失为 0.034，取 \lambda_{\text{ic}} = 1 时在 7,500 步达到 0.0004 的相对误差，比加权版本最终的 0.0039 低十倍。在固定学习率下它不会停在那里：误差在相邻步之间于约 0.0001 和 0.016 之间跳动，10,000 步时读数为 0.0073，所以要保留最好的检查点，或者让学习率衰减。

通过构造施加条件。 Lagaris、Likas 和 Fotiadis（1998）把解写成一个试函数，无论网络做什么，它都满足条件：u_\theta(t) = 1 + t^2 N_\theta(t) 给出 u_\theta(0) = 1 和 u_\theta'(0) = [2tN_\theta + t^2N_\theta']_{t=0} = 0。初始条件项消失了，零也不再可达。然而在有量纲的单位下，10,000 步后误差仍约为 0.09（在实验 4 “动手试试”第 1 项的一次运行中为 0.090），因为尺度糟糕的残差仍然主导着训练。这种构造消除的是平凡解，而不是尺度问题。先做缩放。

反问题

当方程中的某个量未知时，PINN 才真正发挥所长。把它设为可训练参数即可。实验 4 把 \zeta 当作未知量，从初始值 0.5 出发优化 \log\zeta（这使 \zeta 保持为正，并让步长成为相对量），并加入一个权重为 10 的数据项，数据是在随机时刻采集的 12 个读数，噪声标准差为 0.02。10,000 步后 \zeta = 0.0987，真实值为 0.1。

公平的比较是用闭式解对同样的 12 个读数做最小二乘拟合（SciPy 的 curve_fit），结果为 0.0977 \pm 0.0015。PINN 与经典拟合相当，并没有胜过它。当存在闭式解或廉价的求解器时，把它包进最小二乘拟合里。PINN 的代价物有所值的地方，是两者都不存在，或者必须从偏微分方程中恢复整个场的情形：从一个没有闭式解的结构上的十几个加速度计读数中求阻尼比，从几个热电偶求热导率，或者从少数几个应变片求材料参数。

平凡解之外的失败模式

• 刚性与多尺度行为。 大小相差悬殊的项给出大小相差悬殊的梯度，最大的项在训练，其他项则停滞不前。Wang、Teng 和 Perdikaris（2021）分析了这种梯度失衡，并在训练中自适应地调整权重。
• 谱偏差。 网络先拟合低频（Rahaman 等人，2019），所以振荡或陡峭的解收敛得很慢。傅里叶特征输入，即对少数几个 k 取 [\sin k\hat t, \cos k\hat t]，会有帮助（Tancik 等人，2020）。
• 因果性。 在长的时间窗口上，优化器会在晚期时刻所依赖的早期解稳定下来之前就去拟合晚期时刻。在各个子区间上做时间推进可以恢复这个顺序。
• 训练彻底失败的情形。 Krishnapriyan 等人（2021）表明，高速对流就是其中之一。
• 手工调出的权重。 每一个试出来的 \lambda 都是一个超参数，而它是对着一个你本无法用其他方式检验的答案调出来的。
• 不适定的设定。 去掉一个边界条件，就有无穷多个函数的残差为零；PINN 会以很小的损失返回其中一个（练习 12）。

公平的比较

对于已知几何上的正问题，有限元或有限差分求解器通常更快也更准确。对这个振子，SciPy 的 solve_ivp（RK45，相对容差 10^{-8}，绝对容差 10^{-10}）在运行这些实验的机器上用约 30 ms 就达到约 7 \times 10^{-9} 的相对误差；上面最好的 PINN 用了 7,500 步才达到 4 \times 10^{-4}，按实验 4 每步 4.8 ms 计约为 36 s。关于学习型偏微分方程求解器的论断，应当对照以同样精度运行的强经典基线来检验（McGreivy 和 Hakim，2024）。PINN 的用武之地在于数据同化和反问题，在于网格划分成为瓶颈的复杂区域，以及需要一个关于参数可微的模型的场合。

一个 PINN 求解的是一个实例。换一个载荷或热导率，它就要重新训练。设计工作要用不同的输入把同一个方程问上几百次，而且希望每个答案都来得快。算子学习（operator learning）以映射本身为目标：算子 \mathcal{G} 把一个输入函数 a（一个系数场，例如热导率 k(\mathbf{y})，或一个外力、一个边界条件、一种几何）映射为解函数 u = \mathcal{G}(a)。神经算子（neural operator）\mathcal{G}_\theta 在求解器生成的数据对 (a_i, u_i) 上训练，通常是最小化 \frac{1}{N}\sum_i \|\mathcal{G}_\theta(a_i) - u_i\|^2 / \|u_i\|^2，之后只需一次前向传播就能回答一个新实例。

代理模型是一个老想法

几十年来，工程界一直在为昂贵的程序构建代理模型（surrogate model）。响应面（response surface）用设计变量的低阶多项式拟合少量运行结果。克里金法（Kriging），即高斯过程回归，在运行结果之间插值，并报告自己的不确定性。基于本征正交分解（proper orthogonal decomposition，POD）的降阶模型（reduced-order model）收集解的快照，取它们的前几个奇异向量作为模态 \phi_k(\mathbf{y})，并把一个新解写成

对每个新输入只需求出系数 c_k。神经算子是同一个想法，只是基是学出来的。

DeepONet

DeepONet（Lu 等人，2021）是两半都学出来的 POD 形式。分支（branch）网络在 m 个固定的传感器位置上读取输入函数，返回 p 个系数，\mathbf{b}(a) = b\big(a(\mathbf{y}_1), \dots, a(\mathbf{y}_m)\big) \in \R^p。主干（trunk）网络读取查询点，返回 p 个基函数值，\mathbf{t}(\mathbf{y}) \in \R^p。输出是二者的点积：

主干扮演 POD 模态的角色，分支扮演它们的系数。它的理论可以追溯到 Chen 和 Chen（1995），他们为这种分支-主干形式的算子证明了一个通用近似定理。

傅里叶神经算子

傅里叶神经算子（Fourier neural operator，FNO；Li 等人，2021）在网格上工作。它把输入逐点提升为宽度为 d_v 的场，\mathbf{v}_0 = P(a)，再应用 L 个傅里叶层，最后投影回去，u = Q(\mathbf{v}_L)。每一层为

其中 \mathcal{F} 是沿空间方向的离散傅里叶变换，R_l 把最低的 k_{\max} 个模态各乘以一个学出来的 d_v \times d_v 复数矩阵，其余模态置零，\mathbf{W} 是一个逐点线性映射。谱乘积是一个卷积，其卷积核与整个区域一样宽，所以一层就把每个点与其他每个点耦合起来。

import torch, torch.nn as nn, torch.nn.functional as F

class FourierLayer1d(nn.Module):
    def __init__(self, width=32, k_max=16):
        super().__init__()
        self.k_max = k_max
        self.R = nn.Parameter(torch.randn(k_max, width, width, dtype=torch.cfloat)
                              / width**2)
        self.W = nn.Conv1d(width, width, kernel_size=1)    # pointwise W v + bias

    def forward(self, v):                                  # v: (batch, width, n_grid)
        v_hat = torch.fft.rfft(v)                          # (batch, width, n_grid//2 + 1)
        out = torch.zeros_like(v_hat)
        out[..., :self.k_max] = torch.einsum("bik,kio->bok",
                                             v_hat[..., :self.k_max], self.R)
        spectral = torch.fft.irfft(out, n=v.shape[-1])     # back to the same grid
        return F.gelu(self.W(v) + spectral)

（PyTorch 的 numel 把一个复数元素计为一个，报告 17,440。）同一个层接受 n_grid = 64 或 256：权重不依赖于网格，所以训练好的 FNO 可以在另一种分辨率下求值。但这些限定条件是实实在在的。高于 k_{\max} 的频率从未被建模，粗糙的训练网格会把精细的细节混叠到保留的模态中，而更细的网格并不会拓宽训练分布。图 5.15 画出了一个这样的层。

在非结构网格上，MeshGraphNets 这类图网络模拟器（第 8 节）扮演算子的角色：网格就是图，消息传递取代了傅里叶变换。

有效性

代理模型在它训练时所用的输入分布上是有效的。在这个分布之外，即换了一个几何族、载荷工况或材料，它的误差就无从量化，而它自己不会说出来：它返回的是一个光滑、自信的场。使用代理模型之前：

• 逐个变量地检查每个新输入是否在训练范围之内；
• 在将要使用它的区域内，对照新的求解器运行结果做验证；
• 按工况分别报告误差，而不是只报告一个平均值，并说明有效域；
• 速度只与以同样精度运行的求解器比较（McGreivy 和 Hakim，2024）。

这样使用时，神经算子是为一族设计构建快速近似有限元求解器的实用途径：在几千次求解器运行上训练之后，它一次前向传播就能回答这一族中的新成员，而且只在那些运行所覆盖的族之内有效。

标签很贵，结构却是免费的。一座工厂记录了几个月的振动数据，却只有少数几个标注过的故障。自监督学习（self-supervised learning）在一个由无标签数据自身定义的前置任务（pretext task）上训练编码器，再用线性探测（linear probe）评判所得的表示：在冻结的特征上用少量标签拟合一个逻辑回归。对比学习（contrastive learning）就是这样一种前置任务：在若干候选中分辨出哪一个是同一输入的另一个视图。

InfoNCE 是一个分类损失

取一个锚点嵌入 \mathbf{z}_i、一个正样本 \mathbf{z}_i^+（同一输入的另一个视图）和 N - 1 个负样本（其他输入的视图），共 N 个候选。用 \mathrm{sim}(\mathbf{z}_i, \mathbf{z}_j)/\tau 给每个候选打分，其中 \mathrm{sim} 是经过 L2 归一化的嵌入之间的余弦相似度，\tau 是温度。对这些得分做 softmax，就得到一个必须挑出正样本的分类器，它的交叉熵就是 InfoNCE 损失：

其中正样本也在求和的 N 项之中。余弦得分位于 [-1, 1] 内，所以正是小的 \tau 让 softmax 能够变得自信。

van den Oord、Li 和 Vinyals（2018）证明，这个损失给出了两个视图之间互信息的一个界：I(\mathbf{x}; \mathbf{x}^+) \ge \log N - \mathcal{L}_{\text{InfoNCE}}。损失不能低于零，所以这个估计永远不会超过 \log N；更多的负样本，也就是更大的 batch，允许更大的值。

SimCLR，以及为什么要在投影头之前做探测

SimCLR（Chen 等人，2020）对一个 batch 中 B 个输入的每一个施加两次随机增强，得到 2B 个视图。每个视图的正样本是它的孪生视图；其余 2B - 2 个视图是负样本，所以对每个视图，InfoNCE 在 N = 2B - 1 个候选上计算。这就是 NT-Xent 损失。编码器 f 给出表示 \mathbf{h}，一个小的投影头（projection head）g 给出 \mathbf{z}，损失就在它上面计算。要探测 \mathbf{h}，而不是 \mathbf{z}：投影头学会丢弃增强所改变的一切，而其中可能包括下游任务所需要的东西。在实验 5 中，每类 5 个标签时，在 \mathbf{z} 上的探测得分为 0.62，在 \mathbf{h} 上为 0.96。图 5.16 展示了这条流程。

Wang 和 Isola（2020）把损失的作用拆成两部分：对齐性（alignment）把正样本拉到一起，均匀性（uniformity）把所有嵌入铺满整个球面。如果所有嵌入都相同，每个候选的可能性都一样，损失为 \log(2B - 1) = \log N：这个界证明不了任何信息。阻止这种坍塌的正是负样本。BYOL（Grill 等人，2020）这类非对比方法则改用两个不对称的网络来避免使用负样本。

增强就是监督

增强规定了编码器必须忽略哪些差异，从而决定了它学到什么。在实验 5 中，四类机器振动的区别在于频谱，但每个窗口都从一个随机相位开始。增强中包含随机时间平移时，编码器学到相位不变性；没有它，5 个标签的探测从 0.96 降到 0.62，比未训练编码器的 0.53 高不了多少。

只有当各类别在增强保持不变的那些方面毫无差别时，增强才会移除任务信息。旋转 180 度会让 6 和 9 成为同一个数字；颜色抖动会移除识别腐蚀所依赖的颜色。在实验 5 中，16 倍的增益范围并没有造成损害（“动手试试”第 1 项：0.956、0.990 和 0.997），因为各类别在谐波比和信噪比上也有差别。要根据任务真实的不变性来选择增强。

CLIP、掩码建模，以及其中的教训

CLIP（Radford 等人，2021）在 4 亿个图像-说明文字对上同时训练一个图像编码器和一个文本编码器。对一个含 B 对的 batch，它构造 B \times B 的相似度矩阵，正样本位于对角线上，并对每一行（一张图像对 B 条说明文字）和每一列（一条说明文字对 B 张图像）应用 InfoNCE，再取两者的平均。零样本分类把“a photo of a {label}”这类提示词嵌入，然后选出最近的那个。这些嵌入为文本到图像模型提供条件，也支撑着许多检索系统；检索增强生成在 AI 智能体系列中讲解。

另一个分支是掩码建模（masked modelling）：遮住输入的一部分，再预测它。BERT 预测被掩码的 token（模块 06）；掩码自编码器遮住一幅图像 75% 的图像块，再把它们重构出来（He 等人，2022）。下一个 token 预测，即模块 07 和 08 的目标，在同样的意义上也是自监督的。

普遍的教训是：一个没有标签的前置任务，可以产生一种能迁移到只有少量标签的任务上的表示。在工程中：先在几个月的无标签传感器日志上做预训练，再在少数几个标注过的故障上拟合分类器；或者把事故报告嵌入，以检索相似的历史案例。

稠密网络对每个输入都运行它的全部参数，所以容量和计算量一起增长。混合专家（mixture of experts，MoE）把二者分开：保留 E 个专家网络和一个小的路由器（router），把每个输入送到其中 k 个。只有被选中的专家运行。

这一层

在 Transformer 中（模块 06 第 5 节），每个块都包含一个前馈网络，即一个对每个 token 的向量单独作用的两层 MLP。MoE 层用 E 个这样的 MLP 和一个路由器取代它：

其中 \mathbf{W}_r \in \R^{E \times d}，k = 1 或 2。重新归一化后的权重 \tilde g_e 在被选中的专家上之和为 1。对专家的选择不可微，但权重可微，路由器正是通过它们来学习的。

这个想法由来已久。Jacobs、Jordan、Nowlan 和 Hinton（1991）训练了局部专家的自适应混合：一个在各个专家网络之间做选择的软门控，每个专家专攻输入空间的一个区域。在不同运行工况之间切换模型的工程师，例如做增益调度（gain scheduling）时，就已经手工构建过一个这样的系统，调度变量就是门控。Shazeer 等人（2017）用带噪声的 top-k 门控把门控变得稀疏，并把数千个专家放进了一个语言模型。

import torch, torch.nn as nn

class MoE(nn.Module):
    def __init__(self, d, d_ff, n_experts=8, k=2):
        super().__init__()
        self.k = k
        self.router = nn.Linear(d, n_experts, bias=False)
        self.experts = nn.ModuleList(
            nn.Sequential(nn.Linear(d, d_ff), nn.GELU(), nn.Linear(d_ff, d))
            for _ in range(n_experts))

    def forward(self, x):                                   # x: (n_tokens, d)
        probs = self.router(x).softmax(dim=-1)              # (n_tokens, E)
        top_p, top_e = probs.topk(self.k, dim=-1)
        top_p = top_p / top_p.sum(dim=-1, keepdim=True)     # renormalise over the k
        y = torch.zeros_like(x)
        for e, expert in enumerate(self.experts):
            token, slot = (top_e == e).nonzero(as_tuple=True)
            if len(token):                                  # run e on its tokens only
                y[token] += top_p[token, slot, None] * expert(x[token])
        return y

参数量与计算量

参数量随 E 增长；每个 token 的计算量随 k 增长。

路由坍塌与均衡损失

路由中有一个反馈回路。一个专家收到的 token 越多，得到的梯度就越多，于是变得更好，又被更多地选中；挨饿的专家则永远得不到改善。症状是大部分 token 集中在一两个专家上，模型悄然变成了一个小的稠密模型。

Switch Transformer（Fedus、Zoph 和 Shazeer，2022）加入了一个负载均衡损失（load-balancing loss）。在一个 batch 的 token 上，令 f_e 为分派给专家 e 的比例，P_e 为 e 的平均路由概率：

当同一批专家既拿到了 token 又拿到了概率时，这个和就很大。当路由遵循概率，即 f_e \approx P_e 时，它变成 E\sum_e P_e^2；由柯西-施瓦茨不等式，1 = \big(\sum_e P_e\big)^2 \le E\sum_e P_e^2，所以它至少为 1，在均匀路由时取等号。f_e 来自一个硬选择，没有梯度；这个损失通过 P_e 训练路由器，按每个专家已经收到的 token 数成比例地降低它的概率。Switch 使用 \lambda_{\text{bal}} = 0.01。

容量（capacity）是另一个控制手段。每个专家在每个 batch 中最多处理 \text{容量因子} \times (\text{token 数}/E) 个 token。溢出的 token 被丢弃：它们跳过这一层，经由残差连接继续传下去。什么都不会报错，所以被丢弃的 token 是一种无声的质量损失；要统计它们。

其他难点

当专家分布在多个加速器上时，每个 MoE 层都要把 token 送到它们的专家那里再送回来，这需要通信开销；训练在早期也容易不稳定。DeepSeek-V3（总参数 671B，每个 token 激活 37B）把许多细粒度专家与每个 token 都会使用的共享专家结合起来，并通过调整路由得分中每个专家的偏置来均衡负载，不需要辅助损失。截至 2026 年，MoE 是公开资料最充分的那些最大语言模型中的一种常见设计；它在这一规模上的工程实现，包括专家并行，属于模块 08 的内容。

专家的专门化并不像名字所暗示的那样有语义：Mixtral 的作者报告，没有发现按主题分配专家的明显规律。而且节省的是计算，不是内存。尽管每个 token 只用其中 k 个专家，所有专家都必须加载：Mixtral 的 46.7B 个参数按每个 2 字节计约占 93 GB，而一个 13B 的稠密模型只占 26 GB（模块 10）。

各类模型满足的是不同的需求，而每一类都有一条必须先胜过的基线，然后才算物有所值。下表保留了九种需求，并把基线和代价都加了进来。

图 5.18 把同样的选择画成一张决策流程图。

三个实例

对一个没有故障标签的泵群做振动监测。 在正常运行的窗口上训练一个自编码器，把阈值设在留出的正常数据误差的某个百分位数上，并在同样的数据上与 PCA 的 Q 统计量比较（第 2 节）。一旦有几个故障被标注，就像实验 5 那样，在无标签日志上用对比学习预训练一个编码器，再拟合一个探测器；FFT 幅值探测器就是它的基线。

为给定载荷工况生成候选支架几何。 用一个以载荷工况为条件的条件扩散模型；如果几何是图像或体素网格，就在潜在空间中进行（第 6 节）。生成的支架是一个提案，而不是一个设计：每个候选都要像其他任何设计一样，用求解器校核。

一族开孔板上的应力场。 用在求解器运行上训练的神经算子，如果网格各不相同，则用网格 GNN（第 10 节）。它只对它见过的孔径和载荷有效，而且首先要与用同样的运行结果构建的 POD 代理模型比较。

整个系列的规则是：先说明基线。一个胜不过基线的模型没有学到任何有用的东西，无论它的损失曲线看起来多么好。

下面每一条都给出你看到的症状、通常的原因和修复方法。其中有几条是特意安排在实验中出现的。

VAE 的样本看起来都一样

症状：样本是几乎一模一样的模糊图像，每个维度上的 KL 项读数都是 0.00 奈特。原因：后验坍塌：相对于似然项，KL 权重太强（\beta > 1，或者在 [0, 1] 像素上使用相当于 \sigma_x^2 = \tfrac12 的 MSE 求和损失），或者解码器不需要 \mathbf{z} 就能建模 \mathbf{x}。在实验 1 中，\beta = 4 使 8 个维度中活跃的维度为 0 个。修复：使用 \beta \le 1 并配以尺度正确的似然，让 KL 权重从 0 开始预热，或使用自由比特（free bits），并监测每个维度的 KL 和活跃单元的数目。

AUC 很高却漏掉故障的异常检测器

症状：一个 AUC 不错的重构误差检测器，投入使用后漏掉了许多真实故障。原因：AUC 是在所有阈值上的平均，而与正常数据相似的异常能被很好地重构。在实验 1 中，AUC 为 0.95，但在 4% 的误报率下只检测出 58% 的异常。修复：报告在工作阈值下的检出率，与 PCA 的 Q 统计量比较，在已知故障上评估，并在运行条件改变时重新校准。

只会生成同样六张图像的 GAN

症状：训练“一切顺利”，每个样本都令人信服，但彼此不同的样本很少。原因：模式坍塌：对抗损失中没有任何东西奖励对数据的覆盖。修复：除了质量，还要衡量多样性（覆盖的模式数；从留出数据到样本的、类似召回率的最近邻距离）。用梯度惩罚或谱归一化来稳定训练，或者改用扩散模型。

扩散样本颜色发淡，或者第一个反向步骤就爆炸

症状：样本始终达不到极端值，或者采样一开始就发散。原因：噪声调度的某个端点不对。使用 Ho 等人的 \beta_t 范围而 T = 200 而不是 1,000 时，\bar\alpha_T = 0.13，所以采样从网络从未见过的噪声开始。\beta_T 接近 1 时（余弦调度在 0.999 处截断），第一次更新会把误差放大 1/\sqrt{\alpha_T} = 31.6 倍，引导又会进一步放大它们。修复：训练前检查 \bar\alpha_T 和最大的 \beta_t；使用余弦调度或零终端信噪比重缩放；给 \beta_t 设上限（实验 2 中为 0.5），或者把 \hat{\mathbf{x}}_0 截断到数据范围内。

快速采样器给出模糊的样本，强引导给出千篇一律的样本

症状：少步数的样本落在数据之外；强引导的样本过饱和且彼此相似。原因：步数太少（实验 2 中 DDIM 类似精确率的距离：200 步时 0.022，5 步时 0.046，1 步时 1.71），或者引导系数太大（类似召回率的距离：w = 1 时 0.021，w = 7 时 0.043）。修复：使用 20–50 步 DDIM 或蒸馏过的采样器，并且在选择 w 时既看外观，也看一个多样性指标。

训练样本的近似副本

症状：生成的样本与训练样本几乎完全相同。原因：记忆，训练集小或有重复、训练时间长时更容易发生（Carlini 等人，2023）。修复：对数据去重，把每个样本到最近训练样本的距离与留出样本的这一距离作比较，并对生成的设计做与其他任何设计相同的校核。

只预测一个类别的深层 GNN

症状：一个十层的 GNN 给每个节点几乎相同的特征。原因：过平滑：反复的平均以 |\lambda_2|^k 的速率把所有节点向量推向同一个方向，而且深的朴素堆叠本身也训练得不好。在实验 3 中，12 层和 16 层的朴素网络得分 0.781，即多数类的比例；带残差连接的 16 层网络得分 0.948。修复：使用 2–4 层或残差连接，并逐层测量特征相似度。

在依赖方向的属性上停滞不前的 GNN

症状：在一个依赖边的方向或类型的属性上，例如故障树中某个节点上方是什么，精度停滞不前。原因：对称归一化的邻接矩阵把父节点、子节点和兄弟节点混在一起。修复：给每种边方向或类型各自的权重（关系图卷积网络），或者加入边特征。在实验 3 中：无向时最好为 0.959（八层），考虑方向时为 1.000（四层）。

在新图上失败的 GNN

症状：评估得分很好，在新图上的结果却很差。原因：同一张图的节点被分到了训练集和测试集两边，所以测试节点的邻域在训练中已经见过。修复：只要模型将来会遇到新图，就按图划分（整棵故障树、整个网格）。

收敛到零的 PINN

症状：PINN 以很小的损失返回 u = 0，或者一条忽略了初始数据或边界数据的光滑曲线。原因：零满足齐次方程，而条件项的分量被压倒了：在实验 4 中，残差损失开始时为 52.5，而条件损失为 1.36，最终误差为 0.999。修复：无量纲化（最好时误差为 0.0004，不需要权重），提高条件权重（权重为 100 时为 0.004），或者通过构造施加条件。

残差很小、答案却错误的 PINN

症状：缓慢的趋势拟合上了，振荡、陡峭的锋面或晚期时刻却没有；或者残差极小而解是错的。原因：谱偏差、刚性和损失失衡，或者设定不适定：没有边界条件时热方程有无穷多个解，PINN 在残差为 4 \times 10^{-5} 时找到了其中一个，误差为 0.75（练习 12）。修复：傅里叶特征输入、自适应损失权重、时间推进、检查问题是否完全适定，或者改用经典求解器，对正问题它通常更快也更准确。

对新设计自信地给出错误结果的代理模型

症状：神经算子或其他代理模型返回一个看似合理的场，而求解器运行的结果与之矛盾。原因：输入位于训练族之外：另一个几何族、载荷工况或材料。修复：对照训练范围检查输入，在使用代理模型的地方对照新的求解器运行结果做验证，并按工况报告误差和有效域。

没有帮助的对比学习编码器

症状：预训练特征在下游毫无用处，或者损失停在 \log(2B - 1)。原因：增强移除了任务需要的信息（180 度旋转让 6 和 9 成为一类；颜色抖动移除了腐蚀的颜色），或者嵌入坍塌了。修复：根据任务真实的不变性选择增强，并检查每个类别的探测精度。增强就是监督：在实验 5 中，去掉时间平移后，5 个标签的探测得分为 0.62，而不是 0.96。

只使用一两个专家的 MoE 层

症状：大部分 token 都去了一两个专家。原因：富者愈富的反馈导致路由坍塌。修复：加入负载均衡损失或基于偏置的均衡，并监测每个专家的 token 比例和被丢弃的 token 数。

目标。你在 8×8 的手写数字图像上构建第 2 节和第 3 节的三种工具，并把每一种都与一个基线比较。首先，在相同的编码维数下，把一个非线性自编码器与主成分分析（PCA）比较。然后训练第 3 节的变分自编码器，画出它的二维潜在空间，并用它的解码器采样和插值。接着故意提高 KL 权重，制造后验坍塌，这样当它意外出现时，你就能认出它的数字特征。最后，构建一个基于重构误差的异常检测器，其阈值在留出的正常数据上选定，并诚实地把它的检出率与经典的 PCA 监测器比较。数据随 scikit-learn 提供，无需下载，本实验在笔记本电脑 CPU 上约需两分钟。你需要 NumPy、scikit-learn、PyTorch 和 matplotlib。打印出的数字可能与你的结果在最后几位上不同。

步骤 1：加载数字图像并划分

load_digits 数据集有 1,797 张 8×8 像素的图像，取值为 0 到 16 的整数。除以 16 后像素落在 [0, 1] 内，这正是 sigmoid 输出和伯努利似然所要求的。划分是分层的，所以每个数字在两部分中都保持各自的占比；在模型完成之前，测试集不会被动用。

import time

import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from sklearn.datasets import load_digits
from sklearn.decomposition import PCA
from sklearn.metrics import roc_auc_score
from sklearn.model_selection import train_test_split
from sklearn.neighbors import KNeighborsClassifier

np.random.seed(0)
torch.manual_seed(0)

X, y = load_digits(return_X_y=True)
X = (X / 16.0).astype(np.float32)
X_tr, X_te, y_tr, y_te = train_test_split(X, y, test_size=0.2, stratify=y, random_state=0)
print(X_tr.shape, X_te.shape)
print(f"pixel range {X.min():.1f} to {X.max():.1f}, mean pixel {X_tr.mean():.3f}")

# A model that outputs the mean training image for every input: the floor to beat.
mean_img = X_tr.mean(axis=0)
print(f"MSE per pixel of the mean image on the test set: {((X_te - mean_img) ** 2).mean():.4f}")
Xtr_t, Xte_t = torch.from_numpy(X_tr), torch.from_numpy(X_te)

(1437, 64) (360, 64)
pixel range 0.0 to 1.0, mean pixel 0.305
MSE per pixel of the mean image on the test set: 0.0740

最后一行是零个成分的重构：一个对任何具体图像都一无所知的模型的误差。下面的每个模型都必须胜过它，而这一比较告诉你每个模型解释了多少像素方差。后面会看到，两个 PCA 成分能消除这一误差的 28%（1 - 0.0532/0.0740），八个成分能消除约三分之二。

步骤 2：PCA 基线

有 d_z 个成分的 PCA 是最优的线性自编码器（第 2 节），所以它是任何非线性模型在相同编码维数下必须胜过的基线。它在训练集上拟合、在测试集上评分：重构是 inverse_transform(transform(X))，误差是每像素的均方差。

def mse_per_pixel(a, b):
    return float(((a - b) ** 2).mean())


pca_mse = {}
for d_z in (2, 8):
    pca = PCA(n_components=d_z).fit(X_tr)
    recon = pca.inverse_transform(pca.transform(X_te))
    pca_mse[d_z] = mse_per_pixel(recon, X_te)
    print(f"PCA d_z={d_z}: test MSE per pixel {pca_mse[d_z]:.4f}")

PCA d_z=2: test MSE per pixel 0.0532
PCA d_z=8: test MSE per pixel 0.0246

两个成分消除了均值图像误差的 28%，八个成分消除了 67%。剩下的误差是像素细节，这么大的平坦子空间捕捉不到它们。

步骤 3：欠完备自编码器

图 5.2 的自编码器在编码两侧各有一个 128 个单元的隐藏层。两个模型共用一个 train() 辅助函数：Adam，学习率 10^{-3}，batch 大小为 64，200 个轮次，以及一个设定了种子、固定打乱顺序的 torch.Generator，这样重新运行会得到相同的数字。辅助函数以函数的形式接收损失，因为下一步的 VAE 需要另一种损失。

def train(model, data, loss_fn, epochs=200, batch=64, lr=1e-3, seed=0):
    """Adam training with seeded shuffling. loss_fn(model, xb) returns a scalar."""
    gen = torch.Generator().manual_seed(seed)
    opt = torch.optim.Adam(model.parameters(), lr=lr)
    n = len(data)
    for _ in range(epochs):
        order = torch.randperm(n, generator=gen)
        for start in range(0, n, batch):
            xb = data[order[start:start + batch]]
            loss = loss_fn(model, xb)
            opt.zero_grad()
            loss.backward()
            opt.step()
    return model


class AutoEncoder(nn.Module):
    def __init__(self, d_in=64, d_z=2, h=128):
        super().__init__()
        self.enc = nn.Sequential(nn.Linear(d_in, h), nn.GELU(), nn.Linear(h, d_z))
        self.dec = nn.Sequential(nn.Linear(d_z, h), nn.GELU(), nn.Linear(h, d_in), nn.Sigmoid())

    def forward(self, x):
        return self.dec(self.enc(x))


def ae_loss(model, xb):
    return F.mse_loss(model(xb), xb)


aes = {}
for d_z in (2, 8):
    torch.manual_seed(0)
    t0 = time.time()
    ae = train(AutoEncoder(d_z=d_z), Xtr_t, ae_loss)
    aes[d_z] = ae
    with torch.no_grad():
        test_mse = mse_per_pixel(ae(Xte_t).numpy(), X_te)
    print(f"AE d_z={d_z}: test MSE per pixel {test_mse:.4f} "
          f"(PCA {pca_mse[d_z]:.4f}), {time.time() - t0:.0f} s")

AE d_z=2: test MSE per pixel 0.0368 (PCA 0.0532), 12 s
AE d_z=8: test MSE per pixel 0.0110 (PCA 0.0246), 10 s

自编码器在两种维数下都胜过 PCA。数字图像位于像素空间中一个弯曲的低维曲面上，非线性解码器能沿着它走，而平坦子空间做不到。输出上的 sigmoid 把重构保持在 [0, 1] 内。不要把这个差距当作普遍规律：在近似线性的数据上，PCA 一样好，而且几乎没有代价。

步骤 4：变分自编码器

这里的 VAE 就是第 3 节的那个类，只有三处改动。解码器输出伯努利 logits，所以重构项是 64 个像素上的二元交叉熵之和。KL 项按潜在维度保留，这样就能逐维诊断坍塌。一个 beta 属性乘在 KL 上，由此得到 beta-VAE；\beta = 1 就是证据下界（ELBO）本身。两项都以每张图像的奈特数计，在 batch 上取平均，所以损失就是负 ELBO。

一个报告函数给出关键的四个数：负 ELBO、它的两个组成部分，以及每个维度的 KL。重构项对每张图像用一个 \mathbf{z} 样本，以固定的种子抽取。

class VAE(nn.Module):
    def __init__(self, d_in=64, d_z=2, h=128, beta=1.0, gaussian=False):
        super().__init__()
        self.enc = nn.Sequential(nn.Linear(d_in, h), nn.GELU(), nn.Linear(h, 2 * d_z))
        self.dec = nn.Sequential(nn.Linear(d_z, h), nn.GELU(), nn.Linear(h, d_in))
        self.beta = beta          # KL weight; 1 gives the ELBO itself
        self.gaussian = gaussian  # True: summed squared error, i.e. sigma_x^2 = 1/2

    def encode(self, x):
        mu, logvar = self.enc(x).chunk(2, dim=-1)
        return mu, logvar

    def terms(self, x):
        """Reconstruction (nats) and KL per dimension, both averaged over the batch."""
        mu, logvar = self.encode(x)
        z = mu + torch.exp(0.5 * logvar) * torch.randn_like(mu)  # reparameterisation
        out = self.dec(z)
        if self.gaussian:
            recon = F.mse_loss(out, x, reduction="sum") / len(x)
        else:  # Bernoulli decoder: out holds logits
            recon = F.binary_cross_entropy_with_logits(out, x, reduction="sum") / len(x)
        kl_dim = (-0.5 * (1 + logvar - mu**2 - logvar.exp())).mean(0)  # closed form, per dim
        return recon, kl_dim


def vae_loss(model, xb):
    recon, kl_dim = model.terms(xb)
    return recon + model.beta * kl_dim.sum()


def vae_report(model, data):
    """Negative ELBO, reconstruction, total KL and KL per dimension on `data`."""
    torch.manual_seed(1)
    with torch.no_grad():
        recon, kl_dim = model.terms(data)
    return float(recon + kl_dim.sum()), float(recon), float(kl_dim.sum()), kl_dim.numpy()


torch.manual_seed(0)
vae2 = train(VAE(d_z=2), Xtr_t, vae_loss)
neg_elbo, recon, kl, kl_dim = vae_report(vae2, Xte_t)
print(f"test -ELBO {neg_elbo:.1f} = reconstruction {recon:.1f} + KL {kl:.1f} nats")
print("KL per dimension:", np.round(kl_dim, 2))

test -ELBO 24.9 = reconstruction 22.6 + KL 2.3 nats
KL per dimension: [1.09 1.18]

把这些数字当作信息量来读。KL 是编码平均携带的关于一张图像的信息，以奈特计；除以 \ln 2 就得到比特数。这里的 2.3 奈特就是 3.3 比特，相当于在十个等可能的选项中做一次选择所含的信息（\ln 10 = 2.30）：大约是一个类别标签所携带的信息，不多不少。重构项是解码器从这个编码仍然无法预测的部分。两个维度都在使用，KL 相近。KL 为 0.00 的维度就是坍塌了。

步骤 5：并排比较潜在空间

评判一个潜在空间，要看它保留了什么，以及它的点与点之间是什么。第一个检验是近邻可分性：在训练集的编码上拟合一个 5 近邻分类器，在测试集的编码上评分。对 VAE，编码取均值 \boldsymbol{\mu}(\mathbf{x})。原始的 64 个像素给出上限。第二个检验是可视化的：解码一个 10 × 10 的编码网格。对 VAE，网格覆盖 [-2.5, 2.5]^2，先验几乎把全部概率质量都放在这里。普通自编码器没有先验，所以它的网格覆盖它自己训练编码的包围盒。

def codes_ae(model, data):
    with torch.no_grad():
        return model.enc(data).numpy()


def codes_vae(model, data):
    with torch.no_grad():
        return model.encode(data)[0].numpy()


pca2 = PCA(n_components=2).fit(X_tr)
code_sets = {
    "PCA-2": (pca2.transform(X_tr), pca2.transform(X_te)),
    "AE-2": (codes_ae(aes[2], Xtr_t), codes_ae(aes[2], Xte_t)),
    "VAE-2": (codes_vae(vae2, Xtr_t), codes_vae(vae2, Xte_t)),
}
for name, (c_tr, c_te) in code_sets.items():
    knn = KNeighborsClassifier(5).fit(c_tr, y_tr)
    print(f"{name}: 5-NN test accuracy {knn.score(c_te, y_te):.3f}")
raw = KNeighborsClassifier(5).fit(X_tr, y_tr)
print(f"raw 64 pixels: {raw.score(X_te, y_te):.3f}  (reference)")


def tile(images, n_rows, n_cols):
    """Arrange n_rows * n_cols 8x8 images in one mosaic, row by row."""
    imgs = images.reshape(n_rows, n_cols, 8, 8)
    return imgs.transpose(0, 2, 1, 3).reshape(n_rows * 8, n_cols * 8)


def decode_grid(decode, lo, hi, n=10):
    """Decode an n x n grid of 2D codes; the top row has the largest second coordinate."""
    g1 = np.linspace(lo[0], hi[0], n)
    g2 = np.linspace(hi[1], lo[1], n)
    zz = np.array([[a, b] for b in g2 for a in g1], dtype=np.float32)
    with torch.no_grad():
        return decode(torch.from_numpy(zz)).numpy()


ae_codes = code_sets["AE-2"][0]
ae_grid = decode_grid(aes[2].dec, ae_codes.min(0), ae_codes.max(0))
vae_grid = decode_grid(lambda z: torch.sigmoid(vae2.dec(z)), (-2.5, -2.5), (2.5, 2.5))

fig, axes = plt.subplots(2, 2, figsize=(10, 10))
for ax, name in zip(axes[0], ("AE-2", "VAE-2")):
    c = code_sets[name][1]
    sc = ax.scatter(c[:, 0], c[:, 1], c=y_te, cmap="tab10", s=12)
    ax.set_title(f"{name}: test-set codes coloured by digit")
    ax.set_xlabel("code dimension 1")
    ax.set_ylabel("code dimension 2")
fig.colorbar(sc, ax=axes[0], label="digit", ticks=range(10))
axes[1][0].imshow(tile(ae_grid, 10, 10), cmap="gray_r")
axes[1][0].set_title("AE: decoded grid over the box of its training codes")
axes[1][1].imshow(tile(vae_grid, 10, 10), cmap="gray_r")
axes[1][1].set_title("VAE: decoded grid over $[-2.5, 2.5]^2$")
for ax in axes[1]:
    ax.set_xticks([])
    ax.set_yticks([])
plt.show()

PCA-2: 5-NN test accuracy 0.617
AE-2: 5-NN test accuracy 0.839
VAE-2: 5-NN test accuracy 0.750
raw 64 pixels: 0.978  (reference)

自编码器把各个数字分得最开。二维 PCA 编码最差，因为平坦的投影会把几个类别叠在一起。VAE 介于两者之间，散点图说明了原因：KL 项把每一团点都拉向原点、拉向单位宽度，于是各个簇彼此相接。解码网格说明了这些簇换来了什么。VAE 网格的每一格都是一个像样的数字，数字在平面上平滑地变化。自编码器的网格在簇附近数字更清晰，而在簇之间的空隙里则是污迹或不合理的形状，因为从来没有任何东西要求解码器在那里给出合理的输出。两种二维编码都远不及原始像素：十个类别无法无损地装进两个数里。

步骤 6：采样与插值

从 VAE 生成只需两行：从先验抽取 \mathbf{z} \sim \mathcal{N}(0, \mathbf{I})，然后解码。插值取两张真实测试图像（一个“1”和一个“7”）的潜在均值，解码二者连线上的八个点。由于 KL 项已使聚合后验接近先验，这条直线始终位于解码器熟悉的区域。

torch.manual_seed(3)
z = torch.randn(16, 2)
with torch.no_grad():
    samples = torch.sigmoid(vae2.dec(z)).numpy()

i_one = int(np.where(y_te == 1)[0][0])
i_seven = int(np.where(y_te == 7)[0][0])
mu_ends = codes_vae(vae2, Xte_t[[i_one, i_seven]])
alphas = np.linspace(0, 1, 8, dtype=np.float32)[:, None]
z_path = torch.from_numpy((1 - alphas) * mu_ends[0] + alphas * mu_ends[1])
with torch.no_grad():
    path = torch.sigmoid(vae2.dec(z_path)).numpy()
print("latent means of the two ends:", np.round(mu_ends.astype(float), 2).tolist())

fig, axes = plt.subplots(1, 2, figsize=(11, 3.6), gridspec_kw={"width_ratios": [1, 2]})
axes[0].imshow(tile(samples, 4, 4), cmap="gray_r")
axes[0].set_title("16 samples, $z \\sim N(0, I)$")
axes[1].imshow(tile(path, 1, 8), cmap="gray_r")
axes[1].set_title("interpolation from a test '1' to a test '7' (8 steps)")
for ax in axes:
    ax.set_xticks([])
    ax.set_yticks([])
plt.show()

latent means of the two ends: [[-2.09, -0.77], [0.24, -0.33]]

在 8×8 像素下，样本很粗糙，而且都是模糊的：解码器输出的是与一个编码相容的所有图像的平均（第 3 节，样本模糊）。要找的是像数字的笔画，而不是噪声；有些样本是两个类别的含糊混合，正如散点图中的重叠所预示的。插值一步一步地改变形状，而不是让一张图像淡出、另一张淡入。在像素空间中混合“1”和“7”，得到的会是两张叠在一起的淡图像。这就是平滑潜在空间的实际含义。

步骤 7：故意制造后验坍塌

现在编码维数取 d_z = 8，并改变 KL 权重：\beta = 0.5、1 和 4。对每个模型，实验打印总 KL、每个维度的 KL、重构项以及活跃单元（active units）的数目，即测试集上均值编码 \mu_j(\mathbf{x}) 的方差超过 0.01 的维度 j。第四个模型使用第 3 节那个精简类的损失，即平方误差之和，它相当于 \sigma_x^2 = 1/2 的高斯解码器。它的重构项处在另一个尺度上，所以只有它的 KL 和活跃单元能与其他模型比较。

def active_units(model, data, threshold=0.01):
    with torch.no_grad():
        mu = model.encode(data)[0]
    return int((mu.var(0) > threshold).sum())


runs = {}
settings = [("beta=0.5", 0.5, False), ("beta=1", 1.0, False), ("beta=4", 4.0, False),
            ("beta=1, squared error", 1.0, True)]
for name, beta, gaussian in settings:
    torch.manual_seed(0)
    model = train(VAE(d_z=8, beta=beta, gaussian=gaussian), Xtr_t, vae_loss)
    _, recon, kl, kl_dim = vae_report(model, Xte_t)
    runs[name] = (model, kl_dim)
    print(f"{name:22s} KL {kl:5.2f}  active {active_units(model, Xte_t)}  "
          f"reconstruction {recon:5.1f}")
    print(f"{'':22s} KL per dim {np.round(kl_dim, 2)}")

fig, axes = plt.subplots(1, 4, figsize=(13, 3.2), sharey=True)
for ax, (name, (model, kl_dim)) in zip(axes, runs.items()):
    ax.bar(range(8), kl_dim)
    ax.set_title(name)
    ax.set_xlabel("latent dimension")
axes[0].set_ylabel("KL per dimension (nats)")
plt.show()

beta=0.5               KL  6.51  active 8  reconstruction  18.9
                       KL per dim [1.21 0.9  0.46 1.25 0.42 0.38 0.69 1.2 ]
beta=1                 KL  3.57  active 6  reconstruction  21.0
                       KL per dim [0.84 0.53 0.   0.91 0.   0.01 0.37 0.9 ]
beta=4                 KL  0.00  active 0  reconstruction  27.2
                       KL per dim [0. 0. 0. 0. 0. 0. 0. 0.]
beta=1, squared error  KL  0.48  active 3  reconstruction   4.3
                       KL per dim [0.1  0.   0.   0.21 0.16 0.   0.   0.  ]

KL 权重决定了编码可以携带多少信息。\beta = 0.5 时，编码携带 6.5 奈特，八个维度全都在用。\beta = 1 时，也就是真正的 ELBO，编码携带 3.6 奈特，其中几乎全部由五个维度承担（3.57 奈特中的 3.55）；另外三个维度停在先验附近，其中一个仍随输入略有变化，这就是为什么有六个单元算作活跃。\beta = 4 时，每个维度的 KL 都是零：编码器对每张图像都输出先验，解码器无论 \mathbf{z} 是什么都产生同一个平均数字。这就是后验坍塌，而在这里它不是训练中的意外。是目标函数本身偏好它：使用编码带来的重构收益小于其 KL 代价的四倍（第 3 节给出了具体数字）。平方误差损失出于同样的原因、以较温和的形式让大多数维度坍塌，因为 \sigma_x^2 = 1/2 使得放弃重构的代价很小。

要根据每个维度的 KL 来诊断坍塌，如上所示，而不是根据损失：坍塌的模型损失非常平稳。

步骤 8：阈值选得诚实的异常检测器

最后一项任务是第 2 节中的用途。检测器只在正常数据（数字 0 到 8）上训练，数字 9 扮演一个未曾预料的故障。训练用的数字再划分一次：80% 用来拟合自编码器，20% 留出作为正常数据的验证集，阈值就从验证集读出。阈值取验证误差的第 95 百分位数，所以按照构造，正常数据上的误报率应当接近 5%。数字 9 上的检出率不由任何东西设定，而是测出来的。测试集只在最后使用一次。

基线是工程师首先会构建的那种监测器：在同样的正常数据上拟合 8 个成分的 PCA，用 Q 统计量（即平方重构误差）评分，并用同一验证集上它自己的第 95 百分位数作为阈值。

normal_tr = y_tr <= 8
X_norm = X_tr[normal_tr]
X_fit, X_val = train_test_split(X_norm, test_size=0.2, random_state=0)
X_test_normal = X_te[y_te <= 8]
X_test_nine = X_te[y_te == 9]
print(f"fit {len(X_fit)}, validation {len(X_val)}, test normal {len(X_test_normal)}, "
      f"test nines {len(X_test_nine)}")

torch.manual_seed(0)
ae_norm = train(AutoEncoder(d_z=8), torch.from_numpy(X_fit), ae_loss)
pca_norm = PCA(n_components=8).fit(X_fit)


def error_ae(data):
    with torch.no_grad():
        t = torch.from_numpy(data)
        return ((ae_norm(t) - t) ** 2).mean(1).numpy()


def error_pca(data):
    return ((pca_norm.inverse_transform(pca_norm.transform(data)) - data) ** 2).mean(1)


results = {}
for name, score in (("autoencoder", error_ae), ("PCA Q statistic", error_pca)):
    threshold = np.percentile(score(X_val), 95)
    e_norm, e_nine = score(X_test_normal), score(X_test_nine)
    auc = roc_auc_score(np.r_[np.zeros(len(e_norm)), np.ones(len(e_nine))],
                        np.r_[e_norm, e_nine])
    results[name] = (e_norm, e_nine, threshold)
    print(f"{name:16s} threshold {threshold:.4f}  false alarms {np.mean(e_norm > threshold):.3f}"
          f"  9s detected {np.mean(e_nine > threshold):.3f}  AUC {auc:.3f}")

fig, axes = plt.subplots(1, 2, figsize=(11, 3.8), sharey=True)
for ax, (name, (e_norm, e_nine, threshold)) in zip(axes, results.items()):
    bins = np.linspace(0, max(e_norm.max(), e_nine.max()), 40)
    ax.hist(e_norm, bins=bins, alpha=0.6, label="normal test digits 0-8")
    ax.hist(e_nine, bins=bins, alpha=0.6, label="test 9s (anomalies)")
    ax.axvline(threshold, color="k", linestyle="--", label="95th-percentile threshold")
    ax.set_title(name)
    ax.set_xlabel("reconstruction error (MSE per pixel)")
axes[0].set_ylabel("number of test images")
axes[0].legend()
plt.show()

fit 1034, validation 259, test normal 324, test nines 36
autoencoder      threshold 0.0261  false alarms 0.040  9s detected 0.583  AUC 0.952
PCA Q statistic  threshold 0.0395  false alarms 0.077  9s detected 0.222  AUC 0.791

这张表要按顺序读出三点。正常测试数字上的误报率接近阈值所设定的 5%，这证实验证集划分起了作用（正常测试图像只有 324 张，差一两个百分点只是噪声）。检出率是异常本身的性质，而不是阈值的性质：尽管 AUC 为 0.95，数字 9 也只有约 58% 被捕获。AUC 是在所有阈值上的平均，包括那些没人会用的阈值，它对你选定的那个阈值什么也没说。而 PCA 监测器明显更差，这正是证明网络有其价值的比较。直方图显示了造成漏检的重叠：许多 9 的误差落在正常范围之内，因为它们与模型能重构得很好的数字相似。

如果你在测试误差上选阈值以捕获更多的 9，报告的检出率就只是这一选择的产物。在验证数据上选择，在测试数据上报告。

你应该看到什么

• 在相同编码维数下，自编码器胜过 PCA：d_z = 2 时约好 30%，d_z = 8 时好一倍以上，因为数字图像位于一个弯曲的流形上。
• VAE 的二维编码比自编码器的重叠更多（5 近邻准确率约 0.75，对比 0.84），因为 KL 项把每个编码都拉向 \mathcal{N}(0, \mathbf{I})。作为交换，VAE 网格的每一点都解码为一个像样的数字，而自编码器的网格在簇与簇之间有不合理的区域。
• KL 权重控制坍塌。\beta = 4 时每个维度的 KL 都是 0.00，解码器输出一个平均数字。平方误差之和的损失（\sigma_x^2 = 1/2）只让八个维度中约三个保持活跃。
• AUC 为 0.95 并不意味着 95% 的检出率。在误报约 5% 的阈值下，约 58% 的 9 被捕获。PCA 监测器在这两方面都更差。

动手试试

1. KL 预热。在 d_z = 8 的 VAE 中，让 \beta 在前 50 个轮次内从 0 线性升到 1，并统计活跃单元：预计会从 6 个升到 8 个，多出的两个只携带很少的 KL（这一扩展的一次运行得到 8 个活跃单元，KL 为 3.48 奈特，两个新维度分别为 0.01 和 0.04）。若改为升到 \beta = 4，模型仍然坍塌（KL 0.01，没有活跃单元）。预热能修复由优化路径造成的坍塌，却修复不了作为目标函数最优解的坍塌。
2. 二维中的簇。从三个二维高斯簇生成 3,000 个点，均值为 (-2, 0)、(2, 0)、(0, 2.5)，标准差为 0.3。用 d_x = 2、d_z = 2 和高斯解码器训练 VAE，画出按簇着色的潜在均值。然后设 d_z = 1，看各个簇是否仍然分开。
3. 另一种异常。改为留出数字 8 而不是 9，再改为留出数字 0，阈值规则不变。其中一个的检出率会急剧上升。对这个模型来说，哪些数字是难以检出的异常？为什么这取决于正常类中剩下了哪些数字？

目标。你在小到可以直接观察的数据上实现第 5 节和第 6 节的各个部件：采用余弦调度的前向加噪过程、用简化损失训练的噪声预测网络、从反向过程进行的祖先采样（ancestral sampling）、无分类器引导，以及步数更少的确定性 DDIM 采样器。数据是 make_moons 生成的两个交错的月牙，这是一个二维分布，其质量可以用最近邻距离来测量，而不必靠肉眼判断。你会看到样本的结构在反向过程的后期才出现，看到引导以多样性换取保真度，看到步数减少时质量下降。数据是合成的，无需下载，本实验在笔记本电脑 CPU 上约需两分钟（所示运行用了 105 秒）：网络有 37,858 个参数，其中训练占 95 秒。在第一个代码块中设 QUICK = True，就只训练 8,000 步而不是 20,000 步，约一分钟即可完成。打印出的数字可能与你的结果在最后几位上不同。

步骤 1：噪声调度，以及为什么要给它设上限

第 5 节的前向过程在 T 步内把数据变成噪声。记 \alpha_t = 1 - \beta_t、\bar\alpha_t = \prod_{s \le t}\alpha_s，加噪后的点为 \mathbf{x}_t = \sqrt{\bar\alpha_t}\,\mathbf{x}_0 + \sqrt{1 - \bar\alpha_t}\,\boldsymbol\epsilon。 调度采用 Nichol 和 Dhariwal 的余弦调度：\bar\alpha_t 正比于 \cos^2\!\big(\tfrac{t/T + s}{1 + s}\cdot\tfrac{\pi}{2}\big)，其中 s = 0.008，并归一化使 \bar\alpha_0 = 1，每个 \beta_t = 1 - \bar\alpha_t/\bar\alpha_{t-1}。本实验用 T = 200，而不是论文中的 1,000，因为双月数据需要的步数少得多。

与论文唯一的不同是一个上限：\beta_t \le 0.5，而不是 0.999。不设上限的调度以 \beta_T = 1 和 \bar\alpha_T = 0 结束。反向步要除以 \sqrt{\alpha_t}，所以在最后一步，\alpha_T = 1 - \beta_T 极小，网络犯下的任何误差都会被乘以 1/\sqrt{\alpha_T}。\beta_T = 0.999 时这个因子是 31.6，有上限时是 1.41。第一个代码块打印了两种情况下的因子，让你看到这一影响有多大。上限使 \bar\alpha_T 停在一个很小的正数上，而采样器从 \mathcal{N}(0, \mathbf{I}) 出发，这样起始分布就有了极其微小的偏差；对这份数据，这一误差看不出来。

import time

import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from sklearn.datasets import make_moons
from sklearn.neighbors import KNeighborsClassifier, NearestNeighbors

np.random.seed(0)
torch.manual_seed(0)
torch.set_num_threads(2)

QUICK = False  # True: 8,000 training steps instead of 20,000
T = 200


def cosine_alpha_bar(T, s=0.008):
    """alpha-bar_t for t = 0..T from the cosine schedule, with alpha-bar_0 = 1."""
    t = np.arange(T + 1) / T
    f = np.cos((t + s) / (1 + s) * np.pi / 2) ** 2
    return f / f[0]


ab_cos = cosine_alpha_bar(T)
beta_uncapped = 1 - ab_cos[1:] / ab_cos[:-1]  # entry t-1 is beta_t
beta = np.minimum(beta_uncapped, 0.5)  # the cap
alpha = 1 - beta
alpha_bar = np.concatenate([[1.0], np.cumprod(alpha)])  # alpha_bar[t] for t = 0..T

print(f"alpha_bar_1 {alpha_bar[1]:.5f}  alpha_bar_100 {alpha_bar[100]:.4f}  "
      f"alpha_bar_200 {alpha_bar[200]:.2e}")
print("last three uncapped betas:", np.round(beta_uncapped[-3:], 4))
print(f"error factor 1/sqrt(alpha_T): capped {1 / np.sqrt(alpha[-1]):.2f}, "
      f"beta_T = 0.999 gives {1 / np.sqrt(1 - 0.999):.1f}")

plt.figure(figsize=(6, 3.5))
plt.plot(np.arange(T + 1), alpha_bar)
plt.xlabel("step t")
plt.ylabel(r"$\bar\alpha_t$ (fraction of signal variance left)")
plt.title("Cosine noise schedule, T = 200, beta capped at 0.5")
plt.show()

alpha_bar_1 0.99975  alpha_bar_100 0.4938  alpha_bar_200 6.83e-05
last three uncapped betas: [0.5555 0.75   1.    ]
error factor 1/sqrt(alpha_T): capped 1.41, beta_T = 0.999 gives 31.6

信号起初下降得慢，后来下降得快：t = 50 时还剩 0.85 的方差，t = 100 时剩一半，t = 150 时剩 0.14，最后五十步把剩下的部分降到几乎为零。选用余弦调度正是为了这种形状：它在中等噪声水平上花很多步，此时数据的结构部分可见，而不是一冲而过。下面的训练和采样始终使用设了上限的值，所以前向过程与反向过程彼此一致。

步骤 2：画出前向过程

闭式解让你可以直接跳到任意一步。每个点用一次固定的噪声抽样 \boldsymbol\epsilon，并在每个 t 上重复使用，这样每个点都沿着一条直线从原处滑向噪声，各个子图之间也可以比较。数据先做标准化，使每个坐标的均值为零、方差为一，这样终态 \mathcal{N}(0, \mathbf{I}) 与数据的尺度相同。标准化使用训练集自身的均值和标准差；同样的数值也会用于评估集。

X_raw, y_raw = make_moons(n_samples=10000, noise=0.05, random_state=0)
mean, std = X_raw.mean(0), X_raw.std(0)
data = torch.tensor((X_raw - mean) / std, dtype=torch.float32)
labels = torch.tensor(y_raw)
X_ref_raw, y_ref = make_moons(n_samples=5000, noise=0.05, random_state=1)
X_ref = ((X_ref_raw - mean) / std).astype(np.float32)  # reference set for evaluation
print("standardised data: mean", np.round(data.mean(0).numpy(), 3),
      " std", np.round(data.std(0).numpy(), 3))

ab = torch.tensor(alpha_bar, dtype=torch.float32)


def noise_to(x0, t, eps):
    """Closed-form forward process: x_t from x_0 in one step (t is a tensor of step indices)."""
    a = ab[t][:, None]
    return a.sqrt() * x0 + (1 - a).sqrt() * eps


eps_fixed = torch.randn(1500, 2)
x0_show, y_show = data[:1500], labels[:1500].numpy()
steps = (0, 50, 100, 150, 175, 200)
fig, axes = plt.subplots(1, 6, figsize=(16, 2.9), sharex=True, sharey=True)
for ax, t in zip(axes, steps):
    xt = noise_to(x0_show, torch.full((1500,), t), eps_fixed).numpy()
    ax.scatter(xt[:, 0], xt[:, 1], c=y_show, cmap="coolwarm", s=3)
    ax.set_title(f"t = {t}, $\\bar\\alpha_t$ = {alpha_bar[t]:.3f}")
    ax.set_xlabel("$x_1$")
axes[0].set_ylabel("$x_2$")
plt.show()

standardised data: mean [-0. -0.]  std [1. 1.]

到 t = 50 时，细细的月牙已经变成两团宽大且相互重叠的云，每个月牙一团。到 t = 100 时，它们是两团严重重叠的斑块；从 t = 150 起，两个类别混在一起，只剩下一团高斯云。数据的精细几何，也就是月牙的弧线，在前向过程的早期就被破坏了。反向过程必须把它重建出来，这是步骤 5 中“结构后期才出现”这一观察的第一个提示。

步骤 3：去噪网络

网络 \boldsymbol\epsilon_\theta(\mathbf{x}_t, t, c) 预测被加入的噪声。它需要知道步数 t，因为正确答案取决于输入有多嘈杂；它还需要知道条件 c：月牙标签 0 或 1，或者第三个值 null，表示“无条件”。null 值使步骤 6 中的无分类器引导成为可能，因为同一个网络于是既给出条件预测，也给出无条件预测。

步数编码为 32 维的正弦嵌入，与 Transformer 中一样；条件由一个有三个条目的可学习 32 维嵌入编码。两个嵌入相加，其和与 \mathbf{x}_t 的 2 个坐标拼接。随后是三个 128 单元、使用 SiLU 激活的隐藏层，以及输出 2 个数的线性层。参数量为 34 \cdot 128 + 128 + 2(128 \cdot 128 + 128) + 128 \cdot 2 + 2 + 3 \cdot 32 = 37{,}858。

def time_embedding(t, dim=32):
    """Sinusoidal embedding of the step index, shape (B, dim)."""
    freqs = torch.exp(-np.log(10000.0) * torch.arange(dim // 2) / (dim // 2))
    angles = t.float()[:, None] * freqs[None, :]
    return torch.cat([angles.sin(), angles.cos()], dim=-1)


NULL = 2  # condition index meaning "no condition"


class Denoiser(nn.Module):
    def __init__(self, hidden=128, d_emb=32):
        super().__init__()
        self.cond_emb = nn.Embedding(3, d_emb)  # moon 0, moon 1, null
        self.net = nn.Sequential(
            nn.Linear(2 + d_emb, hidden), nn.SiLU(),
            nn.Linear(hidden, hidden), nn.SiLU(),
            nn.Linear(hidden, hidden), nn.SiLU(),
            nn.Linear(hidden, 2),
        )

    def forward(self, x, t, c):
        emb = time_embedding(t) + self.cond_emb(c)
        return self.net(torch.cat([x, emb], dim=-1))


model = Denoiser()
print("parameters:", sum(p.numel() for p in model.parameters()))

parameters: 37858

步骤 4：用简化损失训练

每个训练步抽取一个 batch 的 512 个干净点，为每个点抽取在 \{1, \dots, 200\} 上均匀分布的步数 t 和噪声 \boldsymbol\epsilon；构造 \mathbf{x}_t；并最小化 \boldsymbol\epsilon 与 \boldsymbol\epsilon_\theta(\mathbf{x}_t, t, c) 之间的均方误差。这就是第 5 节的目标函数，一个普通的回归。条件以 0.2 的概率被替换为 NULL，使同一个网络同时学会条件噪声和无条件噪声。使用 Adam，学习率 10^{-3}，余弦衰减到零，完成整个训练。

损失不会降到零，也不应该降到零：噪声是随机的，网络只能预测其中能由 \mathbf{x}_t 和 t 揭示的那部分。t 很小时，\mathbf{x}_t 几乎就是 \mathbf{x}_0，噪声几乎无法从单个点中恢复。t 很大时，\mathbf{x}_t 几乎全是噪声，任务就容易了。因此打印出的损失是在难度差别极大的情形上的平均，一条平坦的曲线对样本质量说明不了多少。

n_steps = 8000 if QUICK else 20000
opt = torch.optim.Adam(model.parameters(), lr=1e-3)
sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=n_steps)
torch.manual_seed(0)
t0 = time.time()
fifth = n_steps // 5
running = []
for step in range(1, n_steps + 1):
    idx = torch.randint(0, len(data), (512,))
    x0, c = data[idx], labels[idx].clone()
    c[torch.rand(512) < 0.2] = NULL  # condition dropout for classifier-free guidance
    t = torch.randint(1, T + 1, (512,))
    eps = torch.randn(512, 2)
    loss = F.mse_loss(model(noise_to(x0, t, eps), t, c), eps)
    opt.zero_grad()
    loss.backward()
    opt.step()
    sched.step()
    running.append(loss.item())
    if step % fifth == 0:
        print(f"step {step:6d}  mean loss over the last {fifth} steps "
              f"{np.mean(running[-fifth:]):.4f}")
print(f"training time {time.time() - t0:.0f} s")

step   4000  mean loss over the last 4000 steps 0.3646
step   8000  mean loss over the last 4000 steps 0.3324
step  12000  mean loss over the last 4000 steps 0.3294
step  16000  mean loss over the last 4000 steps 0.3277
step  20000  mean loss over the last 4000 steps 0.3260
training time 95 s

损失在前五分之一为 0.36，之后稳定在 0.33 附近。这是一个无法低于噪声下限的误差在所有噪声水平上的平均，所以早在样本停止改进之前，曲线就几乎平了。0.33 的损失并不能说明月牙是否清晰；接下来几步的最近邻测量才能说明。训练时间取决于机器，在繁忙的机器上可能是安静的笔记本电脑上一分半钟的两倍。

步骤 5：祖先采样，以及它产生了什么

采样从 \mathbf{x}_T \sim \mathcal{N}(0, \mathbf{I}) 开始，对 t = T, \dots, 1 应用 Ho 等人的更新（算法 2）：

最后一步取 \mathbf{z} = 0。第一项去掉预测的噪声并重新缩放；第二项加回少量新的噪声，方差为 \beta_t，这使链条成为从反向过程中抽取的真正样本，而不是一次确定性的滑动。下面的函数接收一个条件和一个引导强度（guidance scale）w；这一步使用 null 条件，所以 w 不起作用，步骤 6 会同时用到两者。

样本用最近邻距离评分，单位是标准化后的单位。从每个样本到 5,000 点参考集中最近点的平均距离是一个类似精确率的量：样本落在月牙上时它就小。从每个参考点到最近样本的平均距离是一个类似召回率的量：样本覆盖了月牙的每一部分时它就小。两把标尺为它们提供校准：从真实分布新抽取的 2,000 个点，代表任何采样器所能达到的最好结果；以及从 \mathcal{N}(0, \mathbf{I}) 抽取的 2,000 个点，即起点。

@torch.no_grad()
def eps_hat(x, t_int, c, w):
    """Noise prediction; with a condition c in {0, 1}, classifier-free guidance of scale w."""
    t = torch.full((len(x),), t_int, dtype=torch.long)
    e_null = model(x, t, torch.full((len(x),), NULL))
    if c == NULL:
        return e_null
    e_cond = model(x, t, torch.full((len(x),), c))
    return e_null + w * (e_cond - e_null)


beta_t, alpha_t, ab_t = (torch.tensor(a, dtype=torch.float32) for a in (beta, alpha, alpha_bar))


@torch.no_grad()
def ancestral(n, c=NULL, w=0.0, seed=0, snapshots=()):
    """Ho et al. Algorithm 2. Returns x_0 and {t: x_t before the update at step t}."""
    gen = torch.Generator().manual_seed(seed)
    x = torch.randn(n, 2, generator=gen)
    snaps = {}
    for t in range(T, 0, -1):
        if t in snapshots:
            snaps[t] = x.clone().numpy()
        e = eps_hat(x, t, c, w)
        x = (x - beta_t[t - 1] / (1 - ab_t[t]).sqrt() * e) / alpha_t[t - 1].sqrt()
        if t > 1:
            x = x + beta_t[t - 1].sqrt() * torch.randn(n, 2, generator=gen)
    return x.numpy(), snaps


def nn_dist(a, b):
    """Mean distance from each row of a to its nearest row of b."""
    return float(NearestNeighbors(n_neighbors=1).fit(b).kneighbors(a)[0].mean())


fresh_raw, _ = make_moons(n_samples=2000, noise=0.05, random_state=2)
fresh = ((fresh_raw - mean) / std).astype(np.float32)
noise_pts = np.random.default_rng(0).standard_normal((2000, 2)).astype(np.float32)
for name, pts in (("fresh data", fresh), ("N(0, I) noise", noise_pts)):
    print(f"{name:14s} precision-like {nn_dist(pts, X_ref):.3f}   "
          f"recall-like {nn_dist(X_ref, pts):.3f}")

t0 = time.time()
snap_steps = (200, 150, 100, 50, 20, 5)
x_gen, snaps = ancestral(2000, snapshots=snap_steps)
print(f"unconditional samples  precision-like {nn_dist(x_gen, X_ref):.3f}   "
      f"recall-like {nn_dist(X_ref, x_gen):.3f}   ({time.time() - t0:.1f} s)")
print("precision-like distance of x_t at t =", snap_steps, ":",
      [round(nn_dist(snaps[t], X_ref), 3) for t in snap_steps])

fig, axes = plt.subplots(1, 7, figsize=(18, 2.8), sharex=True, sharey=True)
for ax, t in zip(axes, snap_steps):
    ax.scatter(snaps[t][:, 0], snaps[t][:, 1], s=2)
    ax.set_title(f"$x_t$ at t = {t}")
    ax.set_xlabel("$x_1$")
axes[-1].scatter(x_gen[:, 0], x_gen[:, 1], s=2, color="tab:green")
axes[-1].set_title("final sample $x_0$")
axes[-1].set_xlabel("$x_1$")
axes[0].set_ylabel("$x_2$")
plt.show()

fresh data     precision-like 0.013   recall-like 0.021
N(0, I) noise  precision-like 0.242   recall-like 0.051
unconditional samples  precision-like 0.022   recall-like 0.022   (0.5 s)
precision-like distance of x_t at t = (200, 150, 100, 50, 20, 5) : [0.246, 0.234, 0.212, 0.134, 0.057, 0.027]

先读两把标尺。新抽取的数据到参考集的距离约为 0.013，参考集到它们的距离约为 0.021，这是由样本量和月牙本身的噪声决定的下限：没有哪个采样器能低于它们。噪声的两个距离为 0.24 和 0.05。模型的样本在两项上都接近下限：它们落在月牙上，并且覆盖了两个月牙。六张快照显示了这是在哪里发生的。\mathbf{x}_t 的类精确率距离在 t = 200 时仍约为 0.25，与纯噪声相同，直到 t = 100 只缓慢下降（0.212），t = 50 时为 0.134，并在最后五十步内降到最终的 0.022。结构在后期才出现：t = 100 和 t = 50 的快照是没有形状的斑块，只隐约可见两个月牙，月牙只在 t = 50 到 t = 5 之间才变清晰。这就是最后几步对质量最重要的原因，也是步骤 7 中减少步数代价高昂的原因。

设 QUICK = True 时，预计类精确率距离约为 0.04，类召回率距离约为 0.025（一次 QUICK = True 的运行得到 0.039 和 0.025）：月牙明显更模糊，月牙附近有零散的点。模型训练的步数更少，而它在反向链条中走的每一步都会累积它的误差。

步骤 6：无分类器引导

条件模型从 p(\mathbf{x} \mid c) 中采样。引导使它更集中。更新中使用的噪声预测为

其中 \varnothing 是 null 条件。w = 0 时它就是无条件模型；w = 1 时它就是条件模型；w > 1 时，它沿着从“任意”指向“类别 c”的方向走得比单独的条件模型更远。这里引导强度的写法与代码一致，w = 1 表示普通的条件模型。Ho 和 Salimans 用 (1 + w) 写同一件事，所以他们的 w = 0 就是本实验的 w = 1。

代码块对 w = 0, 1, 3, 7 各采样 1,000 个类别 0 的点。一个在带标签参考集上拟合的 15 近邻分类器判断每个样本落在哪个月牙上；落在月牙 0 上的比例衡量条件被遵守的程度。类精确率距离使用所有参考点。类召回率距离只使用类别 0 的参考点，因为被采样的正是这个分布。

knn = KNeighborsClassifier(15).fit(X_ref, y_ref)
ref_class0 = X_ref[y_ref == 0]
guided = {}
for w in (0, 1, 3, 7):
    x_w, _ = ancestral(1000, c=0, w=float(w), seed=10 + w)
    guided[w] = x_w
    frac = float((knn.predict(x_w) == 0).mean())
    print(f"w = {w}: fraction on moon 0 {frac:.2f}   precision-like {nn_dist(x_w, X_ref):.3f}"
          f"   recall-like (class-0 points) {nn_dist(ref_class0, x_w):.3f}")

fig, axes = plt.subplots(1, 4, figsize=(15, 3.4), sharex=True, sharey=True)
for ax, w in zip(axes, guided):
    ax.scatter(X_ref[:, 0], X_ref[:, 1], s=1, color="lightgrey")
    ax.scatter(guided[w][:, 0], guided[w][:, 1], s=3, color="tab:red")
    ax.set_title(f"class 0, guidance w = {w}")
    ax.set_xlabel("$x_1$")
axes[0].set_ylabel("$x_2$")
plt.show()

w = 0: fraction on moon 0 0.50   precision-like 0.023   recall-like (class-0 points) 0.032
w = 1: fraction on moon 0 1.00   precision-like 0.014   recall-like (class-0 points) 0.021
w = 3: fraction on moon 0 1.00   precision-like 0.013   recall-like (class-0 points) 0.029
w = 7: fraction on moon 0 1.00   precision-like 0.020   recall-like (class-0 points) 0.043

w = 0 时条件被忽略，一半样本落在每个月牙上：这就是无条件模型。它到类别 0 各点的类召回率距离较高（0.032），因为它只有一半样本靠近这些点。w = 1 时，每个样本都落在所要求的月牙上，类召回率距离也达到最好，为 0.021。落在月牙上的比例不可能超过 1.00，但其他数字还在变化。类召回率距离随 w 增大（w = 3 时为 0.029，w = 7 时为 0.043），因为样本集中到月牙最密的部分，它的两端无人问津，如图所示。w = 7 时类精确率距离也变大了（0.020，而 w = 3 时为 0.013），因为更强的推力冲过了头：有些样本越过月牙的左上端，离开了月牙。这是引导最纯粹形式的权衡。对条件的保真度是用多样性换来的，而超过某个强度后，代价还包括更差的拟合。

步骤 7：DDIM 与步数的代价

祖先采样每个样本需要 200 次网络计算。DDIM（Song 等人）复用同一个训练好的网络，采用一种可以跳步的确定性更新。每一步它先由噪声预测估计干净的点，

这是前向公式的逆，然后用同一个预测噪声把它重新加噪到下一个更低的噪声水平 t'：

这里不注入新的噪声，所以从 \mathbf{x}_T 到 \mathbf{x}_0 的映射是确定性的。 由于更新只需要 \bar\alpha_t 和 \bar\alpha_{t'}，t' 不必是 t - 1：从 200 到 0 均匀取 K 步，就得到一个无需重新训练的 K 步采样器。代码块对 K = 200, 50, 20, 10, 5, 2, 1 测量无条件样本，以及每种情况所用的时间。K = 1 那一行是在 t = 200 处只做一次计算得到的 \hat{\mathbf{x}}_0。

@torch.no_grad()
def ddim(n, K, seed=0):
    gen = torch.Generator().manual_seed(seed)
    x = torch.randn(n, 2, generator=gen)
    ts = np.round(np.linspace(T, 0, K + 1)).astype(int)
    for t, t_next in zip(ts[:-1], ts[1:]):
        e = eps_hat(x, int(t), NULL, 0.0)
        x0_hat = (x - (1 - ab_t[t]).sqrt() * e) / ab_t[t].sqrt()
        x = ab_t[t_next].sqrt() * x0_hat + (1 - ab_t[t_next]).sqrt() * e
    return x.numpy()


print("   K  precision-like  recall-like  seconds")
ddim_samples = {}
for K in (200, 50, 20, 10, 5, 2, 1):
    t0 = time.time()
    xs = ddim(2000, K)
    ddim_samples[K] = xs
    print(f"{K:4d}  {nn_dist(xs, X_ref):13.3f}  {nn_dist(X_ref, xs):11.3f}  {time.time() - t0:7.2f}")

fig, axes = plt.subplots(1, 4, figsize=(14, 3.3), sharex=True, sharey=True)
for ax, K in zip(axes, (50, 10, 5, 1)):
    xs = ddim_samples[K]
    outside = int((np.abs(xs).max(1) > 3).sum())
    ax.scatter(xs[:, 0], xs[:, 1], s=2)
    ax.set_xlim(-3, 3)
    ax.set_ylim(-3, 3)
    ax.set_title(f"DDIM, K = {K} steps ({outside} of 2000 outside the box)")
    ax.set_xlabel("$x_1$")
axes[0].set_ylabel("$x_2$")
plt.show()

   K  precision-like  recall-like  seconds
 200          0.022        0.023     0.84
  50          0.022        0.024     0.21
  20          0.025        0.029     0.12
  10          0.035        0.042     0.07
   5          0.046        0.067     0.03
   2          0.152        0.120     0.02
   1          1.714        0.142     0.02

二十到五十步 DDIM 给出的样本接近 200 步祖先采样器的样本，而网络计算次数只有它的四分之一到十分之一；时间一列几乎与 K 成正比地下降。少于约十步时，质量迅速下降。一步则是另一种失败。在 t = 200 处信号权重为 \sqrt{\bar\alpha_{200}} = 0.008，所以估计 \hat{\mathbf{x}}_0 要把网络的输出除以 0.008，即放大 121 倍，\boldsymbol\epsilon_\theta 中的任何误差都会被放大这么多倍。因此一步得到的样本散布在远离数据的地方，类精确率距离远大于 1，尽管它们的类召回率距离只有 0.14：数据区域确实被覆盖了，但覆盖它的那团云还覆盖了大得多的范围。扩散的代价在于它的步数，而这些步数换来的，正是你在步骤 5 中看到的对结构的逐步确定。

你应该看到什么

• 结构在反向过程的后期才出现。中间样本直到约 t = 50 都没有形状（类精确率距离 0.134，而噪声为 0.246），月牙在 t = 50 到 t = 5 之间形成。
• 无条件样本与数据的接近程度几乎与新抽取的数据相当（约 0.022，对比 0.013），并覆盖了两个月牙（类召回率距离 0.022，对比 0.021）。
• 引导以多样性换取保真度。w = 1 已经让每个样本都落在所要求的月牙上。w = 3 和 w = 7 把样本挤向月牙最密的部分，到 w = 7 时类召回率距离翻了一倍；w = 7 时，由于样本冲过了头，类精确率距离开始变差。
• 20 到 50 步的 DDIM 接近 200 步的祖先采样器。少于约 10 步时质量迅速下降，而单步只返回一个远离数据、被抹开的估计。
• \beta_t 的上限是一道保险。按论文的上限 0.999，第一个反向步会把网络的误差乘以 31.6，而上限为 0.5 时只乘以 1.41。这是否造成损害取决于具体的运行：在规划本模块时做的一个原型中，w = 7 的样本发散了（类精确率距离 1.34），而在当前环境中用 0.999 上限重新训练的本实验副本，在 w = 7 时得到 0.019，与设了上限的模型一样好。上限以零代价消除了这一风险。

动手试试

1. 恢复原上限。把上限设为 0.999，重复步骤 6。然后保持 0.999 的上限，但计算 \hat{\mathbf{x}}_0，把它裁剪到 [-3, 3]，再改用由 \hat{\mathbf{x}}_0 得到的后验均值更新 \tilde\mu_t。最初的实现正是这样在不设上限的调度下存活下来的。
2. 同一数据上的 GAN。在同样标准化的月牙数据上，把第 4 节的非饱和 GAN（生成器和判别器：三个 128 单元隐藏层的多层感知机，Adam，学习率 10^{-3}，\beta = (0.5, 0.999)）训练 6,000 步。把它的类精确率和类召回率距离与扩散模型比较。它每个样本只需一次前向传播；扩散模型需要 200 次。
3. 线性调度。把调度换成 Ho 等人的线性范围，\beta_t 从 10^{-4} 到 0.02，仍取 T = 200，它结束于 \bar\alpha_T = 0.13（用它重新训练的本实验副本得到类精确率距离 0.020、类召回率距离 0.022）。在这些二维标准化数据上，损害很小。解释为什么二维标准化数据掩盖了一个对图像很重要的问题：\bar\alpha_T = 0.13 在起始样本中留下了什么？

目标。你从边列表出发，用约三十行代码构建消息传递，并把它用于一个答案精确已知的任务：找出故障树的单点故障，即仅凭自身失效就能导致顶事件发生的基本事件。你编写一个随机故障树生成器和一个给出真实标签的标注器，测量模型必须胜过的两个基线，训练深度逐渐增加的图卷积网络（第 7 节），并亲眼看到第 8 节所说的深度限制：按距顶事件的距离统计的准确率、无需任何训练即可测量的过平滑，以及用残差连接进行的修复。最后，你把对称邻接换成一个方向感知层，它利用了朴素 GCN 丢掉的那一点结构。数据是合成的，本实验不使用任何图神经网络库；它在笔记本电脑 CPU 上约需三分钟（所示运行用了 168 秒），需要 NumPy、PyTorch 和 matplotlib。它让 PyTorch 在单线程上运行：在多线程上，消息传递中的 scatter-add 求和顺序每次运行都会变，准确率也随之而变。即使在单线程上，准确率也会随 PyTorch 版本变动一两个百分点，所以要把它们当作现象来读，而不是当作确切的数字。

步骤 1：故障树生成器及其标签

故障树由门（或门 OR：任一输入失效则输出失效；与门 AND：只有全部输入失效时才失效）和基本事件（即叶节点）组成。如果一个基本事件到顶事件路径上的每个门都是或门，它就是单点故障：此时它自身的失效会一路向上传播。路径上任何地方只要有一个与门，就意味着该门的其他输入也必须失效。

生成器遵循固定的配方。顶事件是一个门，以 0.6 的概率为或门，否则为与门。每个门有两到四个输入，均匀抽取。深度为 0 或 1 的门，其输入本身以 0.6 的概率是一个门；深度为 2 时概率为 0.35；深度为 3 的门只以基本事件为输入，所以没有哪棵树在顶事件之下深于四层。每个非顶层的门以 0.6 的概率为或门。每个节点有四个特征：其类型（或门、与门、基本事件）的独热编码，以及一个标记顶事件的标志位。门不参与评分；基本事件的标签为 1 表示单点故障。

实验用一个设定了种子的生成器生成 300 棵树，200 棵用于训练、100 棵用于测试，并且按树划分，使测试树是新的图，而不是某棵训练树中的节点。工程模型要用在新的模型上，而不是旧模型的新节点上（第 7 节）。

import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

np.random.seed(0)
torch.manual_seed(0)
torch.set_num_threads(1)    # scatter-adds sum in a fixed order only on one thread

OR, AND, BASIC = 0, 1, 2


def make_tree(rng):
    """One random fault tree as arrays: node type, parent index (-1 for the top), depth."""
    types, parent, depth = [], [], []

    def add(kind, par, d):
        types.append(kind)
        parent.append(par)
        depth.append(d)
        return len(types) - 1

    top = add(OR if rng.random() < 0.6 else AND, -1, 0)
    open_gates = [top]
    while open_gates:
        g = open_gates.pop()
        d = depth[g]
        p_gate = 0.6 if d <= 1 else (0.35 if d == 2 else 0.0)
        for _ in range(rng.integers(2, 5)):  # 2 to 4 inputs
            if rng.random() < p_gate:
                open_gates.append(add(OR if rng.random() < 0.6 else AND, g, d + 1))
            else:
                add(BASIC, g, d + 1)
    return {"type": np.array(types), "parent": np.array(parent), "depth": np.array(depth)}


def label_tree(tree):
    """1 for a basic event whose gates up to the top are all OR, 0 for the other basic
    events, -1 for gates (not scored)."""
    label = np.full(len(tree["type"]), -1)
    for v in np.where(tree["type"] == BASIC)[0]:
        label[v], u = 1, tree["parent"][v]
        while u != -1:
            if tree["type"][u] == AND:
                label[v] = 0
            u = tree["parent"][u]
    return label


rng = np.random.default_rng(0)
trees = [make_tree(rng) for _ in range(300)]
for tree in trees:
    tree["label"] = label_tree(tree)
train_trees, test_trees = trees[:200], trees[200:]

sizes = [len(t["type"]) for t in trees]
print(f"nodes per tree: mean {np.mean(sizes):.1f}, min {min(sizes)}, max {max(sizes)}")
for name, group in (("train", train_trees), ("test", test_trees)):
    n_nodes = sum(len(t["type"]) for t in group)
    lab = np.concatenate([t["label"] for t in group])
    dep = np.concatenate([t["depth"] for t in group])
    basic = lab >= 0
    counts = {d: int(((dep == d) & basic).sum()) for d in (1, 2, 3, 4)}
    print(f"{name}: {n_nodes} nodes, {basic.sum()} basic events, "
          f"{lab[basic].mean():.1%} single points of failure; by depth {counts}")

nodes per tree: mean 28.1, min 3, max 92
train: 5579 nodes, 3766 basic events, 22.1% single points of failure; by depth {1: 252, 2: 400, 3: 1169, 4: 1945}
test: 2850 nodes, 1934 basic events, 21.9% single points of failure; by depth {1: 123, 2: 193, 3: 640, 4: 978}

在训练树和测试树中，大约每五个基本事件中就有一个是单点故障，所以一个从不预测单点故障的模型大约五次中能对四次：这就是要胜过的基线。大多数基本事件位于顶事件之下三到四层，而它们恰恰需要最多的信息才能标注。换一个种子或换一种抽取顺序，计数会有些不同（各项比例变动一两个百分点），所以不要指望用不同的生成器得到这些数字。

步骤 2：在一棵可以手工验证的树上检验标注器

没人检验过的标注器，是错误基准最常见的来源。第 8 节的冷却系统故障树有十二个节点，它的答案可以直接从图上读出。顶事件“失去冷却”是三个输入的或门：门 G1“所有泵失效”（三个泵的与门）；门 G2“流道堵塞”（一个阀门、一根管道和门 G3“两个控制器都失效”的或门，G3 是两个控制器的与门）；以及基本事件 E1，即电源。因此单点故障是直接相连的 E1，以及经由或门 G2 的阀门 E5 和管道 E6。泵和控制器都位于与门之下。

NAMES = ["TOP", "G1", "G2", "E1", "E2", "E3", "E4", "E5", "E6", "G3", "E7", "E8"]
cooling = {
    "type": np.array([OR, AND, OR, BASIC, BASIC, BASIC, BASIC, BASIC, BASIC, AND, BASIC, BASIC]),
    "parent": np.array([-1, 0, 0, 0, 1, 1, 1, 2, 2, 2, 9, 9]),
}
cooling["depth"] = np.array([0, 1, 1, 1, 2, 2, 2, 2, 2, 2, 3, 3])
cooling["label"] = label_tree(cooling)

found = [NAMES[v] for v in np.where(cooling["label"] == 1)[0]]
print("single points of failure found by the labeller:", found)
assert found == ["E1", "E5", "E6"], "the labeller disagrees with the hand answer"

single points of failure found by the labeller: ['E1', 'E5', 'E6']

这个断言就是这一步的意义所在。如果它失败了，后面的一切都不可信。

步骤 3：两个基线

模型在测试基本事件上的准确率，若没有某个平凡方法的准确率作对照，就毫无意义。这里有两个平凡的预测器。多数类：对每个事件都说“不是单点故障”。以及一条听起来正确的规则，“事件自己的门是或门”：如果正上方的门是或门，这个事件看起来就像单点故障。恰好在更高处某个地方有与门时，这条规则就错了。两者都在测试基本事件上计算。

def basic_events(group):
    """Concatenate trees; return labels, own-gate types and depths of the basic events."""
    lab = np.concatenate([t["label"] for t in group])
    own_gate = np.concatenate([t["type"][np.maximum(t["parent"], 0)] for t in group])
    dep = np.concatenate([t["depth"] for t in group])
    m = lab >= 0
    return lab[m], own_gate[m], dep[m]


lab_te, gate_te, dep_te = basic_events(test_trees)
majority = np.mean(lab_te == 0)
own_gate_or = np.mean((gate_te == OR).astype(int) == lab_te)
print(f"majority class (never a single point of failure): {majority:.3f}")
print(f"'own gate is OR' rule:                            {own_gate_or:.3f}")

majority class (never a single point of failure): 0.781
'own gate is OR' rule:                            0.616

这条看似合理的规则比预测多数类还差。许多自己的门是或门的事件，在更高处有一个与门，所以这条规则会产生大量误报。下面的每个模型都必须胜过这两个数字，一个只胜过这条规则的模型，对单点故障什么也没学到。

步骤 4：边列表上的消息传递

消息沿边发送，整个层就是一次 scatter-add。图存储为边列表：两个数组 i 和 j，表示存在一条从 j 到 i 的边（一条从 j 发送到 i 的消息）。GCN 的传播矩阵是 \hat{\mathbf{A}} = \tilde{\mathbf{D}}^{-1/2}(\mathbf{A} + \mathbf{I})\tilde{\mathbf{D}}^{-1/2}， 所以边列表包含每条树边的两个方向，以及每个节点的一个自环，每条边的权重为 1/\sqrt{\tilde d_i \tilde d_j}，其中 \tilde d 是节点的邻居数加上它自己。于是传播特征 \mathbf{H} 就是

out[i] = sum over edges (i, j) of  w_ij * H[j]

index_add_ 在 O(|\mathcal{E}|\,d) 时间内算出它，从不构造 n \times n 矩阵。代码为任意一组树构建一个 Graph（这些树合并成一张大图，树与树之间没有边），并在十二节点的树上把边列表传播与稠密公式相核对。

class Graph:
    """A batch of trees as one disconnected graph, with everything a model needs."""

    def __init__(self, group):
        offset, parts = 0, []
        for t in group:
            n = len(t["type"])
            par = np.where(t["parent"] >= 0, t["parent"] + offset, -1)
            parts.append((t["type"], par, t["depth"], t["label"]))
            offset += n
        types = np.concatenate([p[0] for p in parts])
        parent = np.concatenate([p[1] for p in parts])
        self.n = len(types)
        onehot = np.eye(3, dtype=np.float32)[types]
        top_flag = (parent == -1).astype(np.float32)[:, None]
        self.x = torch.from_numpy(np.concatenate([onehot, top_flag], axis=1))
        self.y = torch.from_numpy(np.concatenate([p[3] for p in parts]))
        self.depth = np.concatenate([p[2] for p in parts])
        self.basic = self.y >= 0
        # directed edges child -> parent, used by the direction-aware layer of Step 7
        child_idx = np.where(parent >= 0)[0]
        self.child = torch.from_numpy(child_idx)
        self.par = torch.from_numpy(parent[child_idx])
        n_children = np.bincount(parent[child_idx], minlength=self.n)
        self.n_children = torch.from_numpy(np.maximum(n_children, 1).astype(np.float32))
        # symmetric edge list with self-loops and GCN weights
        loops = np.arange(self.n)
        self.i = torch.from_numpy(np.concatenate([child_idx, parent[child_idx], loops]))
        self.j = torch.from_numpy(np.concatenate([parent[child_idx], child_idx, loops]))
        deg = np.bincount(self.i.numpy(), minlength=self.n).astype(np.float32)  # includes loop
        self.w = torch.from_numpy(1 / np.sqrt(deg[self.i.numpy()] * deg[self.j.numpy()]))


def propagate(g, H):
    """A_hat @ H from the edge list: one scatter-add."""
    return torch.zeros_like(H).index_add_(0, g.i, g.w[:, None] * H[g.j])


g_cool = Graph([cooling])
n = g_cool.n
A = torch.zeros(n, n)
A[g_cool.child, g_cool.par] = 1.0
A = A + A.T + torch.eye(n)  # A + I
d_inv_sqrt = A.sum(1).pow(-0.5)
A_hat_dense = d_inv_sqrt[:, None] * A * d_inv_sqrt[None, :]

H = torch.randn(n, 5)
diff = (propagate(g_cool, H) - A_hat_dense @ H).abs().max().item()
print(f"edge-list propagation against the dense matrix: max abs difference {diff:.1e}")
print("row sums of A_hat on the cooling tree (not 1: it is not a mean):",
      np.round(A_hat_dense.sum(1).numpy()[:4], 3))

g_train, g_test = Graph(train_trees), Graph(test_trees)

edge-list propagation against the dense matrix: max abs difference 1.2e-07
row sums of A_hat on the cooling tree (not 1: it is not a mean): [1.051 1.372 1.28  0.854]

两者吻合到 float32 的舍入误差。行和不是 1，因为对称归一化按两端的度来给边加权，而不是只按接收方的度；正是这一点使反复传播不会让特征爆炸或缩小（第 7 节）。

步骤 5：深度逐渐增加的 GCN

模型是第 7 节的 GCN。一个线性层把四个特征映射为 32 个数；L 个层各自用自己的 32 \times 32 矩阵计算 \mathbf{H} \leftarrow \mathrm{ReLU}(\hat{\mathbf{A}}\mathbf{H}\mathbf{W})；一个线性层把每个节点映射为两个类别得分。训练采用全 batch 方式，因为整个训练集就是一张几千个节点的图：Adam，学习率 10^{-2}，权重衰减 5 \cdot 10^{-4}，200 个轮次，交叉熵损失只在训练基本事件上计算，因为门没有标签。同一个 fit 辅助函数训练本实验中的每个模型。表格给出 L = 1, 2, 3, 4, 6, 8, 12, 16 时的总体测试准确率以及按事件深度统计的准确率。

def accuracy_by_depth(pred, g):
    """Accuracy on the basic events, overall and for each depth 1-4."""
    ok = (pred == g.y).numpy()
    out = [ok[g.basic.numpy()].mean()]
    for d in (1, 2, 3, 4):
        out.append(ok[g.basic.numpy() & (g.depth == d)].mean())
    return np.array(out)


def fit(model, g, epochs=200, lr=1e-2, weight_decay=5e-4):
    opt = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=weight_decay)
    for _ in range(epochs):
        loss = F.cross_entropy(model(g)[g.basic], g.y[g.basic])
        opt.zero_grad()
        loss.backward()
        opt.step()
    return model


def evaluate(model, g):
    model.eval()
    with torch.no_grad():
        pred = model(g).argmax(1)
    return accuracy_by_depth(pred, g)


class GCN(nn.Module):
    def __init__(self, L, hidden=32, residual=False):
        super().__init__()
        self.inp = nn.Linear(4, hidden)
        self.layers = nn.ModuleList(nn.Linear(hidden, hidden, bias=False) for _ in range(L))
        self.out = nn.Linear(hidden, 2)
        self.residual = residual

    def forward(self, g):
        h = self.inp(g.x)
        for lin in self.layers:
            m = torch.relu(propagate(g, lin(h)))  # ReLU(A_hat H W)
            h = h + m if self.residual else m
        return self.out(h)


gcn_acc = {}
print("  L   overall   depth1   depth2   depth3   depth4")
for L in (1, 2, 3, 4, 6, 8, 12, 16):
    torch.manual_seed(0)
    gcn_acc[L] = evaluate(fit(GCN(L), g_train), g_test)
    print(f"{L:3d}   " + "   ".join(f"{a:.3f} " for a in gcn_acc[L]))

  L   overall   depth1   depth2   depth3   depth4
  1   0.810    1.000    0.663    0.769    0.843
  2   0.832    1.000    0.886    0.769    0.843
  3   0.897    1.000    0.969    0.920    0.854
  4   0.938    1.000    0.964    0.933    0.929
  6   0.937    1.000    0.938    0.923    0.939
  8   0.959    0.984    0.959    0.958    0.956
 12   0.781    0.537    0.663    0.769    0.843
 16   0.781    0.537    0.663    0.769    0.843

先按列读这张表。至少有一层的模型都能正确标注深度 1 的事件，而每多一层就解锁下一个深度：L = 1 的模型在深度 1 上已经完美，但在深度 3 和 4 上接近多数类的比率；到 L = 4 时，每个深度都在 0.92 以上。这就是感受野。深度为 d 的基本事件上方有深度为 0, \dots, d-1 的门，最顶上的门在 d 跳之外，所以有 L 层的模型看不到足够的信息，无法准确标注深于 L 的事件。

再按行读。准确率在 L = 4 和 6 时保持在 0.94 附近，在 L = 8 时最好（0.959），因为对称层把兄弟节点和门混在一起，额外的层有助于把它们重新分开；然后准确率就断崖式下跌。12 层和 16 层的模型在每个深度上都得到多数类的准确率：它们对每个事件都预测“不是单点故障”。这不是过拟合，下一步会探究原因。

步骤 6：无需训练即可测量的过平滑

用 \hat{\mathbf{A}} 反复传播，会让相连节点的特征越来越相似（第 8 节）。要在没有任何网络的情况下看到这一点，就把 \hat{\mathbf{A}}^k 作用于最大那棵测试树的原始特征，并测量所有节点对之间的平均余弦相似度：1 表示每个节点都指向同一个方向。这里不涉及任何权重，所以这只是图和归一化本身的性质。

这里尝试的补救方法来自一般的深度网络（模块 03 第 8 节）：残差连接，\mathbf{H} \leftarrow \mathbf{H} + \mathrm{ReLU}(\hat{\mathbf{A}}\mathbf{H}\mathbf{W})， 使每个节点在平滑后的特征之外还保留自己的特征，梯度也有一条穿过 16 层的直接通路。

sizes_test = [len(t["type"]) for t in test_trees]
biggest = Graph([test_trees[int(np.argmax(sizes_test))]])
Z = biggest.x.clone()
print(f"largest test tree: {biggest.n} nodes")
curve_k, curve_cos = [], []
for k in range(1, 65):
    Z = propagate(biggest, Z)
    U = F.normalize(Z, dim=1)
    cos = float(((U @ U.T).sum() - biggest.n) / (biggest.n * (biggest.n - 1)))
    curve_k.append(k)
    curve_cos.append(cos)
    if k in (1, 2, 4, 8, 16, 32, 64):
        print(f"k = {k:2d}: mean pairwise cosine similarity of A_hat^k X = {cos:.3f}")

# Is the 16-layer failure underfitting? Compare train and test accuracy, with and without
# residual connections.
for name, residual in (("plain", False), ("residual", True)):
    torch.manual_seed(0)
    model16 = fit(GCN(16, residual=residual), g_train)
    tr_acc, te_acc = evaluate(model16, g_train)[0], evaluate(model16, g_test)[0]
    print(f"16 layers, {name:8s}: train accuracy {tr_acc:.3f}, test accuracy {te_acc:.3f}")

plt.figure(figsize=(6, 3.8))
plt.semilogx(curve_k, curve_cos, marker=".")
plt.xlabel("propagation steps k")
plt.ylabel("mean pairwise cosine similarity")
plt.title(f"Over-smoothing: node features of one {biggest.n}-node tree under $\\hat{{A}}^k X$")
plt.show()

largest test tree: 69 nodes
k =  1: mean pairwise cosine similarity of A_hat^k X = 0.840
k =  2: mean pairwise cosine similarity of A_hat^k X = 0.906
k =  4: mean pairwise cosine similarity of A_hat^k X = 0.932
k =  8: mean pairwise cosine similarity of A_hat^k X = 0.957
k = 16: mean pairwise cosine similarity of A_hat^k X = 0.976
k = 32: mean pairwise cosine similarity of A_hat^k X = 0.990
k = 64: mean pairwise cosine similarity of A_hat^k X = 0.997
16 layers, plain   : train accuracy 0.779, test accuracy 0.781
16 layers, residual: train accuracy 0.967, test accuracy 0.948

相似度起点就很高，因为不同节点的四维特征本来就相似，然后它向 1 攀升：16 步之后，\hat{\mathbf{A}}^k\mathbf{X} 的各行几乎平行。只有节点的度还有差别，它通过主特征向量中的因子 \sqrt{\tilde d_i} 体现，而节点的类型已经被平均掉了。过平滑是图和归一化的性质，任何权重都无法完全消除它。

朴素 16 层模型的训练准确率等于其测试准确率，两者都等于多数类比率：它没有学会训练集。这是一次优化失败。经过许多带 ReLU 和权重衰减的层平均之后，信号留下的梯度指向毫无用处的方向，训练就停留在把每个事件都判为负类的平台上。残差连接给每个节点的自身特征一条直接通路，也给梯度一条穿过整个堆叠的直接通路，于是同样的 16 层就能训练起来，达到 0.948，介于 4 层和 8 层的朴素模型之间（0.938 和 0.959）。问题不在于深度，而在于平滑时没有为节点自身的特征留出通路。

步骤 7：方向感知层

GCN 的对称邻接对一个节点的门、它的兄弟节点和它的输入一视同仁。而要学习的性质并非如此：它只取决于事件之上的门。一个把两个方向分开的层，可以直接表达“我上方的每个门都是或门”。方向感知层有三个权重矩阵：

一个用于节点自身，一个用于来自节点所在门（它的父节点；顶事件为零）的消息，一个用于来自其输入（它的子节点；基本事件为零）的消息的均值。它就是在有向边列表上的两次 index_add_ 调用。经过 L 层之后，关于上方 L 层那个门的信息只通过 \mathbf{W}_p 到达节点，丝毫没有被兄弟节点稀释。

class DirGNN(nn.Module):
    def __init__(self, L, hidden=32):
        super().__init__()
        self.inp = nn.Linear(4, hidden)
        self.Ws = nn.ModuleList(nn.Linear(hidden, hidden) for _ in range(L))
        self.Wp = nn.ModuleList(nn.Linear(hidden, hidden, bias=False) for _ in range(L))
        self.Wc = nn.ModuleList(nn.Linear(hidden, hidden, bias=False) for _ in range(L))
        self.out = nn.Linear(hidden, 2)

    def forward(self, g):
        h = self.inp(g.x)
        for Ws, Wp, Wc in zip(self.Ws, self.Wp, self.Wc):
            from_gate = torch.zeros_like(h).index_add_(0, g.child, h[g.par])
            from_inputs = torch.zeros_like(h).index_add_(0, g.par, h[g.child])
            from_inputs = from_inputs / g.n_children[:, None]
            h = torch.relu(Ws(h) + Wp(from_gate) + Wc(from_inputs))
        return self.out(h)


dir_acc = {}
print("direction-aware model")
print("  L   overall   depth1   depth2   depth3   depth4")
for L in (2, 3, 4):
    torch.manual_seed(0)
    dir_model = fit(DirGNN(L), g_train)
    dir_acc[L] = evaluate(dir_model, g_test)
    print(f"{L:3d}   " + "   ".join(f"{a:.3f} " for a in dir_acc[L]))
dir4 = dir_model

depths = [1, 2, 3, 4]
plt.figure(figsize=(7, 4))
for L in (1, 2, 4):
    plt.plot(depths, gcn_acc[L][1:], marker="o", label=f"GCN, L = {L}")
plt.plot(depths, dir_acc[4][1:], marker="s", color="k", label="direction-aware, L = 4")
plt.axhline(majority, color="grey", linestyle=":", label="majority class")
plt.xticks(depths)
plt.xlabel("depth of the basic event below the top event")
plt.ylabel("test accuracy")
plt.title("Accuracy by depth: receptive field and direction")
plt.legend(loc="lower left")
plt.show()

direction-aware model
  L   overall   depth1   depth2   depth3   depth4
  2   0.853    1.000    1.000    0.797    0.843
  3   0.941    1.000    1.000    1.000    0.882
  4   1.000    1.000    1.000    1.000    1.000

L = 2 时，方向感知模型优于同样深度的 GCN，而且恰好在其感受野所及的那些事件上准确率完美：深度 1 和 2 的事件，它们的门在上方一跳和两跳处。L = 3 时它在深度 3 以内都完美，L = 4 时它正确标注了每一个基本事件。这就是具有正确归纳偏置的架构的表现：网络的容量不必花在学习“方向很重要”上，它的极限正是感受野所预言的那些。

步骤 8：把模型用于冷却系统

训练好的 L = 4 方向感知模型从未见过这棵手工构建的树。应用它是最后一项检验：它必须找出 E1、E5 和 E6，即步骤 2 的答案。代码块还按第 8 节的布局画出这棵树，并标出预测的单点故障。

dir4.eval()
with torch.no_grad():
    pred_cool = dir4(g_cool).argmax(1).numpy()
predicted = [NAMES[v] for v in np.where((pred_cool == 1) & (cooling["type"] == BASIC))[0]]
print("single points of failure predicted for the cooling system:", predicted)

pos = {0: (340, 40), 1: (130, 130), 2: (530, 130), 3: (340, 130), 4: (50, 230), 5: (130, 230),
       6: (210, 230), 7: (450, 230), 8: (530, 230), 9: (610, 230), 10: (570, 320),
       11: (650, 320)}
fig, ax = plt.subplots(figsize=(8, 4.2))
for v, p in enumerate(cooling["parent"]):
    if p >= 0:
        ax.plot([pos[v][0], pos[p][0]], [-pos[v][1], -pos[p][1]], color="grey", zorder=1)
for v, (px, py) in pos.items():
    kind = ["OR", "AND", "event"][cooling["type"][v]]
    hit = pred_cool[v] == 1 and cooling["type"][v] == BASIC
    ax.scatter(px, -py, s=900, zorder=2, marker="s" if kind != "event" else "o",
               color="tab:red" if hit else ("white" if kind == "event" else "lightgrey"),
               edgecolor="k")
    ax.text(px, -py, f"{NAMES[v]}\n{kind}" if kind != "event" else NAMES[v],
            ha="center", va="center", fontsize=8, zorder=3)
ax.set_title("Cooling-system fault tree: predicted single points of failure in red")
ax.set_xlabel("layout position (arbitrary units)")
ax.set_ylabel("level (top at the top)")
ax.set_yticks([])
plt.show()

single points of failure predicted for the cooling system: ['E1', 'E5', 'E6']

模型指出了 E1、E5 和 E6，别无其他。一棵小树判断正确，本身并不是有力的证据。证据是测试集，这棵树只是一次合理性检查，确认生成器生成的树与手绘的树遵循同样的规则。

你应该看到什么

• 按深度统计的准确率显示了感受野：L 层模型只对位于顶事件之下至多 L 跳的事件可靠。方向感知模型把这一点变得精确：在深度不超过 L 的每个深度上准确率都是 1.0。
• 无向 GCN 一直改进到约 4 层（0.938），在 8 层时达到峰值（0.959）。随后朴素的 12 层和 16 层模型对每个事件都预测多数类。它们的节点特征几乎完全相同（未经训练的 \hat{\mathbf{A}}^k\mathbf{X} 相似度向 1 攀升），而且没有残差通路，深的堆叠也很难训练。残差连接使它恢复。
• 对称邻接给每个事件的是它的门、它的兄弟节点以及两跳之外它的门的其他输入的混合。而这一性质只取决于上方的门。分开的父节点权重和子节点权重让网络能精确计算“上方的每个门都是或门”。
• 听起来正确的规则（“它的门是或门”）比多数类还差。每个模型都必须与这两者比较。

动手试试

1. 按节点划分。不按树划分，而是对全部 300 棵树的所有节点做随机的 70/30 划分（所有节点在一张图中都可见，只隐藏标签），并比较准确率。解释为什么按节点划分会美化一个将被用于新树的模型。
2. 深度特征。把每个节点的深度作为第五个特征加入，用 L = 4 重新训练无向 GCN。它能弥补与方向感知模型之间的差距吗？这个答案说明对称层缺少了什么？
3. 预测哪些节点是基本事件。训练一个 2 层 GCN 来预测一个节点是否为基本事件，并与你在练习 9 中给出的基线比较。GCN 学到了什么？这一题最好在做完该练习之后再尝试。

目标。 你为一个精确解已知的问题构建第 9 节的物理信息神经网络，即阻尼振子

其中 \omega_0 = 2\pi rad/s，\zeta = 0.1，区间为 t \in [0, 2] s。由于精确解已知，你可以测量误差，失败也正因此才看得见。你先用显而易见的方式训练网络，看着它收敛到平凡解 u = 0：它精确满足方程，却完全不满足初始条件。然后你用两种方式修复它：改变损失权重，以及改变单位（第三种修复是把初始条件直接构建进网络，留作扩展）。最后你把 \zeta 当作未知量，从十二个带噪声的读数中把它恢复出来，并与对闭式解做经典最小二乘拟合的结果比较。数据在实验中生成，无需下载，整个实验在笔记本电脑的 CPU 上大约运行三分钟（所示运行用时 190 s）。你需要 NumPy、SciPy、PyTorch 和 matplotlib。打印出的数字与你的结果在最后几位上可能不同。

步骤 1：问题及其精确解

当 \zeta < 1 时，满足这两个初始条件的解为

你可以用这两个条件来检验它：在 t = 0 处它等于 1，它在 0 处的导数为 -\zeta\omega_0 + \zeta\omega_0 = 0。代码还确定了本实验的工作约定。线程数设为一：这么小的网络每一步的运算量太少，无法从多线程中获益，而单线程可以避免其他程序争抢核心时出现的严重减速（在一个原型中，机器空闲时无论线程数多少，每一步约需 1.2 ms；机器繁忙时用四个线程，每一步约需 30 ms）。误差指标是在 1,000 个测试时刻上的相对 L2 误差 \lVert u_\theta - u\rVert_2 / \lVert u\rVert_2，对于输出恒为零的网络，它等于 1。

import time

import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn as nn
from scipy.optimize import curve_fit

torch.manual_seed(0)
torch.set_num_threads(1)

ZETA, W0, T_END = 0.1, 2 * np.pi, 2.0


def exact(t, zeta=ZETA):
    """Closed-form solution of u'' + 2 zeta w0 u' + w0^2 u = 0, u(0) = 1, u'(0) = 0."""
    wd = W0 * np.sqrt(1.0 - zeta**2)
    return np.exp(-zeta * W0 * t) * (np.cos(wd * t) + (zeta * W0 / wd) * np.sin(wd * t))


def rel_l2(u_pred, u_true):
    return float(np.linalg.norm(u_pred - u_true) / np.linalg.norm(u_true))


t_test = np.linspace(0.0, T_END, 1000)
u_test = exact(t_test)
print(f"damped frequency {W0 * np.sqrt(1 - ZETA**2) / (2 * np.pi):.4f} Hz")
print(f"u(0) = {exact(0.0):.4f}, u(0.25 s) = {exact(0.25):.4f}, u(2 s) = {exact(2.0):.4f}")
print(f"largest |u| after 1 s: {np.abs(u_test[t_test > 1.0]).max():.4f}")

plt.figure(figsize=(7, 3.2))
plt.plot(t_test, u_test, color="black")
plt.xlabel("time t (s)")
plt.ylabel("displacement u")
plt.title("Exact solution: damped oscillator, w0 = 2 pi rad/s, zeta = 0.1")
plt.show()

damped frequency 0.9950 Hz
u(0) = 1.0000, u(0.25 s) = 0.0926, u(2 s) = 0.2822
largest |u| after 1 s: 0.5318

振子在这个区间内完成两个周期，其包络衰减到初始幅值的约 29%。记住这两个事实：网络必须重现两次振荡，而它要重现的幅值量级为 1。

步骤 2：网络、它的导数与损失

网络 u_\theta(t) 是一个多层感知机，有三个隐藏层，每层 32 个 tanh 单元。它把输入除以区间长度，使第一层看到的数总在 [0, 1] 之内，无论采用什么单位。光滑的 tanh 很重要：残差需要网络的二阶导数，而 ReLU 网络的二阶导数几乎处处为零。

求导的辅助函数就是 PINN 的全部机制。torch.autograd.grad 对输出关于输入时间求导，create_graph=True 使导数本身仍然可微，这样它既可以再求一次导（得到 u''），由它构建的损失也可以对权重求导。方程的残差对任意系数 c_1, c_2 写成 u'' + c_1 u' + c_2 u；有量纲问题中 c_1 = 2\zeta\omega_0、c_2 = \omega_0^2，步骤 6 会使用另一组系数。第 9 节的损失为

其中 N = 200 个配点（collocation points）均匀分布。fit 函数用学习率为 10^{-3} 的 Adam 训练给定的步数，并在你指定的步上记录两个损失项、测试误差，以及测试时刻上最大的 |u_\theta|。

class PINN(nn.Module):
    def __init__(self, t_scale, width=32):
        super().__init__()
        self.t_scale = t_scale  # input is divided by the interval length
        self.net = nn.Sequential(
            nn.Linear(1, width), nn.Tanh(),
            nn.Linear(width, width), nn.Tanh(),
            nn.Linear(width, width), nn.Tanh(),
            nn.Linear(width, 1),
        )

    def forward(self, t):
        return self.net(t / self.t_scale)


def d(u, t):
    """du/dt by autograd; create_graph keeps the result differentiable."""
    return torch.autograd.grad(u, t, torch.ones_like(u), create_graph=True)[0]


def loss_terms(model, t_col, t_zero, c1, c2):
    u = model(t_col)
    u_t = d(u, t_col)
    u_tt = d(u_t, t_col)
    residual = (u_tt + c1 * u_t + c2 * u).pow(2).mean()
    u0 = model(t_zero)
    initial = (u0 - 1.0).pow(2).mean() + d(u0, t_zero).pow(2).mean()
    return residual, initial


def fit(t_end, c1, c2, lam_ic, steps, log_at=(), lr=1e-3, seed=0):
    """Train a PINN; return the model and a log of (step, residual, ic, error, max|u|)."""
    torch.manual_seed(seed)
    model = PINN(t_end)
    opt = torch.optim.Adam(model.parameters(), lr=lr)
    t_col = torch.linspace(0.0, t_end, 200).reshape(-1, 1).requires_grad_(True)
    t_zero = torch.zeros(1, 1, requires_grad=True)
    t_eval = torch.tensor(t_test / T_END * t_end, dtype=torch.float32).reshape(-1, 1)
    log, curve = [], []
    t0 = time.time()
    for step in range(steps + 1):
        residual, initial = loss_terms(model, t_col, t_zero, c1, c2)
        if step % 100 == 0 or step in log_at:
            with torch.no_grad():
                u_hat = model(t_eval).numpy().ravel()
            err = rel_l2(u_hat, u_test)
            curve.append((step, err))
            if step in log_at:
                log.append((step, residual.item(), initial.item(), err, np.abs(u_hat).max()))
        if step == steps:
            break
        loss = residual + lam_ic * initial
        opt.zero_grad()
        loss.backward()
        opt.step()
    ms = 1000 * (time.time() - t0) / steps
    return model, log, curve, ms


def show(log):
    print("  step   residual   ic term   rel. L2 error   max|u|")
    for step, res, ic, err, umax in log:
        print(f"{step:6d}   {res:8.2e}   {ic:7.4f}   {err:13.4f}   {umax:6.4f}")


# one forward pass at initialisation: the sizes of the two loss terms
model0 = PINN(T_END)
t_col0 = torch.linspace(0.0, T_END, 200).reshape(-1, 1).requires_grad_(True)
t_zero0 = torch.zeros(1, 1, requires_grad=True)
res0, ic0 = loss_terms(model0, t_col0, t_zero0, 2 * ZETA * W0, W0**2)
print(f"at initialisation: residual {res0.item():.2f}, initial conditions {ic0.item():.2f}")
print(f"parameters: {sum(p.numel() for p in model0.parameters())}")

at initialisation: residual 52.46, initial conditions 1.36
parameters: 2209

还没有任何训练，残差项就已经比初始条件项大约四十倍。这并不是初始化的偶然结果。方程中含有 \omega_0^2 \approx 39.5，所以 u 上一个单位的误差会带来量级为 40 的残差、量级为 1,500 的残差平方，而 u(0) 上同样一个单位的误差只带来 1。两项所处的尺度相差三个数量级，优化器跟着大的那一项走。

步骤 3：显而易见的损失，以及平凡解

现在在问题本身的单位下训练，两项权重相等（\lambda_{\text{ic}} = 1），训练 3,000 步。使用有量纲的系数和区间 [0, 2]。

C1_DIM, C2_DIM = 2 * ZETA * W0, W0**2
model_a, log_a, curve_a, ms = fit(T_END, C1_DIM, C2_DIM, 1.0, 3000, log_at=(0, 1000, 3000))
show(log_a)
print(f"{ms:.1f} ms per step")

with torch.no_grad():
    u_a = model_a(torch.tensor(t_test, dtype=torch.float32).reshape(-1, 1)).numpy().ravel()
plt.figure(figsize=(7, 3.2))
plt.plot(t_test, u_test, color="black", label="exact")
plt.plot(t_test, u_a, color="tab:red", label="PINN, lambda_ic = 1")
plt.xlabel("time t (s)")
plt.ylabel("displacement u")
plt.title("Dimensional units, equal weights: the trivial solution")
plt.legend()
plt.show()

  step   residual   ic term   rel. L2 error   max|u|
     0   5.25e+01    1.3626          1.0914   0.1951
  1000   3.14e-03    0.9922          0.9993   0.0039
  3000   5.92e-03    0.9877          0.9986   0.0063
5.1 ms per step

损失降到了一个很小的数，答案却是错的。残差项很小，初始条件项几乎没有离开它约为 1 的初始值，网络输出的是一条大小约 0.01 的曲线，与振荡毫无关系。相对误差接近 1。这就是第 9 节所述失败模式最纯粹的形式：u \equiv 0 是该微分方程的一个解，所以把网络“关掉”就能把残差压到零；而初始条件是区分所要的解与零解的唯一依据，它们的梯度太弱，抵挡不住。打印出的损失中没有任何迹象提醒你，只有与已知答案（或与测量值）比较才能发现。

步骤 4：彻底去掉初始条件

为了确认初始条件是唯一的屏障，令 \lambda_{\text{ic}} = 0。这时已经没有任何东西能把所要的解与零区分开来。

model_b, log_b, curve_b, _ = fit(T_END, C1_DIM, C2_DIM, 0.0, 3000, log_at=(0, 3000))
show(log_b)

  step   residual   ic term   rel. L2 error   max|u|
     0   5.25e+01    1.3626          1.0914   0.1951
  3000   5.07e-06    1.0002          1.0000   0.0001

残差降到约 5\times10^{-6}，比上一次运行低三个数量级，网络的幅值量级为 10^{-4}。误差为 1.000：网络找到的恰恰是残差为零的函数中最容易找到的那一个。一个适定问题需要它的条件，而一个不让条件成为硬性要求的损失存在平凡的极小值。

步骤 5：第一种修复，给条件一个大权重

把 \lambda_{\text{ic}} 提高到 100，使网络以关键方式出错时两项的大小相当，并训练 10,000 步。训练更长，是因为这个问题比看上去更难：要用一个起初几乎平坦的函数去拟合衰减振荡的两个周期。

model_c, log_c, curve_c, _ = fit(
    T_END, C1_DIM, C2_DIM, 100.0, 10000, log_at=(0, 1000, 5000, 10000)
)
show(log_c)

  step   residual   ic term   rel. L2 error   max|u|
     0   5.25e+01    1.3626          1.0914   0.1951
  1000   5.24e+00    0.0086          0.3306   0.9071
  5000   1.07e-01    0.0000          0.0152   0.9973
 10000   1.89e-02    0.0000          0.0039   0.9992

权重起了作用，但很慢：1,000 步后误差仍为 33%，5,000 步后约 1.5%，10,000 步后 0.4%。残差项在数千步内一直远高于零，说明优化器在两项之间被拉扯得多么厉害。代价是多了一个超参数。这里选 100，是因为我们在权重为 1 时看到了失败；而当解未知时，没有误差可看，权重只能通过试错或自适应方案来确定（见第 9 节）。

步骤 6：第二种修复，无量纲化

更深层的原因在于单位。定义无量纲时间 \hat t = \omega_0 t，对我们的区间，它的取值范围是 [0, 4\pi]。由链式法则得 \mathrm{d}/\mathrm{d}t = \omega_0\, \mathrm{d}/\mathrm{d}\hat t，于是

初始条件不变，仍为 u(0) = 1 和 u_{\hat t}(0) = 0。现在系数为 c_1 = 2\zeta = 0.2 和 c_2 = 1，所有项的量级都是 1，残差项与初始条件项无需任何手工调节的权重就大小相当。网络相同，权重回到 1，变的只是坐标。

C1_ND, C2_ND = 2 * ZETA, 1.0
model_d, log_d, curve_d, ms = fit(
    4 * np.pi, C1_ND, C2_ND, 1.0, 10000, log_at=(0, 2500, 5000, 7500, 10000)
)
show(log_d)
print(f"{ms:.1f} ms per step")

  step   residual   ic term   rel. L2 error   max|u|
     0   3.37e-02    1.3623          1.0914   0.1951
  2500   4.67e-03    0.0000          0.3892   0.9946
  5000   2.84e-04    0.0000          0.0603   1.0004
  7500   4.35e-07    0.0000          0.0004   1.0002
 10000   9.52e-06    0.0000          0.0073   1.0039
4.8 ms per step

现在初始残差是 0.034 而不是 52，所以步骤 2 中的失衡消失了，甚至反了过来：初始条件项（1.36）现在是较大的一项，优化器先满足它，再去拟合振荡。训练在开始阶段并没有更快（2,500 步后误差为 0.39，5,000 步后为 0.06，而加权运行在 5,000 步时为 0.015），但它持续改进，在 7,500 步时达到 0.0004，比加权运行的最终值 0.0039 低十倍，而且不需要调节任何权重。它并没有停留在那里。在第 10,000 步，打印出的误差为 0.0073，因为固定学习率 10^{-3} 的 Adam 不断把网络从极小值附近踢开：在本实验一个打印了 curve_d 的副本中，每 100 步记录一次的误差在第 6,000 步之后在约 0.0001 和 0.016 之间跳动，最低值（0.00006）出现在第 8,700 步。加权运行同样会振荡，在最后 1,000 步内介于 0.004 和 0.013 之间。常用的补救办法是使用衰减的学习率，或按某个验证指标保留最好的检查点；这里要说明的是，无量纲形式不用权重就能达到很低的误差。（两次运行的残差不能直接比较，因为两个方程相差一个因子 \omega_0^2。）这个教训具有普遍性，不局限于这个方程：在动用损失平衡方案之前，先把方程写成各项量级都为 1 的形式。

三次运行中记录的收敛曲线直接展示了这一点。

fig, ax = plt.subplots(1, 2, figsize=(10, 3.6))
for curve, label, colour in [
    (curve_a, "dimensional, weight 1", "tab:red"),
    (curve_c, "dimensional, weight 100", "tab:orange"),
    (curve_d, "non-dimensional, weight 1", "tab:blue"),
]:
    steps_, errs = zip(*curve)
    ax[0].semilogy(steps_, errs, label=label, color=colour)
ax[0].set_xlabel("training step")
ax[0].set_ylabel("relative L2 error")
ax[0].set_title("Error against the exact solution")
ax[0].legend()

with torch.no_grad():
    u_d = model_d(torch.tensor(W0 * t_test, dtype=torch.float32).reshape(-1, 1))
ax[1].plot(t_test, u_test, color="black", label="exact")
ax[1].plot(t_test, u_d.numpy().ravel(), "--", color="tab:blue", label="PINN, non-dimensional")
ax[1].set_xlabel("time t (s)")
ax[1].set_ylabel("displacement u")
ax[1].set_title("The repaired network")
ax[1].legend()
plt.tight_layout()
plt.show()

步骤 7：反问题

现在阻尼比未知。在 [0, 2] s 内的随机时刻，由 \zeta = 0.1 的精确解加上标准差为 0.02 的高斯噪声，生成十二个位移读数，就像由加速度计数据得到的十二个样本那样。网络与步骤 6 一样是无量纲的。未知量作为一个可训练的标量进入残差，参数化为 \log\zeta 以保证它为正，初始值为 \zeta = 0.5，比真值大五倍。损失中加入一个数据项，即读数处均方误差的十倍；权重 10 表示：对读数的信任程度，略高于对单个配点上满足物理方程的要求。网络权重和 \log\zeta 由同一个 Adam 一起优化。

g = torch.Generator().manual_seed(1)
t_meas = T_END * torch.rand(12, generator=g)
y_meas = torch.tensor(exact(t_meas.numpy()), dtype=torch.float32) + 0.02 * torch.randn(
    12, generator=g
)
th_meas = (W0 * t_meas).reshape(-1, 1)  # readings in the non-dimensional time


def fit_inverse(steps=10000, lr=1e-3, seed=0, w_data=10.0):
    torch.manual_seed(seed)
    model = PINN(4 * np.pi)
    log_zeta = nn.Parameter(torch.log(torch.tensor(0.5)))
    opt = torch.optim.Adam(list(model.parameters()) + [log_zeta], lr=lr)
    t_col = torch.linspace(0.0, 4 * np.pi, 200).reshape(-1, 1).requires_grad_(True)
    t_zero = torch.zeros(1, 1, requires_grad=True)
    t_eval = torch.tensor(W0 * t_test, dtype=torch.float32).reshape(-1, 1)
    path = []
    for step in range(steps + 1):
        zeta = log_zeta.exp()
        residual, initial = loss_terms(model, t_col, t_zero, 2 * zeta, 1.0)
        data = (model(th_meas).squeeze(1) - y_meas).pow(2).mean()
        if step % 1000 == 0:
            with torch.no_grad():
                err = rel_l2(model(t_eval).numpy().ravel(), u_test)
            path.append((step, zeta.item(), err))
        if step == steps:
            break
        loss = residual + initial + w_data * data
        opt.zero_grad()
        loss.backward()
        opt.step()
    return model, zeta.item(), path


model_e, zeta_hat, path = fit_inverse()
print("  step   zeta    solution error")
for step, z, err in path:
    print(f"{step:6d}   {z:.4f}   {err:.4f}")
print(f"recovered zeta = {zeta_hat:.4f} (true {ZETA})")

  step   zeta    solution error
     0   0.5000   1.0914
  1000   0.4235   0.2269
  2000   0.2372   0.1180
  3000   0.1736   0.0806
  4000   0.1400   0.0599
  5000   0.1201   0.0473
  6000   0.1087   0.0345
  7000   0.1030   0.0305
  8000   0.1007   0.0169
  9000   0.0996   0.0181
 10000   0.0987   0.0192
recovered zeta = 0.0987 (true 0.1)

估计值从高出五倍的位置出发，到第 5,000 步降到 0.12，最终稳定在 0.0987，比真值 0.1 低 1.3%。解的误差最终约为 0.019。在均方根值为 0.44 的信号上，0.02 的噪声相当于十二个读数处约 0.046 的相对误差，所以网络给出的解比读数本身更接近真值：方程滤掉了噪声。十二个读数加一个方程就够了，因为方程提供了曲线的形状，读数只需要确定一个数。解的误差没有步骤 6 的最好结果（0.0004）那么小，是因为数据项把网络拉向带噪声的点。

步骤 8：经典基线

对于一个有闭式解的单参数反问题，公允的比较对象是用同样的读数对闭式解做最小二乘拟合。curve_fit 最小化 exact(t, zeta) 与读数之差的平方，并根据拟合的曲率给出不确定度。

popt, pcov = curve_fit(
    lambda t, z: exact(t, z), t_meas.numpy(), y_meas.numpy(), p0=[0.5], bounds=(0.01, 0.99)
)
print(f"curve_fit: zeta = {popt[0]:.4f} +/- {np.sqrt(pcov[0, 0]):.4f}")
print(f"PINN:      zeta = {zeta_hat:.4f}")

fig, ax = plt.subplots(1, 2, figsize=(10, 3.6))
ax[0].plot(t_test, u_test, color="black", label="exact (zeta = 0.1)")
with torch.no_grad():
    u_e = model_e(torch.tensor(W0 * t_test, dtype=torch.float32).reshape(-1, 1)).numpy()
ax[0].plot(t_test, u_e.ravel(), "--", color="tab:blue", label="PINN, inverse problem")
ax[0].scatter(t_meas.numpy(), y_meas.numpy(), color="tab:red", zorder=3, label="12 readings")
ax[0].set_xlabel("time t (s)")
ax[0].set_ylabel("displacement u")
ax[0].set_title("Twelve noisy readings and the recovered solution")
ax[0].legend()
steps_p, zetas, _ = zip(*path)
ax[1].plot(steps_p, zetas, "o-", color="tab:blue", label="PINN estimate")
ax[1].axhline(ZETA, color="black", label="true value")
ax[1].axhline(popt[0], color="tab:green", linestyle="--", label="curve_fit")
ax[1].set_xlabel("training step")
ax[1].set_ylabel("damping ratio zeta")
ax[1].set_title("The estimate of zeta during training")
ax[1].legend()
plt.tight_layout()
plt.show()

curve_fit: zeta = 0.0977 +/- 0.0015
PINN:      zeta = 0.0987

两个估计值 0.0977 和 0.0987 相差约为拟合自身标准误差 0.0015 的三分之二，而真值 0.1 与最小二乘估计相差约 1.5 个标准误差：两者都与真值一致，彼此也一致。这是正确的结论，而不是令人失望的结论。当解是只含一个未知量的闭式解时，最小二乘拟合更快，能给出误差棒，也不会陷入平凡解。PINN 的价值在没有闭式解时才体现出来：非线性方程、不规则几何、作为空间场的未知量。本实验的价值在于，你在一个小到可以核对的问题上看到了每一步。

你应该看到什么

• 在有量纲单位下，残差项一开始约比初始条件项大 40 倍（52.5 对 1.4）。权重相等时，优化器把网络推向几乎为零：损失很小（残差 0.006），相对误差却是 0.999。没有初始条件时结果相同，残差为 5\times10^{-6}。平凡解精确满足方程。
• 给初始条件加权 100 可以修复它（10,000 步后误差 0.0039）；无量纲化不用权重也能修复它，因为方程各项的量级都是 1，而且能达到更低的误差（7,500 步时 0.0004），只是在固定学习率下误差会来回跳动，在 10,000 步时读数为 0.0073。
• 反问题从十二个带噪声的读数中恢复出 \zeta = 0.0987，与真值相差 1.3% 以内，并与对闭式解的最小二乘拟合（0.0977 \pm 0.0015）一致。有闭式解时就用闭式解；没有闭式解时，PINN 才物有所值。
• 通过自动微分求二阶导数需要多次前向和反向传播，但这个网络每一步只需约 5 ms，所以整个实验在 CPU 上大约三分钟就能跑完。开销在于步数（达到 0.0004 需要 7,500 步），而不在于每一步。

动手试试

1. 硬约束。 把初始条件构建进网络，u_\theta(t) = 1 + (t/t_{\text{end}})^2\,N_\theta(t)：无论 N_\theta 是什么，它在 t = 0 处都等于 1，且导数为零。去掉初始条件项，在有量纲单位下训练 10,000 步。平凡解不再可能出现，所以步骤 3 的失败不会发生。误差是否仍然很大？为什么？（尺度问题并没有消失。）
2. 谱偏差。 令 \omega_0 = 8\pi（2 s 内有八个周期），再次训练无量纲形式；网络现在需要八次振荡。然后加入傅里叶特征 [\sin k\hat t, \cos k\hat t]（k = 1, 2, 4）作为额外输入，比较两次运行。为什么输入量级为 1 的 tanh 网络难以处理高频？
3. 从常微分方程到偏微分方程。 练习 e12 从这段代码出发，转向热传导方程，那里的配点变成空间和时间上的网格。

目标。 你在机器振动窗口上用第 11 节的 InfoNCE 损失预训练一个编码器，训练中从不向它展示这些窗口的标签，然后用线性探测（linear probe）衡量预训练带来了什么：一个只用每类 5、20 或 100 个有标签窗口训练的逻辑回归。比较对象是工程师首先会尝试的两样东西，即原始波形及其频谱幅值，以及一个从未训练过的编码器。随后你逐一去掉数据增强，看着表示变差，这正是本实验的要点：在对比学习中，数据增强就是监督。振动数据在实验中生成（四个类别，随机相位），所以无需下载，整个实验在笔记本电脑的 CPU 上运行一到两分钟。你需要 NumPy、scikit-learn、PyTorch 和 matplotlib。打印出的数字与你的结果在最后几位上可能不同。

步骤 1：振动窗口生成器

真实的旋转机械会用加速度计记录。这里一个窗口是以 256 Hz 采样的一秒钟，即 256 个样本，对应一根以转速 f_r 旋转的轴，f_r 在 9 到 11 Hz 之间均匀抽取。四个类别的区别在于它们包含 f_r 的哪些谐波：

所有相位都是随机的，传感器增益在 0.8 到 1.2 上均匀分布，并加上标准差为 0.3 的高斯噪声。随机相位是这个问题的核心。当窗口从周期中任意时刻开始时，真实的记录就是这个样子；这也意味着同一类别的两个窗口没有任何共同的样本值。标签取决于各谐波的幅值，而不取决于相位。

冲击序列由距上一次冲击的时间 \tau 构建，它等于 (t - t_0) 对冲击周期 1/(3.57 f_r) 取模；这样只用一行向量化代码，就能得到每次冲击之后完整的衰减振荡。

import time

import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import confusion_matrix
from sklearn.preprocessing import StandardScaler

np.random.seed(0)
torch.manual_seed(0)
rng = np.random.default_rng(0)

FS, N_SAMPLES = 256, 256
T = np.arange(N_SAMPLES) / FS
CLASS_NAMES = ["healthy", "imbalance", "misalignment", "bearing defect"]


def make_windows(n, rng):
    labels = rng.integers(0, 4, size=n)
    f_r = rng.uniform(9.0, 11.0, size=(n, 1))
    ph = rng.uniform(0.0, 2 * np.pi, size=(n, 3))

    def wave(k, i):  # k-th harmonic of the shaft speed with its own random phase
        return np.sin(2 * np.pi * k * f_r * T + ph[:, i:i + 1])

    amp = np.array([[1.0, 0.2, 0.0],      # healthy: amplitudes of 1x, 2x, 3x
                    [2.5, 0.2, 0.0],      # imbalance
                    [1.0, 1.5, 0.5],      # misalignment
                    [1.0, 0.2, 0.0]])[labels]  # bearing defect: healthy + impulses
    x = amp[:, 0:1] * wave(1, 0) + amp[:, 1:2] * wave(2, 1) + amp[:, 2:3] * wave(3, 2)
    period = 1.0 / (3.57 * f_r)
    t0 = rng.uniform(0.0, 1.0, size=(n, 1)) * period
    tau = np.mod(T - t0, period)  # time since the most recent impulse
    impulses = 1.5 * np.exp(-30.0 * tau) * np.sin(2 * np.pi * 60.0 * tau)
    x = x + (labels == 3)[:, None] * impulses
    x = x * rng.uniform(0.8, 1.2, size=(n, 1)) + 0.3 * rng.standard_normal((n, N_SAMPLES))
    return x.astype(np.float32), labels


X_pool, y_pool = make_windows(4000, rng)  # pretraining pool; labels kept aside
X_test, y_test = make_windows(2000, rng)
print(X_pool.shape, X_test.shape)
print("class counts in the pool:", np.bincount(y_pool))
print(f"signal standard deviation {X_pool.std():.3f}")

fig, axes = plt.subplots(2, 4, figsize=(13, 4.8))
spec_axis = np.fft.rfftfreq(N_SAMPLES, 1 / FS)
for c in range(4):
    x = X_pool[np.flatnonzero(y_pool == c)[0]]
    axes[0, c].plot(T, x, color="black", linewidth=0.8)
    axes[0, c].set_title(CLASS_NAMES[c])
    axes[0, c].set_xlabel("time (s)")
    axes[1, c].plot(spec_axis, np.abs(np.fft.rfft(x)) / N_SAMPLES * 2, color="tab:blue")
    axes[1, c].set_xlabel("frequency (Hz)")
    axes[1, c].set_xlim(0, 128)
axes[0, 0].set_ylabel("acceleration (a.u.)")
axes[1, 0].set_ylabel("spectral magnitude")
fig.suptitle("One window of each class (top) and its magnitude spectrum (bottom)")
plt.tight_layout()
plt.show()

(4000, 256) (2000, 256)
class counts in the pool: [ 986  985  991 1038]
signal standard deviation 1.306

先看频谱再看波形，并注意各纵轴的刻度不同。健康窗口在约 f_r 处有一个峰，在 2f_r 处有一个小峰。不平衡窗口在 f_r 处的峰更高（1x 幅值大 2.5 倍，不过由于窗口里不是整数个周期，频谱泄漏也会降低峰高）。不对中窗口的最大峰在 2f_r 处，在 3f_r 处还有第三个峰。轴承缺陷窗口多出一排位于冲击频率 3.57 f_r 整数倍处的梳状峰，在这个窗口中靠近 39 和 78 Hz。在时域里，这些类别很难用肉眼区分，而且相位在窗口之间各不相同。频谱是工程师习惯计算的东西，步骤 3 把它们用作基线。

步骤 2：探测协议

表示的质量，用一个线性分类器在标签很少时在它上面的表现来衡量。线性探测无法通过学习表示的某个巧妙非线性函数来弥补一个糟糕的表示，所以它的准确率反映的是特征本身。对每类 5、20 和 100 个窗口这三种标签预算，辅助函数从池中为每类抽取这么多有标签窗口，用不同的抽取重复五次，在每次抽取上拟合一个经过标准化的逻辑回归，并对 2,000 个测试窗口上的准确率求平均。在每类 5 个标签时，五次抽取很重要，因为单独一次抽取的 20 个窗口可能运气不好。

BUDGETS = (5, 20, 100)


def probe(feat_pool, feat_test, budgets=BUDGETS, n_draws=5, seed=0, return_model=False):
    """Mean test accuracy of a logistic regression on a few labelled windows per class."""
    draw_rng = np.random.default_rng(seed)
    accs, last = [], None
    for n_per_class in budgets:
        scores = []
        for _ in range(n_draws):
            idx = np.concatenate([
                draw_rng.choice(np.flatnonzero(y_pool == c), n_per_class, replace=False)
                for c in range(4)
            ])
            scaler = StandardScaler().fit(feat_pool[idx])
            clf = LogisticRegression(max_iter=5000).fit(scaler.transform(feat_pool[idx]),
                                                         y_pool[idx])
            pred = clf.predict(scaler.transform(feat_test))
            scores.append((pred == y_test).mean())
            last = pred
        accs.append(float(np.mean(scores)))
    return (accs, last) if return_model else accs


def fmt(accs):
    return "  ".join(f"{a:.3f}" for a in accs)

步骤 3：基线

三条基线给出了门槛。第一条是在 256 个原始样本上的逻辑回归。第二条是在 129 个 FFT 幅值上的逻辑回归：这是旋转机械的经典特征，按构造就丢弃了相位。第三条是步骤 5 的编码器在任何训练之前、权重随机时的结果：随机网络是一个信息量出奇丰富的特征提取器，任何声称预训练有帮助的说法都必须胜过它。

acc_raw = probe(X_pool, X_test)
fft_pool = np.abs(np.fft.rfft(X_pool, axis=1)).astype(np.float32)
fft_test = np.abs(np.fft.rfft(X_test, axis=1)).astype(np.float32)
acc_fft = probe(fft_pool, fft_test)
print(f"raw waveform (256 features):   {fmt(acc_raw)}")
print(f"FFT magnitude (129 features):  {fmt(acc_fft)}")


class Encoder(nn.Module):
    """h = f(x) is the representation kept after pretraining; z = g(h) feeds the loss."""

    def __init__(self, d_in=256, d_hidden=256, d_h=128, d_z=64):
        super().__init__()
        self.f = nn.Sequential(nn.Linear(d_in, d_hidden), nn.ReLU(), nn.Linear(d_hidden, d_h))
        self.g = nn.Sequential(nn.ReLU(), nn.Linear(d_h, d_z))  # projection head

    def forward(self, x):
        h = self.f(x)
        return h, F.normalize(self.g(h), dim=1)


def embed(model, X):
    with torch.no_grad():
        h, z = model(torch.from_numpy(X))
    return h.numpy(), z.numpy()


torch.manual_seed(0)
untrained = Encoder()
h_pool0, _ = embed(untrained, X_pool)
h_test0, _ = embed(untrained, X_test)
acc_untrained = probe(h_pool0, h_test0)
print(f"untrained encoder, h (128):    {fmt(acc_untrained)}")

raw waveform (256 features):   0.371  0.448  0.473
FFT magnitude (129 features):  0.850  0.899  0.974
untrained encoder, h (128):    0.532  0.738  0.895

可以看出三个事实。原始波形远远落后于其他两者：在相位随机时，256 个样本的任何固定线性组合都无法识别类别，这正是上面陈述的问题。FFT 幅值要好得多，因为幅值不依赖于相位，但在每类 5 个标签时，它们只达到约 0.85，与步骤 6 中最好的表示之间的差距也在这里最大。未训练的网络已经优于原始波形，因为信号的随机非线性特征在一定程度上对相位不敏感；但在每类 5 个标签时它仍远低于 FFT（0.53 对 0.85），而它就是训练后的编码器必须越过的门槛。

步骤 4：损失

损失是 InfoNCE 的 NT-Xent 形式（第 11 节）。一个含 B 个窗口的 batch 给出 2B 个视图，即每个窗口的两个增强版本。设 \mathbf{z}_1, \dots, \mathbf{z}_{2B} 是它们经过 L2 归一化的嵌入，于是点积就是余弦相似度。对于视图 i，设其孪生视图为 j(i)，损失为

这是在其余 2B - 1 个视图上的交叉熵，孪生视图是正确的类别。在代码中，这是一次矩阵乘法、把对角线设为 -\infty（一个视图永远不是它自己的候选），再以孪生视图的索引为目标调用 F.cross_entropy。温度 \tau 去除相似度：\tau 越小，softmax 越尖锐，对最难的负样本惩罚越重。

合理性检查复用第 11 节的算例：一个锚点与正样本和三个负样本的相似度为 (0.9, 0.2, 0.1, -0.3)。在 \tau = 1 时，softmax 几乎是平的，损失为 0.81；在 \tau = 0.1 时，正样本占主导，损失为 0.0013。

def nt_xent(z1, z2, tau):
    """InfoNCE over 2B views; z1[i] and z2[i] are two views of window i (unit vectors)."""
    b = z1.shape[0]
    z = torch.cat([z1, z2], dim=0)                     # 2B x d
    logits = z @ z.T / tau                             # 2B x 2B cosine similarities / tau
    logits.fill_diagonal_(float("-inf"))               # a view is not its own negative
    target = torch.cat([torch.arange(b, 2 * b), torch.arange(0, b)])  # index of the twin
    return F.cross_entropy(logits, target)


sims = torch.tensor([[0.9, 0.2, 0.1, -0.3]])
for tau in (1.0, 0.1):
    print(f"tau = {tau}: loss {F.cross_entropy(sims / tau, torch.tensor([0])).item():.4f}")

# Collapse check: identical embeddings for every view give log(2B - 1).
B = 256
collapsed = F.normalize(torch.ones(B, 64), dim=1)
print(f"collapsed embeddings: {nt_xent(collapsed, collapsed, 0.2).item():.4f}")
print(f"log(2B - 1) = log({2 * B - 1}) = {np.log(2 * B - 1):.4f}")

tau = 1.0: loss 0.8096
tau = 0.1: loss 0.0013
collapsed embeddings: 6.2364
log(2B - 1) = log(511) = 6.2364

坍塌值是一个不携带任何信息的表示的损失：每个候选看起来都一样，所以 softmax 在 2B - 1 = 511 个视图上是均匀的，损失为 \log 511。训练必须把损失降到远低于这个数；如果损失停在它附近，说明编码器已经坍塌，或者数据增强破坏性太强。

步骤 5：数据增强与预训练循环

每个视图由三个随机操作生成，选择它们是为了说明哪些东西对表示来说不应有影响：随机若干个样本的循环时移（窗口起点是任意的，所以相位不应有影响）、在 0.8 到 1.25 之间按对数均匀分布抽取的增益（传感器灵敏度不应有影响），以及标准差为 0.1 的附加噪声。时移用 torch.gather 实现，使 batch 中每个窗口都有自己的时移量。循环时移几乎完全保留频谱的幅值，而改变相位，这正是我们想教给模型的不变性（环绕把窗口的末尾接到了开头，这是一个小的伪影，真实的流水线会从更长的记录中截取窗口来避免它）。

编码器就是步骤 3 中构建的那个。训练使用 B = 256 的 batch（于是每个视图有一个孪生视图，以及来自其他窗口的 2B - 2 = 510 个视图作为负样本），学习率为 10^{-3} 的 Adam，\tau = 0.2，在 4,000 个池窗口上训练 60 个轮次：每个轮次 15 个 batch，共 900 步。不使用标签。

def augment(x, gen, shift=True, gain=True, noise=True):
    b, n = x.shape
    if shift:
        s = torch.randint(0, n, (b, 1), generator=gen)
        idx = (torch.arange(n).unsqueeze(0) + s) % n
        x = torch.gather(x, 1, idx)
    if gain:
        x = x * torch.exp(torch.empty(b, 1).uniform_(np.log(0.8), np.log(1.25), generator=gen))
    if noise:
        x = x + 0.1 * torch.randn(x.shape, generator=gen)
    return x


def pretrain(seed=0, epochs=60, batch=256, tau=0.2, lr=1e-3, **aug):
    torch.manual_seed(seed)
    gen = torch.Generator().manual_seed(seed)
    model = Encoder()
    opt = torch.optim.Adam(model.parameters(), lr=lr)
    data = torch.from_numpy(X_pool)
    losses = []
    for epoch in range(epochs):
        order = torch.randperm(len(data), generator=gen)
        for start in range(0, len(data) - batch + 1, batch):  # drop the last short batch
            xb = data[order[start:start + batch]]
            _, z1 = model(augment(xb, gen, **aug))
            _, z2 = model(augment(xb, gen, **aug))
            loss = nt_xent(z1, z2, tau)
            opt.zero_grad()
            loss.backward()
            opt.step()
            losses.append(loss.item())
    return model, losses


t0 = time.time()
model_full, losses_full = pretrain(seed=0)
print(f"{len(losses_full)} steps in {time.time() - t0:.1f} s")
print(f"loss at step 0: {losses_full[0]:.3f}, final loss (mean of last 15): "
      f"{np.mean(losses_full[-15:]):.3f}")
print(f"loss of collapsed embeddings: log(2B - 1) = {np.log(511):.3f}")

plt.figure(figsize=(7, 3.2))
plt.plot(losses_full, color="tab:blue")
plt.axhline(np.log(511), color="grey", linestyle="--", label="log(2B - 1), collapsed")
plt.xlabel("training step")
plt.ylabel("InfoNCE loss")
plt.title("Contrastive pretraining loss (tau = 0.2, B = 256)")
plt.legend()
plt.show()

900 steps in 12.4 s
loss at step 0: 6.150, final loss (mean of last 15): 2.376
loss of collapsed embeddings: log(2B - 1) = 6.236

损失从 6.15 开始，接近坍塌值 6.24，然后降到约 2.4，即坍塌值的 38%。它不会降到零，也不应该降到零：在噪声和增益都不同的情况下，一个窗口的两个视图永远不可能在其他 510 个视图中被完美匹配；如果损失趋于零，反倒说明这个任务是平凡的。

步骤 6：在预训练表示上做线性探测

编码器现在用作冻结的特征提取器。探测在与基线相同的有标签抽取上、按相同的协议训练。探测两种表示：\mathbf{h}，即编码器本体的 128 维输出；以及 \mathbf{z}，即经过损失所看到的投影头之后的 64 维输出。另外用两个随机种子再做两次预训练，可以感受一下波动的大小。

h_pool, z_pool = embed(model_full, X_pool)
h_test, z_test = embed(model_full, X_test)
acc_h = probe(h_pool, h_test)
acc_z = probe(z_pool, z_test)
print("labels per class:           ", "   ".join(f"{b:>3d}  " for b in BUDGETS))
print(f"raw waveform:               {fmt(acc_raw)}")
print(f"FFT magnitude:              {fmt(acc_fft)}")
print(f"untrained encoder, h:       {fmt(acc_untrained)}")
print(f"contrastive, h:             {fmt(acc_h)}")
print(f"contrastive, z (after head):{fmt(acc_z)}")

for seed in (1, 2):
    m, _ = pretrain(seed=seed)
    hp, _ = embed(m, X_pool)
    ht, _ = embed(m, X_test)
    print(f"seed {seed}, h, 5 labels per class: {probe(hp, ht, budgets=(5,))[0]:.3f}")

fig, ax = plt.subplots(figsize=(7.5, 3.8))
xs = np.arange(len(BUDGETS))
for k, (label, accs) in enumerate([("raw waveform", acc_raw), ("FFT magnitude", acc_fft),
                                   ("untrained encoder h", acc_untrained),
                                   ("contrastive h", acc_h), ("contrastive z", acc_z)]):
    ax.bar(xs + 0.16 * (k - 2), accs, width=0.16, label=label)
ax.set_xticks(xs, [str(b) for b in BUDGETS])
ax.set_xlabel("labelled windows per class")
ax.set_ylabel("test accuracy of a linear probe")
ax.set_title("What pretraining buys when labels are scarce")
ax.legend(fontsize=8)
plt.show()

labels per class:              5      20     100
raw waveform:               0.371  0.448  0.473
FFT magnitude:              0.850  0.899  0.974
untrained encoder, h:       0.532  0.738  0.895
contrastive, h:             0.956  0.988  0.996
contrastive, z (after head):0.620  0.920  0.996
seed 1, h, 5 labels per class: 0.954
seed 2, h, 5 labels per class: 0.944

按列读这张表。每类 5 个标签时，预训练得到的 \mathbf{h} 得分 0.956，FFT 幅值为 0.850，未训练的编码器为 0.532，原始波形为 0.371：一个在预训练中从未见过标签的网络，比经典特征高出十一个点。另外两个随机种子给出 0.954 和 0.944，所以这个优势并不依赖于某一次幸运的初始化。每类 100 个标签时，优势几乎消失（0.996 对 0.974）：一旦标签充足，一个好的手工特征就够了，预训练带来的是标签效率。投影输出 \mathbf{z} 则是另一回事：在 5 个标签时它只有 0.620，远低于 \mathbf{h}，到 100 个标签时才追上来。投影头被训练成对数据增强所改变的一切都保持不变，它丢掉的信息多于类别所需；这就是 SimCLR 保留 \mathbf{h}、丢弃投影头的原因。这个协议的误差棒并不小（五次有标签窗口的抽取，三个预训练随机种子），所以行与行之间一两个点的差异不应被读作排名。

步骤 7：数据增强就是监督

最后一个实验去掉数据增强。先只去掉时移（保留增益和噪声），再把三者全部去掉。没有时移时，一个窗口的两个视图具有相同的相位，所以匹配它们最容易的方法是记住波形中依赖相位的细节，而这恰恰是对分类毫无用处的特征。损失仍然下降，因为任务仍然被解决了，但它是用错误的特征解决的。无时移模型在每类 100 个标签时的混淆矩阵显示了哪些类别为此付出代价。

model_noshift, loss_noshift = pretrain(seed=0, shift=False)
model_noaug, loss_noaug = pretrain(seed=0, shift=False, gain=False, noise=False)
ablation = {}
for name, m in [("no time shift", model_noshift), ("no augmentation", model_noaug)]:
    hp, _ = embed(m, X_pool)
    ht, _ = embed(m, X_test)
    ablation[name] = probe(hp, ht, return_model=True)
    print(f"{name:<22s}{fmt(ablation[name][0])}")
print(f"{'with all three':<22s}{fmt(acc_h)}")
print(f"final losses: full {np.mean(losses_full[-15:]):.3f}, "
      f"no shift {np.mean(loss_noshift[-15:]):.3f}, none {np.mean(loss_noaug[-15:]):.3f}")

pred = ablation["no time shift"][1]
cm = confusion_matrix(y_test, pred)
print("confusion matrix of the no-shift model (100 labels per class; rows = true):")
print(cm)
print(f"misalignment predicted as imbalance: {cm[2, 1]} of {cm[2].sum()}")
print(f"bearing defect predicted as healthy: {cm[3, 0]} of {cm[3].sum()}")

no time shift         0.622  0.809  0.977
no augmentation       0.581  0.747  0.942
with all three        0.956  0.988  0.996
final losses: full 2.376, no shift 1.746, none 1.692
confusion matrix of the no-shift model (100 labels per class; rows = true):
[[503   0   0   3]
 [  0 531   0   0]
 [  0  18 454   1]
 [ 13   0   0 477]]
misalignment predicted as imbalance: 18 of 473
bearing defect predicted as healthy: 13 of 490

两点观察。没有时移时，每类 5 个标签下的探测准确率从 0.956 降到 0.622，只比未训练编码器的 0.532 高约 0.09，远低于使用时移的模型；完全没有数据增强时为 0.581。在 100 个标签时，差距缩小到 2 个点以内（0.977 对 0.996），所以损害同样体现在标签效率上：特征可用，但没有按类别组织起来。最终的损失才是更有启发性的数字：没有时移时为 1.75，没有任何数据增强时为 1.69，两者都低于完整方案的 2.38。当两个视图共享相位时，匹配它们的任务更容易，所以损失在改善，表示却在变差。对比损失衡量的是编码器解决前置任务（pretext task）的程度，而不是特征对下游任务有多好，这就是为什么由探测而不是损失来做判断。无时移模型的混淆矩阵（一次探测，每类 100 个标签的一次抽取）把错误放在了物理规律所预言的地方：不对中窗口被当作不平衡（473 个中有 18 个），两者都有很强的低次谐波；轴承缺陷被当作健康（490 个中有 13 个），它们的冲击比谐波小。这些计数取决于抽取；要看的是规律，而不是具体数字。

你应该看到什么

• 每类 5 个标签时，原始波形上的线性分类器在四个类别上接近随机水平（0.37，随机水平为 0.25），在 100 个标签时也只达到 0.47：在相位随机时，样本的任何固定线性组合都无法识别类别。频谱幅值这一经典特征按构造去掉了相位，表现很好（从 0.85 升到 0.97）。
• 带时移增强的对比预训练无需标签就能学到相位不变的表示。每类 5 个标签时它达到 0.956，比 FFT 特征高出约十一个点，另外两个预训练随机种子给出 0.95 和 0.94；在 100 个标签时，两者相差不到 2.5 个点。
• 数据增强就是监督。没有时移时，5 个标签下的准确率降到 0.62，只比未训练的编码器（0.53）略高，尽管预训练损失更低（1.75 对 2.38）。对比损失更低并不意味着特征更好。
• 投影头吸收了数据增强所改变的东西：每类 5 个标签时探测 \mathbf{z} 得到 0.62，而 \mathbf{h} 为 0.956，这就是保留 \mathbf{h} 的原因。
• InfoNCE 损失最终约为 2.4，远高于零，也远低于 \log(2B - 1) = 6.24，即不携带任何信息的嵌入所对应的值。

动手试试

1. 更宽的增益增强。 把增益范围从 0.8 到 1.25 换成 0.25 到 4，再次探测 \mathbf{h}。探测结果几乎不变（本实验一个使用更宽范围的副本在 5、20 和 100 个标签时给出 0.956、0.990 和 0.997，原来为 0.956、0.988 和 0.996），因为健康窗口和不平衡窗口在谐波比例和信噪比上也有差异，幅值并不是唯一的线索。设计一种在这些数据中确实会抹掉某个类别区分的数据增强，并用混淆矩阵确认。这就是第 11 节中的概念性失败，即旋转之下的 6 和 9，在这里变得具体。
2. 温度与 batch 大小。 让 \tau 取 0.05、0.5 和 1.0，batch 大小取 64 和 512，保持轮次数不变。哪些设置会改变探测在 5 个标签时的准确率？最终损失能预测这一点吗？
3. 编码器从未见过的类别。 只在测试集中加入第五个类别——机械松动（f_r 的许多次谐波，幅值逐次衰减），并对测试窗口的 \mathbf{h} 画出按类别着色的二维主成分分析（PCA）图。预训练得到的 \mathbf{h} 能否在不做任何重新训练的情况下，把新类别与旧类别分开？与 FFT 幅值的 PCA 图比较。

十五道练习按它们所练习的各节顺序排列，并按工作量分级：★ 是概念题，约 5 分钟，至多需要读出一个比值或一个和；★★ 是推导或计算，需要 10 到 12 分钟；★★★ 是编程，约 25 分钟。第一类有八道，第二类六道，第三类一道，合计 127 分钟。概念题先在脑中或纸上做完，再打开解答：它们之所以短，是因为全部难点都在于定下一个答案。

没有一道练习重复例题、正文检查题或实验的数字或情形。各节在一组数字上讲授每种方法；这里同一种方法遇到另一组数字，使能迁移的是方法而不是答案。解答中引用的每个数字都经过计算，产生它的代码或者列出，或者有所说明。唯一的编程练习（练习 12）从实验 4 的物理信息神经网络出发，并且自成一体：它的代码在一个新的 Python 会话中按原样即可运行。

每个解答在你打开之前都是隐藏的。只有在写下自己的答案之后，哪怕很粗糙，再打开它：真正有用的检查，是你的答案与解答分道扬镳的那个地方。

import numpy as np

p = np.array([0.6, 0.3, 0.1])        # p_data
q = np.array([0.2, 0.3, 0.5])        # p_g
d_star = p / (p + q)
value = (p * np.log(d_star)).sum() + (q * np.log(1 - d_star)).sum()
m = (p + q) / 2
jsd = 0.5 * (p * np.log(p / m)).sum() + 0.5 * (q * np.log(q / m)).sum()
print("D* =", d_star.round(4))
print(f"V = {value:.4f}   -log 4 + 2 JSD = {-np.log(4) + 2 * jsd:.4f}   JSD = {jsd:.4f}")

D* = [0.75   0.5    0.1667]
V = -1.1361   -log 4 + 2 JSD = -1.1361   JSD = 0.1251

import numpy as np
from scipy.spatial import cKDTree

rng = np.random.default_rng(0)
angles = 2 * np.pi * np.arange(8) / 8
centres = 2.0 * np.stack([np.cos(angles), np.sin(angles)], axis=1)   # 8 kinds of section

def draw(n, kinds):
    k = rng.choice(kinds, size=n)
    return centres[k] + 0.1 * rng.standard_normal((n, 2))

train_real = draw(2000, range(8))           # stands in for the real sections
held_out = draw(1000, range(8))             # held-out real sections
full = draw(1000, range(8))                 # a generator that covers every kind
collapsed = draw(1000, [0, 3, 5])           # perfect samples, but only 3 kinds

def median_nn(queries, reference):
    return np.median(cKDTree(reference).query(queries)[0])

print(f"real -> real (reference)   {median_nn(held_out, train_real[:1000]):.3f}")
d_ref = cKDTree(train_real[:1000]).query(held_out)[0]
for name, gen in (("covering generator", full), ("collapsed generator", collapsed)):
    d_gen = cKDTree(gen).query(held_out)[0]
    uncovered = np.mean(d_gen > np.quantile(d_ref, 0.99))
    print(f"{name}: sample->real {median_nn(gen, held_out):.3f}, "
          f"real->sample {median_nn(held_out, gen):.3f}, "
          f"uncovered real sections {uncovered:.3f}")

real -> real (reference)   0.016
covering generator: sample->real 0.016, real->sample 0.016, uncovered real sections 0.011
collapsed generator: sample->real 0.016, real->sample 1.141, uncovered real sections 0.621

import numpy as np

beta = np.array([0.1, 0.2, 0.3])
alpha = 1 - beta
alpha_bar = np.cumprod(alpha)

var = 0.0
for t in range(3):
    var = alpha[t] * var + beta[t]       # variance recursion for fixed x_0
    print(f"t={t + 1}: recursion {var:.4f}   1 - alpha_bar {1 - alpha_bar[t]:.4f}")

rng = np.random.default_rng(0)
x = np.ones(1_000_000)                   # one million chains started at x_0 = 1
for t in range(3):
    x = np.sqrt(alpha[t]) * x + np.sqrt(beta[t]) * rng.standard_normal(x.size)
print(f"sampled mean {x.mean():.4f} (sqrt(alpha_bar) = {np.sqrt(alpha_bar[2]):.4f}), "
      f"sampled variance {x.var():.4f}")

t=1: recursion 0.1000   1 - alpha_bar 0.1000
t=2: recursion 0.2800   1 - alpha_bar 0.2800
t=3: recursion 0.4960   1 - alpha_bar 0.4960
sampled mean 0.7097 (sqrt(alpha_bar) = 0.7099), sampled variance 0.4953

import numpy as np

def report(name, beta):
    alpha_bar = np.prod(1 - beta)
    print(f"{name:34s} sum(beta) {beta.sum():6.3f}  e^-sum {np.exp(-beta.sum()):.3e}  "
          f"alpha_bar_T {alpha_bar:.3e}  sqrt {np.sqrt(alpha_bar):.4f}")

report("T=1000, 1e-4 to 0.02", np.linspace(1e-4, 0.02, 1000))
report("T=100,  1e-4 to 0.02", np.linspace(1e-4, 0.02, 100))
report("T=100,  rescaled 1e-3 to 0.2", np.linspace(1e-3, 0.2, 100))

T, s = 100, 0.008                        # cosine schedule, s = 0.008
t = np.arange(T + 1) / T
f = np.cos((t + s) / (1 + s) * np.pi / 2) ** 2
report("T=100, cosine (clip 0.999)", np.clip(1 - f[1:] / f[:-1], 0, 0.999))

T=1000, 1e-4 to 0.02               sum(beta) 10.050  e^-sum 4.319e-05  alpha_bar_T 4.036e-05  sqrt 0.0064
T=100,  1e-4 to 0.02               sum(beta)  1.005  e^-sum 3.660e-01  alpha_bar_T 3.636e-01  sqrt 0.6030
T=100,  rescaled 1e-3 to 0.2       sum(beta) 10.050  e^-sum 4.319e-05  alpha_bar_T 2.039e-05  sqrt 0.0045
T=100, cosine (clip 0.999)         sum(beta)  7.879  e^-sum 3.787e-04  alpha_bar_T 2.429e-07  sqrt 0.0005

import numpy as np

A = np.zeros((4, 4))
A[0, 1:] = A[1:, 0] = 1                    # node 0 is the centre, 1-3 the leaves
A_tilde = A + np.eye(4)
d_tilde = A_tilde.sum(axis=1)              # (4, 2, 2, 2)
A_hat = A_tilde / np.sqrt(np.outer(d_tilde, d_tilde))
print("A_hat =\n", A_hat.round(4))
print("eigenvalues", np.linalg.eigvalsh(A_hat).round(4))

u = np.sqrt(d_tilde) / np.linalg.norm(np.sqrt(d_tilde))
for label, x in (("centre feature", np.array([1.0, 0, 0, 0])),
                 ("leaf feature  ", np.array([0, 1.0, 0, 0]))):
    limit = u * (u @ x)
    errors, h = [], x.copy()
    for k in range(1, 9):
        h = A_hat @ h
        errors.append(np.linalg.norm(h - limit) / np.linalg.norm(limit))
    print(label, "limit", limit.round(4))
    print("   relative error, k = 1..8:", " ".join(f"{e:.4f}" for e in errors))
    if label.startswith("centre"):
        print("   A_hat x   =", (A_hat @ x).round(4))
        print("   A_hat^2 x =", (A_hat @ A_hat @ x).round(4))

A_hat =
 [[0.25   0.3536 0.3536 0.3536]
 [0.3536 0.5    0.     0.    ]
 [0.3536 0.     0.5    0.    ]
 [0.3536 0.     0.     0.5   ]]
eigenvalues [-0.25  0.5   0.5   1.  ]
centre feature limit [0.4    0.2828 0.2828 0.2828]
   relative error, k = 1..8: 0.3062 0.0765 0.0191 0.0048 0.0012 0.0003 0.0001 0.0000
   A_hat x   = [0.25   0.3536 0.3536 0.3536]
   A_hat^2 x = [0.4375 0.2652 0.2652 0.2652]
leaf feature   limit [0.2828 0.2    0.2    0.2   ]
   relative error, k = 1..8: 0.9354 0.4593 0.2286 0.1142 0.0571 0.0285 0.0143 0.0071

import numpy as np, torch, torch.nn as nn
rng = np.random.default_rng(0)
torch.manual_seed(0)

def random_fault_tree(max_depth=4):
    """Return (edges, is_basic_event). Node 0 is the top gate; gates have 2-3 inputs."""
    edges, is_basic, frontier = [], [False], [(0, 0)]
    while frontier:
        node, depth = frontier.pop()
        for _ in range(rng.integers(2, 4)):          # a gate has 2 or 3 inputs
            child = len(is_basic)
            leaf = depth + 1 >= max_depth or rng.random() < 0.45
            is_basic.append(bool(leaf))
            edges.append((node, child))
            if not leaf:
                frontier.append((child, depth + 1))
    return edges, np.array(is_basic)

def graph_tensors(edges, n):
    A = np.zeros((n, n), dtype=np.float32)
    for a, b in edges:
        A[a, b] = A[b, a] = 1.0
    At = A + np.eye(n, dtype=np.float32)
    dinv = 1.0 / np.sqrt(At.sum(1))
    return torch.tensor(dinv[:, None] * At * dinv[None, :]), A.sum(1)

trees = [random_fault_tree() for _ in range(300)]
data = []
for edges, basic in trees:
    n = len(basic)
    A_hat, deg = graph_tensors(edges, n)
    data.append((A_hat, torch.tensor(deg[:, None] / 4.0, dtype=torch.float32),
                 torch.tensor(basic, dtype=torch.float32)))
sizes = [len(b) for _, b in trees]
print("nodes per tree: min", min(sizes), "max", max(sizes))

# the baseline: a node with exactly one neighbour is a basic event
rule_correct = sum(int(((d[1].squeeze() == 0.25) == d[2].bool()).sum()) for d in data)
print(f"leaf rule accuracy: {rule_correct / sum(sizes):.4f}")

class GCN(nn.Module):
    def __init__(self, hidden=16):
        super().__init__()
        self.w1, self.w2 = nn.Linear(1, hidden), nn.Linear(hidden, 1)
    def forward(self, A_hat, x):
        h = torch.relu(A_hat @ self.w1(x))
        return (A_hat @ self.w2(h)).squeeze(-1)

train, test = data[:200], data[200:]            # split by tree, not by node
model = GCN()
opt = torch.optim.Adam(model.parameters(), lr=1e-2)
for epoch in range(300):
    for A_hat, x, y in train:
        loss = nn.functional.binary_cross_entropy_with_logits(model(A_hat, x), y)
        opt.zero_grad(); loss.backward(); opt.step()

def accuracy(split):
    ok = tot = 0
    with torch.no_grad():
        for A_hat, x, y in split:
            ok += int(((model(A_hat, x) > 0).float() == y).sum()); tot += len(y)
    return ok / tot

share = np.mean(np.concatenate([b for _, b in trees]))
print(f"GCN test accuracy: {accuracy(test):.4f}")
print(f"share of basic events (majority baseline): {max(share, 1 - share):.4f}")

nodes per tree: min 3 max 58
leaf rule accuracy: 1.0000
GCN test accuracy: 0.8306
share of basic events (majority baseline): 0.6237