模块三：可微分推断——融入深度学习框架

3.1 广义 EM 与梯度优化

当发射概率 $p(x_t \mid z_t; \phi)$ 由神经网络给出时，M 步通常不再有闭式解析解。此时我们可采用广义 EM：在固定当前后验估计 $\gamma_t(i)$ 的条件下，对参数 $\phi$ 执行若干步梯度上升。具体地，令

\mathcal{L}_M(\phi) = \sum_{t=1}^T \sum_{i=1}^K \gamma_t(i) \log p_\phi(x_t \mid z_t=i),

然后按照 $\phi \leftarrow \phi + \eta \nabla_\phi \mathcal{L}_M(\phi)$ 更新参数，其中 $\eta > 0$ 为学习率。梯度可通过自动微分框架轻松求得。广义 EM 虽然放弃了 M 步的精确最大化，但仍保证每次迭代不降低似然下界。

3.2 可微分的 HMM 与对数空间实现

若希望将整个 HMM 的学习过程嵌入现代深度学习框架（如 PyTorch）并利用自动微分端到端训练，前向–后向算法必须完全可微。直接使用概率值相乘极易因数值下溢而失败，因为概率均为小于等于 1 的数，长序列连乘后迅速趋向于零。标准解决方案是转入对数空间，并借助 LogSumExp 函数。

对于任意实数向量 $\mathbf{v} = (v_1,\dots,v_K)$，定义：

\text{LogSumExp}(\mathbf{v}) := \log \sum_{i=1}^K e^{v_i} = c + \log \sum_{i=1}^K e^{v_i - c},

其中 $c = \max_i v_i$。减去最大值 $c$ 可避免指数运算溢出，同时保持严格的数学等价。

记对数前向消息 $\hat{\alpha}_t(i) = \log \alpha_t(i)$，对数后向消息 $\hat{\beta}_t(i) = \log \beta_t(i)$。前向递推 (12) 变为：

\hat{\alpha}_t(j) = \log p(x_t \mid z_t=j) + \text{LogSumExp}_{i=1}^K \bigl( \hat{\alpha}_{t-1}(i) + \log A_{ij} \bigr). \tag{22}

后向递推 (14) 变为：

\hat{\beta}_t(i) = \text{LogSumExp}_{j=1}^K \bigl( \log A_{ij} + \log p(x_{t+1} \mid z_{t+1}=j) + \hat{\beta}_{t+1}(j) \bigr). \tag{23}

对数似然直接由 $\hat{\alpha}_T$ 得到：

\log p(\mathbf{X} \mid \theta) = \text{LogSumExp}_{i=1}^K \bigl(\hat{\alpha}_T(i)\bigr). \tag{24}

为了在对数空间中计算 $\gamma_t(i)$ 和 $\xi_t(i,j)$，定义 $s_t(i) := \hat{\alpha}_t(i) + \hat{\beta}_t(i)$，注意到 $\exp(s_t(i)) = \alpha_t(i)\beta_t(i)$ 且 $\sum_i \exp(s_t(i)) = p(\mathbf{X} \mid \theta)$，于是：

\gamma_t(i) = \frac{\exp(s_t(i))}{\sum_{k=1}^K \exp(s_t(k))} = \text{Softmax}(s_t)_i. \tag{25}

类似地，定义 $u_t(i,j) := \hat{\alpha}_t(i) + \log A_{ij} + \log p(x_{t+1} \mid z_{t+1}=j) + \hat{\beta}_{t+1}(j)$，则：

\xi_t(i,j) = \frac{\exp(u_t(i,j))}{\sum_{i',j'} \exp(u_t(i',j'))} = \text{Softmax}(u_t)_{i,j}. \tag{26}

这一整套操作完全规避了下溢问题，并且所有运算均可通过标准的自动微分算子构建。将前向–后向封装为一个 torch.nn.Module，仅需调用 loss = -log_likelihood; loss.backward() 即可完成参数学习。这使得经典的 Baum–Welch 算法在现代深度学习框架中获得了新生。

3.3 计算图的深度展开与隐式微分

3.3.1 问题的提出：双层优化与超梯度

在元学习或超参数优化中，我们常遇到双层优化结构：内层用 EM 迭代求解模型参数 $\theta$，外层优化超参数 $\lambda$（如正则化系数、先验参数）。形式上，外层目标为

\mathcal{L}_{\text{outer}}\bigl(\theta^*(\lambda),\, \lambda\bigr),

其中 $\theta^*(\lambda)$ 是内层 EM 在给定 $\lambda$ 下收敛到的不动点。要对外层做梯度下降，就必须计算超梯度 $\dfrac{d\mathcal{L}_{\text{outer}}}{d\lambda}$。由链式法则：

