面对配分函数 - 估计配分函数篇

序言

配分函数，作为统计物理和热力学中的核心概念，扮演着举足轻重的角色。它不仅是系统微观状态与宏观性质之间的桥梁，更是我们理解和预测物质宏观行为的关键工具。配分函数包含了系统所有可能微观状态的统计信息，通过对其进行分析和计算，我们可以得出系统的能量、熵、自由能等重要热力学量，进而揭示系统的相变、稳定性等宏观特性。

然而，配分函数的精确计算往往是一项艰巨的任务，特别是对于复杂系统而言。因此，估计配分函数成为了研究者们关注的焦点。估计配分函数不仅可以帮助我们近似地获得系统的热力学性质，还可以为新材料的设计、新药物的研发等领域提供理论支持。近年来，随着计算机技术的发展和数值方法的进步，配分函数的估计方法也日益丰富和完善。

估计配分函数

• 估计配分函数可能会很重要，当我们希望计算数据的归一化似然时，我们会需要它。在评估模型，监控训练性能，和比较模型时，这通常是很重要的。

• 例如，假设我们有两个模型：

  • 概率分布为 $p_A(\text{x};{\mathbf{\theta }}_A)=\frac{1}{Z_A}(\text{x};{\mathbf{\theta }}_A)$ 的模型 ${\mathcal{M}}_A$和概率分布为 $p_B(\text{x};{\mathbf{\theta }}_B)=\frac{1}{Z_B}{\tilde{p}}_B(\text{x};{\mathbf{\theta }}_B)$ 的模型 ${\mathcal{M}}_B$。

  • 比较模型的常用方法是评估和比较两个模型分配给独立同分布测试数据集的似然。

  • 假设测试集含 $m$ 个样本 { x ( 1 ) , … , x ( m ) } \{\boldsymbol{x}^{(1)},\dots,\boldsymbol{x}^{(m)}\} { x(1),…,x(m)}。

  • 如果：
    ${\prod }_i{p}_A({\text{x}}^{(i)};{\mathbf{\theta }}_A)>{\prod }_i{p}_B({\text{x}}^{(i)};{\mathbf{\theta }}_B)$，
    或等价地，如果：
    ${\sum }_i\log p_A({\text{x}}^{(i)};{\mathbf{\theta }}_A)-{\sum }_i\log p_B({\text{x}}^{(i)};{\mathbf{\theta }}_B)>0$ $\text{—公式1}$

  • 那么我们说 ${\mathcal{M}}_A$ 是一个比 ${\mathcal{M}}_B$ 更好的模型（或者，至少可以说，它在测试集上是一个更好的模型），这是指它有一个更大的测试对数似然。

  • 然而，测试这个条件是否成立需要知道配分函数。

  • 公式1看起来需要估计模型分配给每个点的对数概率，因而需要估计配分函数。

  • 我们可以通过将公式1重新转化为另一种形式来简化情况，在该形式中我们只需要知道两个模型的配分函数的比率：
    $\sum _i\log p_A({\text{x}}^{(i)};{\mathbf{\theta }}_A)-\sum _i\log p_B({\text{x}}^{(i)};{\mathbf{\theta }}_B)=\sum _i\left(\log \frac{{\tilde{p}}_A({\text{x}}^{(i)};{\mathbf{\theta }}_A)}{p_B({\text{x}}^{(i)};{\mathbf{\theta }}_B)}\right)-m\log \frac{Z({\mathbf{\theta }}_A)}{Z({\mathbf{\theta }}_B)}$ $\text{—公式2}$

• 因此，我们可以在不知道任一模型的配分函数，而只知道它们比率的情况下，判断模型 ${\mathcal{M}}_A$ 是否比模型 ${\mathcal{M}}_B$ 更优。正如我们将很快看到的，在两个模型相似的情况下，我们可以使用重要采样来估计比率。