\frac{d\mathcal{L}_{\text{outer}}}{d\lambda} = \frac{\partial \mathcal{L}_{\text{outer}}}{\partial \lambda} + \frac{\partial \mathcal{L}_{\text{outer}}}{\partial \theta^*}\, \frac{\partial \theta^*}{\partial \lambda}. \tag{27}

第一项可直接由自动微分得到，真正的难点在第二项中的 $\dfrac{\partial \theta^*}{\partial \lambda}$：它要求我们知道"当超参数 $\lambda$ 略微变化时，EM 收敛点 $\theta^*$ 会如何移动"。

3.3.2 朴素方法的困境：为什么不能直接展开计算图？

最直接的想法是把 EM 的每一步迭代都记录到计算图里，让自动微分一路反向传播回 $\lambda$。但 EM 往往需要成百上千步才收敛，这意味着：

内存爆炸：计算图深度等于迭代步数，每步的中间激活值都要保存，$O(T)$ 的内存开销在长迭代下难以承受；
梯度失真：过深的链式求导会遭遇梯度消失或爆炸，传到 $\lambda$ 时信号已不可靠。

这里暴露出一个根本矛盾：我们真正关心的只是收敛后的不动点 $\theta^*$，却被迫为整段迭代历史付出代价。 理想的做法应当是——绕开迭代过程，直接在不动点处求导。这正是隐函数定理发挥作用的地方。

3.3.3 不动点条件：$\theta^*$ 是被方程"隐式"定义的

EM 收敛到的不动点 $\theta^*$ 满足平稳点条件，即期望完全数据对数似然的梯度为零：

F(\theta^*, \lambda) := \nabla_\theta\, \mathcal{L}(\theta^*, \lambda) = \mathbf{0}. \tag{28}

这是一个关于 $(\theta, \lambda)$ 的方程组。关键观察在于：方程 $F(\theta, \lambda) = \mathbf{0}$ 通常无法解出 $\theta$ 的显式表达式——我们写不出形如 $\theta^* = g(\lambda)$ 的封闭公式。然而，从直觉上讲，只要 $\lambda$ 给定，EM 迭代总能收敛到一个确定的 $\theta^*$，这说明 $\theta^*$ 确实"依赖"于 $\lambda$。

这种"被方程约束、却写不出显式公式"的依赖关系，正是隐函数的典型形态。方程 $F(\theta, \lambda) = \mathbf{0}$ 就像一张未解出的方程，把 $\theta$ 隐藏地定义为 $\lambda$ 的函数 $\theta = \theta^*(\lambda)$——它存在，但不显形。

3.3.4 隐函数定理：保证隐函数存在且可导

仅凭直觉不够，我们需要严格的数学保证：$\theta^*(\lambda)$ 是否真的存在？是否可微？导数如何计算？隐函数定理一次性回答了这三个问题。

隐函数定理（多变量情形）

设 $F: \mathbb{R}^p \times \mathbb{R}^q \to \mathbb{R}^p$ 连续可微，且在点 $(\theta_0, \lambda_0)$ 处满足 $F(\theta_0, \lambda_0) = \mathbf{0}$。若 $F$ 关于 $\theta$ 的雅可比矩阵 $J_\theta F = \dfrac{\partial F}{\partial \theta}$ 在该点非奇异（即可逆），则存在 $(\theta_0, \lambda_0)$ 的一个邻域，使得：

1.（存在唯一性）方程 $F(\theta, \lambda) = \mathbf{0}$ 在该邻域内唯一确定了一个连续可微的函数 $\theta = \theta^*(\lambda)$，满足 $\theta^*(\lambda_0) = \theta_0$ 且 $F(\theta^*(\lambda), \lambda) \equiv \mathbf{0}$；

2.（导数公式）该隐函数的导数为

\frac{\partial \theta^*}{\partial \lambda} = -\left( \frac{\partial F}{\partial \theta} \right)^{-1} \frac{\partial F}{\partial \lambda}.

把这一定理对应到我们的场景：

$F(\theta, \lambda) = \nabla_\theta \mathcal{L}(\theta, \lambda)$，平稳点 $(\theta^*, \lambda_0)$ 满足 $F = \mathbf{0}$；
雅可比 $\dfrac{\partial F}{\partial \theta}$ 即对数似然的 Hessian 矩阵 $H$。只要 $\mathcal{L}$ 在 $\theta^*$ 处是严格凹的（局部极大），$H$ 就是负定可逆的，定理的非奇异条件成立；
于是定理保证：在 $\lambda_0$ 附近，确实存在一个可微的隐函数 $\theta^*(\lambda)$，并且给出了它的导数公式。

这就是"隐函数在哪里"的答案——它藏在不动点方程 $F = \mathbf{0}$ 背后，由隐函数定理担保其存在性与可微性。

3.3.5 为什么必须用隐函数定理？

理解了定理的内容，我们就能看清它在此处不可替代的作用：

$\theta^*(\lambda)$ 没有显式表达式。 EM 是迭代算法，不动点是极限，无法写成 $\theta^* = g(\lambda)$ 的封闭形式。既然没有显式公式，常规的"先写出函数再求导"这条路走不通。
隐函数定理把"求不出显式函数"的困境转化为"只需方程本身的信息"。 它的导数公式只用到 $F$ 的两个偏导数 $\dfrac{\partial F}{\partial \theta}$ 和 $\dfrac{\partial F}{\partial \lambda}$，完全不要求知道 $\theta^*(\lambda)$ 长什么样。这正是"隐式微分"得名的原因——我们绕过了隐函数本身，直接从定义它的方程中提取导数。
它同时提供了存在性保障。 在动手计算 $\dfrac{\partial \theta^*}{\partial \lambda}$ 之前，定理先确认了这个导数是有意义的（隐函数确实存在且可微），否则一切计算都是空中楼阁。

简言之，隐函数定理在这里扮演了"桥梁"：一头连着无法显式求解的不动点方程，另一头连着可计算的超梯度公式。

3.3.6 超梯度的推导

现在把定理的导数公式代入链式法则 (27)。对恒等式 $F(\theta^*(\lambda), \lambda) \equiv \mathbf{0}$ 两边关于 $\lambda$ 求全导数：

\frac{\partial F}{\partial \theta}\, \frac{\partial \theta^*}{\partial \lambda} + \frac{\partial F}{\partial \lambda} = \mathbf{0} \quad\Longrightarrow\quad \frac{\partial \theta^*}{\partial \lambda} = -\left( \frac{\partial F}{\partial \theta} \right)^{-1} \frac{\partial F}{\partial \lambda}. \tag{29}

这一步其实就是隐函数定理导数公式的现场推导。将其代入 (27)，得到超梯度：

\begin{aligned} \frac{d\mathcal{L}_{\text{outer}}}{d\lambda} &= \frac{\partial \mathcal{L}_{\text{outer}}}{\partial \lambda} - \frac{\partial \mathcal{L}_{\text{outer}}}{\partial \theta^*} \left( \frac{\partial F}{\partial \theta} \right)^{-1} \frac{\partial F}{\partial \lambda}. \end{aligned} \tag{30}

注意式中每一项都只涉及不动点处的局部信息：$\dfrac{\partial \mathcal{L}_{\text{outer}}}{\partial \theta^*}$、Hessian $H = \dfrac{\partial F}{\partial \theta}$、$\dfrac{\partial F}{\partial \lambda}$，完全不需要 EM 的迭代历史。

3.3.7 伴随向量法：避免显式求 Hessian 逆

公式 (30) 含有 $H^{-1}$。在高维参数空间中（$\theta$ 可能有数百万维），显式求逆代价不可承受。为此引入伴随向量 $\mathbf{v}$，令它满足线性系统：

H^\top \mathbf{v} = \left( \frac{\partial \mathcal{L}_{\text{outer}}}{\partial \theta^*} \right)^{\!\top}. \tag{31}

于是 $\mathbf{v}^\top = \dfrac{\partial \mathcal{L}_{\text{outer}}}{\partial \theta^*}\, H^{-1}$，超梯度 (30) 改写为：

\frac{d\mathcal{L}_{\text{outer}}}{d\lambda} = \frac{\partial \mathcal{L}_{\text{outer}}}{\partial \lambda} + \mathbf{v}^\top \frac{\partial F}{\partial \lambda}. \tag{32}

线性系统 (31) 不必显式构造 $H$，可用共轭梯度法等迭代求解器，每次只需计算 Hessian–向量乘积 $H\mathbf{p}$（这可通过 Pearlmutter 技巧以一次前向加一次反向传播实现），从而彻底回避矩阵求逆。

3.3.8 小结：隐式微分的优势

整个方法的逻辑链可以浓缩为一句话：既然不动点方程 $F=\mathbf{0}$ 把 $\theta^*$ 隐式地绑定为 $\lambda$ 的函数，就用隐函数定理直接求这个隐函数的导数，而不去展开迭代过程。

由此带来的收益是：

内存：从展开计算图的 $O(T)$ 降为 $O(1)$，只保存不动点处的若干向量；
精度：不依赖迭代链的梯度传播，规避了梯度消失/爆炸；
通用性：只要内层算法收敛到不动点（不限于 EM，也包括不动点迭代、求根等），同样的框架即可适用。

这种基于隐函数定理的方法被称为隐式微分（implicit differentiation），它为深度概率编程中的双层优化——元学习、超参数优化、约束优化——提供了坚实的数学基础。

上一节1.2 HMM 的图视角下一节1.4 神经参数化与生成模型