• 然而，如果我们想要计算测试数据在 ${\mathcal{M}}_A$ 或 ${\mathcal{M}}_B$ 上的真实概率，我们需要计算配分函数的真实值。如果我们知道两个配分函数的比率， $r=\frac{Z({\mathbf{\theta }}_B)}{Z({\mathbf{\theta }}_A)}$，并且我们知道两者中一个的实际值，比如说 $Z({\mathbf{\theta }}_A)$，那么我们可以计算另一个的值：
  $Z({\mathbf{\theta }}_B)=rZ({\mathbf{\theta }}_A)=\frac{Z({\mathbf{\theta }}_B)}{Z({\mathbf{\theta }}_A)}Z({\mathbf{\theta }}_A)$ $\text{—公式3}$

• 一种估计配分函数的简单方法是使用蒙特卡罗方法，例如简单重要采样。

  • 以下用连续变量积分来表示该方法，也可以替换积分为求和，很容易将其应用到离散变量的情况。

  • 我们使用提议分布 $p_0(\text{x})=\frac{1}{Z_0}{\tilde{p}}_0(\text{x})$，其在配分函数 $Z_0$ 和未归一化分布${\tilde{p}}_0(\text{x})$ 上易于采样和估计。

  • $\{\begin{array}{l}\begin{array}{rlr}{Z}_1& =\int {\tilde{p}}_1(\text{x})\text{dx}& \text{—公式4}\\ & =\int \frac{p_0(\text{x})}{p_0(\text{x})}{\tilde{p}}_1(\text{x})\text{dx}& \text{—公式5}\\ & =Z_0\int p_0(\text{x})\frac{{\tilde{p}}_1(\text{x})}{{\tilde{p}}_0(\text{x})}\text{dx}& \text{—公式6}\\ & & \end{array}\end{array}$

  • ${\hat{Z}}_1=\frac{Z_0}{K}\sum _{k=1}^K\frac{{\tilde{p}}_1({\text{x}}^{(k)})}{{\tilde{p}}_0({\text{x}}^{(k)})}\text{s.t. : }{\text{x}}^{(k)}\sim p_0$ $\text{—公式7}$

  • 在最后一行，我们使用蒙特卡罗估计，使用从 $p_0(\text{x})$ 中抽取的采样计算积分 ${\hat{Z}}_1$，然后用未归一化的 ${\tilde{p}}_1$ 和提议分布 $p_0$ 的比率对每个采样加权。

• 这种方法使得我们可以估计配分函数之间的比率：
  $\frac{1}{K}\sum _{k=1}^K\frac{{\tilde{p}}_1({\text{x}}^{(k)})}{{\tilde{p}}_0({\text{x}}^{(k)})}\text{s.t. : }{\text{x}}^{(k)}\sim p_0$ $\text{—公式8}$

• 然后该值可以直接比较公式2中的两个模型。

• 如果分布 $p_0$ 接近 $p_1$，那么公式7能够高效地估计配分函数($\text{Minka, 2005}$)。

  • 不幸的是，大多数时候 $p_1$ 都很复杂（通常是多峰的），并且定义在高维空间中。
  • 很难找到一个易求解的 $p_0$，既能易于评估，又能充分接近 $p_1$ 以保持高质量的近似。

• 如果 $p_0$ 和 $p_1$ 不接近，那么 $p_0$ 的大多数采样将在 $p_1$ 中具有较低的概率，从而在公式7的求和中产生（相对的）可忽略的贡献。

• 如果求和中只有少数几个具有显著权重的样本，那么将会由于高方差而导致估计的效果很差。这可以通过估计 ${\hat{Z}}_1$ 的方差来定量地理解：
  $\widehat{\text{Var}}({\hat{Z}}_1)=\frac{Z_0}{K^2}\sum _{k=1}^K{\left(\frac{{\tilde{p}}_1({\text{x}}^{(k)})}{{\tilde{p}}_0({\text{x}}^{(k)})}-{\hat{Z}}_1\right)}^2$ $\text{—公式9}$

• 当重要性权重 $\frac{{\tilde{p}}_1({\text{x}}^{(k)})}{{\tilde{p}}_0({\text{x}}^{(k)})}$ 存在显著偏差时，上式的值是最大的。

• 有两个解决高维空间复杂分布上估计配分函数的方法：退火重要采样和桥式采样。两者都始于上面介绍的简单的重要采样方法，并且都试图通过引入桥接 $p_0$ 和 $p_1$之间差距的中间分布，来解决 $p_0$ 远离 $p_1$ 的问题。

退火重要采样

• 在$D_{KL}(p_0||p_1)$很大的情况下（即，$p_0$和$p_1$之间几乎没有重叠），一种称为退火重要采样（$\text{annealed importance sampling， AIS}$）的方法试图通过引入中间分布来桥接差距 ($\text{Jarzynski, 1997; Neal, 2001}$)。考虑分布序列 $p_{{\eta }_0},\dots ,p_{{\eta }_n}$，其中$0={\eta }_0<{\eta }_1<\cdots <{\eta }_n-1<{\eta }_n=1$，分布序列中的第一个和最后一个分别是 $p_0$ 和$p_1$。

• 这种方法使我们能够估计定义在高维空间多峰分布（例如训练$\text{RBM}$时定义的分布）上的配分函数。我们从一个已知配分函数的简单模型（例如，权重为零的$\text{RBM}$）开始，估计两个模型配分函数之间的比率。该比率的估计基于许多个相似分布的比率估计，例如在零和学习到的权重之间插值一组权重不同的$\text{RBM}$。

• 现在我们可以将比率 $\frac{Z_1}{Z_0}$写作：

$\{\begin{array}{l}\begin{array}{rlr}\frac{Z_1}{Z_0}& =\frac{Z_1}{Z_0}\frac{Z_{{\eta }_1}}{Z_{{\eta }_1}}\dots \frac{Z_{{\eta }_{n-1}}}{Z_{{\eta }_{n-1}}}& \text{—公式10}\\ & & \\ & =\frac{Z_{{\eta }_1}}{Z_{{\eta }_0}}\frac{Z_{{\eta }_2}}{Z_{{\eta }_1}}\dots \frac{Z_{{\eta }_{n-1}}}{Z_{{\eta }_{n-2}}}\frac{Z_1}{Z_{{\eta }_{n-1}}}& \text{—公式11}\\ & & \\ & =\prod _{j=0}^{n-1}\frac{Z_{{\eta }_{j+1}}}{Z_{{\eta }_j}}& \text{—公式12}\end{array}\end{array}$
$\text{—公式12}$

• 如果对于所有的 $0\le j\le n-1$，分布 $p_{{\eta }_j}$ 和 $p_{{\eta }_{j+1}}$ 足够接近，那么我们能够使用简单的重要采样来估计每个因子 $\frac{Z_{{\eta }_{j+1}}}{Z_{{\eta }_j}}$，然后使用这些得到 $\frac{Z_1}{Z_0}$ 的估计。

• 这些中间分布是从哪里来的呢？

  • 正如最先的提议分布 $p_0$ 是一种设计选择，分布序列 $p_{{\eta }_1}\dots p_{{\eta }_{n-1}}$ 也是如此。

  • 也就是说，它们可以被特别设计为特定的问题领域。

  • 中间分布的一个通用目标和流行选择是使用目标分布 $p_1$ 的加权几何平均，起始分布（其配分函数是已知的）为 $p_0$：
    $p_{{\eta }_j}\propto p_1^{{\eta }_j}{p}_0^{1-{\eta }_j}$ $\text{—公式13}$

  • 为了从这些中间分布中采样，我们定义了一组马尔可夫链转移函数 $T_{{\eta }_j}(x^{\prime }\mid x)$，定义了给定 $x$ 转移到 $x^{\prime }$ 的条件概率分布。

  • 转移算子 $T_{{\eta }_j}(x^{\prime }\mid x)$ 定义如下，保持$p_{{\eta }_j}(x)$ 不变：
    $p_{{\eta }_j}(x)=\int p_{{\eta }_j}(x^{\prime })T_{{\eta }_j}(x^{\prime }\mid x)\text{d}{x}^{\prime }$ $\text{—公式14}$

  • 这些转移可以被构造为任何马尔可夫链蒙特卡罗方法（例如， $\text{Metropolis-Hastings，Gibbs}$），包括涉及多次遍历所有随机变量或其他迭代的方法。

  • 然后， $\text{AIS}$采样方法从 $p_0$ 开始生成样本，并使用转移算子从中间分布顺序地生成采样，直到我们得到目标分布 $p_1$ 的采样：

    • 对于$k=1\dots K$

      $-$ 采样$x_{{\eta }_1}^{(k)}\sim p_0(\text{x})$
      $-$ 采样$x_{{\eta }_2}^{(k)}\sim T_{{\eta }_j}({\text{x}}_{{\eta }_2}^{(k)}\mid x_{{\eta }_1}^{(k)})$
      $-\dots$
      $-$ 采样$x_{{\eta }_{n-1}}^{(k)}\sim T_{{\eta }_{n-2}}({\text{x}}_{{\eta }_{n-1}}^{(k)}\mid x_{{\eta }_{n-2}}^{(k)})$
      $-$ 采样$x_{{\eta }_n}^{(k)}\sim T_{{\eta }_{n-1}}({\text{x}}_{{\eta }_n}^{(k)}\mid x_{{\eta }_{n-1}}^{(k)})$

    • 结束

  • 对于采样 $k$，通过连接公式12给出的中间分布之间的重要性权重，我们可以导出目标重要性权重：
    $w^{(k)}=\frac{{\tilde{p}}_{{\eta }_1}(x_{{\eta }_1}^{(k)})}{{\tilde{p}}_0(x_{{\eta }_1}^{(k)})}\frac{{\tilde{p}}_{{\eta }_2}(x_{{\eta }_2}^{(k)})}{{\tilde{p}}_{{\eta }_1}(x_{{\eta }_2}^{(k)})}\dots \frac{{\tilde{p}}_1(x_1^{(k)})}{{\tilde{p}}_{{\eta }_{n-1}}(x_{{\eta }_n}^{(k)})}$ $\text{—公式15}$

• 为了避免诸如上溢的数值问题，最佳方法可能是通过加法或减法计算 $\log w^{(k)}$，而不是通过概率乘法和除法计算 $w^{(k)}$。

• 利用由此定义的抽样过程和公式15中给出的重要性权重，配分函数的比率估计如下所示：
  $\frac{Z_1}{Z_0}\approx \frac{1}{K}\sum _{k=1}^K{w}^{(k)}$ $\text{—公式16}$

• 为了验证该过程定义的重要采样方案是否有效，我们可以展示 ($\text{Neal,2001}$) $\text{AIS}$过程对应着扩展状态空间上的简单重要采样，其中数据点采样自乘积空间 $[x_{{\eta }_1},\dots ,x_{{\eta }_{n-1}},x_1]$。为此，我们将扩展空间上的分布定义为：
  $$\begin{gathered} \tilde{p}(x_{{\eta }_1},\dots ,x_{{\eta }_{n-1}},x_1) \\ ={\tilde{p}}_1(x_1){\tilde{T}}_{{\eta }_{n-1}}(x_{{\eta }_{n-1}}\mid x_1){\tilde{T}}_{{\eta }_{n-2}}(x_{{\eta }_{n-2}}\mid x_{{\eta }_{n-1}})\dots {\tilde{T}}_{{\eta }_1}(x_{{\eta }_1}\mid x_{{\eta }_2}) \end{gathered}$$ $\text{—公式17}$

• 其中 ${\tilde{T}}_a$ 是由 $T_a$ 定义的转移算子的逆（应用贝叶斯规则）：
  ${\tilde{T}}_a(x^{\prime }\mid x)=\frac{p_a(x^{\prime })}{p_a(x)}{T}_a(x\mid x^{\prime })=\frac{{\tilde{p}}_a(x^{\prime })}{p_a(x)}{T}_a(x\mid x^{\prime })$ $\text{—公式18}$

• 将以上代入到公式17给出的扩展状态空间上的联合分布中，我们得到：
  $\{\begin{array}{l}\begin{array}{rlr}& \tilde{p}(x_{{\eta }_1},\dots ,x_{{\eta }_{n-1}},x_1)& \text{—公式19}\\ & ={\tilde{p}}_1(x_1)\frac{{\tilde{p}}_{{\eta }_{n-1}}(x_{{\eta }_{n-1}})}{{\tilde{p}}_{{\eta }_{n-1}}(x_1)}{T}_{{\eta }_{n-1}}(x_1\mid x_{{\eta }_{n-1}})\prod _{i=1}^{n-2}\frac{{\tilde{p}}_{{\eta }_i}(x_{{\eta }_i})}{{\tilde{p}}_{{\eta }_i}(x_{{\eta }_{i+1}})}{T}_{{\eta }_i}(x_{{\eta }_{i+1}}\mid x_{{\eta }_i})& \text{—公式20}\\ & =\frac{{\tilde{p}}_{{\eta }_1}(x_1)}{{\tilde{p}}_{{\eta }_{n-1}}(x_1)}{T}_{{\eta }_{n-1}}(x_1\mid x_{{\eta }_{n-1}}){\tilde{p}}_{{\eta }_1}(x_{{\eta }_1})\prod _{i=1}^{n-2}\frac{{\tilde{p}}_{{\eta }_{i+1}}(x_{{\eta }_{i+1}})}{{\tilde{p}}_{{\eta }_i}(x_{{\eta }_{i+1}})}{T}_{{\eta }_i}(x_{{\eta }_{i+1}}\mid x_{{\eta }_i})& \text{—公式21}\end{array}\end{array}$
  $\text{—公式21}$

• 通过上面给定的采样方案，现在我们可以从扩展样本上的联合提议分布 $q$ 上生成采样，联合分布如下：
  $q(x_{{\eta }_1},\dots ,x_{{\eta }_{n-1}},x_1)=p_0(x_{{\eta }_1})T_{{\eta }_1}(x_{{\eta }_2}\mid x_{{\eta }_1})\dots T_{{\eta }_{n-1}}(x1\mid x_{{\eta }_{n-1}})$ $\text{—公式22}$

• 公式21给出了扩展空间上的联合分布。将 $q(x_{{\eta }_1},\dots ,x_{{\eta }_{n-1}},x_1)$ 作为扩展状态空间上的提议分布（我们会从中抽样），重要性权重如下：
  $w^{(k)}=\frac{\tilde{p}(x_{{\eta }_1},\dots ,x_{{\eta }_{n-1}},x_1)}{q(x_{{\eta }_1},\dots ,x_{{\eta }_{n-1}},x_1)}=\frac{{\tilde{p}}_1(x_1^{(k)})}{{\tilde{p}}_{{\eta }_{n-1}}(x_{{\eta }_{n-1}}^{(k)})}\dots \frac{{\tilde{p}}_{{\eta }_2}(x_{{\eta }_2}^{(k)})}{{\tilde{p}}_{{\eta }_1}(x_{{\eta }_1}^{(k)})}\frac{{\tilde{p}}_{{\eta }_1}(x_{{\eta }_1}^{(k)})}{{\tilde{p}}_0(x_0^{(k)})}$ $\text{—公式23}$

• 这些权重和$\text{AIS}$上的权重相同。因此，我们可以将$\text{AIS}$解释为应用于扩展状态上的简单重要采样，其有效性紧紧来源于重要采样的有效性。

• 退火重要采样首先由$\text{Jarzynski (1997)}$ 发现，然后由$\text{Neal (2001)}$ 再次独立发现。目前它是估计无向概率模型的配分函数的最常用方法。其原因可能与一篇有影响力的论文 ($\text{Salakhutdinov and Murray, 2008}$) 有关，该论文并没有讨论该方法相对于其他方法的优点，而是介绍了将其应用于估计受限玻尔兹曼机和深度信念网络的配分函数。

• 关于$\text{AIS}$估计性质（例如，方差和效率）的讨论，请参看$\text{Neal (2001)}$。

桥式采样

• 类似于$\text{AIS}$， 桥式采样($\text{Bennett, 1976}$) 是另一种处理重要采样缺点的方法。并非将一系列中间分布连接在一起， 桥式采样依赖于单个分布 $p_{\ast }$（被称为桥），在已知配分函数的分布 $p_0$ 和分布 $p_1$（我们试图估计其配分函数 $Z_1$）之间插值。

• 桥式采样估计比率 $Z_1/Z_0$： ${\tilde{p}}_0$ 和 ${\tilde{p}}_{\ast }$ 之间重要性权重期望与 ${\tilde{p}}_1$ 和 ${\tilde{p}}_{\ast }$ 之间重要性权重的比率：
  $\frac{Z_1}{Z_0}\approx \sum _{k=1}^K\frac{{\tilde{p}}_{\ast }(x_0^{(k)})}{{\tilde{p}}_0(x_0^{(k)})}/\sum _{k=1}^K\frac{{\tilde{p}}_{\ast }(x_1^{(k)})}{{\tilde{p}}_1(x_1^{(k)})}$ $\text{—公式24}$

• 如果仔细选择桥式采样$p_{\ast }$，使其与$p_0$和$p_1$都有很大重合的话，那么桥式采样能够允许两个分布（或更正式地，$D_{\text{KL}}(p_0||p_1)$ ）之间有较大差距（相对标准重要采样而言）。

• 可以表示，最优的桥式采样是$p_{\ast }^{(opt)}(\text{x})\approx \frac{{\tilde{p}}_0(x){\tilde{p}}_1(x)}{r{\tilde{p}}_0(x)+{\tilde{p}}_1(x)}$，其中 $r=Z_1/Z_0$。

  • 这似乎是一个不可行的解决方案，因为它似乎需要我们估计数值 $Z_1/Z_0$。
  • 然而，可以从粗糙的 $r$ 开始估计，然后使用得到的桥式采样逐步迭代以改进估计 $\text{(Neal, 2005}$)。
  • 也就是说，我们会迭代地重新估计比率，并使用每次迭代更新 $r$ 的值。

• 链式重要采样 $\text{AIS}$和桥式采样各有优点。

  • 如果 $D_{\text{KL}}(p_0||p_1)$ 不太大（由于 $p_0$ 和$p_1$ 足够接近）的话，那么桥式采样能比$\text{AIS}$更高效地估计配分函数比率。
  • 然而，如果对于单个分布 $p_{\ast }$ 而言，两个分布相距太远难以桥接差距，那么$\text{AIS}$至少可以使用许多潜在中间分布来跨越 $p_0$ 和 $p_1$ 之间的差距。
  • $\text{Neal (2005)}$ 展示链式重要采样方法如何利用桥式采样的优点，桥接$\text{AIS}$中使用的中间分布，并且显著改进了整个配分函数的估计。

• 使用桥式采样，短链$\text{AIS}$和并行回火的组合， $\text{Desjardins et al. (2011) }$设计了一种在训练过程中追踪$\text{RBM}$配分函数的方法。该方法的基础是，在并行回火方法操作的每个温度下， $\text{RBM}$配分函数的独立估计。作者将相邻链（来自并行回火）的配分函数比率的桥式采样估计和跨越时间的$\text{AIS}$估计组合起来，提出一个在每次迭代学习时估计配分函数的（且方差较小的）方法。

• 本章中描述的工具提供了许多不同的方法，以解决难求解配分函数的问题，但是在训练和使用生成模型时，可能会存在一些其他问题。其中最重要的是我们接下来会遇到的难以推断的问题。

总结

通过对配分函数的估计，我们能够更深入地理解系统的微观结构与宏观性质之间的关系，为热力学和统计物理的研究开辟了新的道路。这些估计方法不仅提高了我们计算复杂系统热力学量的精度，还为我们探索未知物质的性质提供了有力的工具。随着相关理论的不断完善和技术的持续发展，相信未来我们将能够更准确地估计配分函数，揭示更多物质的奇妙性质和潜在应用。配分函数之估计，不仅是统计物理领域的一大挑战，更是推动科学技术进步的重要力量。

往期内容回顾

蒙特卡罗方法 -引言篇
深度学习中的结构化概率模型 - 结构化概率模型的深度学习方法篇