Divide-and-Conquer Predictive Coding: a structured Bayesian inference algorithm
分治预测编码：一种结构化贝叶斯推理算法https://arxiv.org/pdf/2408.05834
github：Pyro
摘要
意外刺激会在大脑中诱发“错误”或“惊讶”信号。预测编码理论通过提出大脑皮层在概率图模型中实现变分推断，有望用贝叶斯推断来解释这些观察结果。然而，当应用于机器学习任务时，这类算法在高维、结构化推断问题上的表现尚未与其他变分方法相提并论。为了解决这一问题，我们介绍了一种针对结构化生成模型的新型预测编码算法，我们称之为“分而治之预测编码”（DCPC）；它与其他预测编码公式不同，因为它尊重生成模型的相关性结构，并且可以证明其执行了模型参数的最大似然更新，同时又不牺牲生物合理性。从实证角度来看，DCPC在数值表现上优于其他算法，并且在一些以前未用预测编码解决的问题上提供了准确的推断。我们在Github上的Pyro中公开了DCPC的实现。
1 引言
近几十年来，认知科学、机器学习和理论神经科学领域在通过统计学习建模智能行为方面取得了诸多成功。这些领域各自采取了不同的方法：认知科学研究了每个任务和环境模型中的概率逆推断[Chater等人, 2006; Pouget等人, 2013; Lake等人, 2017]；机器学习采用了误差反向传播[Rumelhart等人, 1986; Lecun等人, 2015; Schmidhuber, 2015]；神经科学假设预测编码（PC）[Srinivasan等人, 1982; Rao和Ballard, 1999; Friston, 2005; Bastos等人, 2012; Spratling, 2017; Hutchinson和Barrett, 2019; Millidge等人, 2021]可以解释感知任务中的神经活动。这些方法共同之处在于都致力于“深层”模型，即任务处理由基本单元的组合而产生。
在机器学习中，基于PC的算法因其理论上可能为训练神经网络提供一种比反向传播更具生物合理性的替代方案而最近备受关注[Salvatori等人, 2023; Song等人, 2024]。然而，由于当前公式的局限性，PC在这些任务中的表现无法与反向传播相提并论。首先，用于梯度计算的预测编码通常将计算图中的每个节点建模为高斯分布，因此无法表达许多常见的生成模型。最近关于PC的工作通过允许用样本近似非高斯能量函数来解决这一问题[Pinchetti等人, 2022]。其次，后验的拉普拉斯近似仅推断每个潜在变量的最大后验（MAP）估计和高斯协方差，导致PC无法捕获多模态或相关分布。第三，这种对后验分布的宽松近似导致生成模型参数的更新不准确且方差较高。
在这项工作中，我们提出了一种新的算法——分而治之预测编码（DCPC），该算法使用蒙特卡洛样本群体来近似结构化目标分布。DCPC超越了高斯假设，将从结构化目标中采样的问题分解为对单个随机变量的局部坐标更新。这些局部更新通过以生物学上合理的预测误差为参数的未调整朗之万提议进行信息更新。将局部更新嵌套在分而治之的顺序蒙特卡洛方法[Lindsten等人, 2017; Kuntz等人, 2024]中，可以确保DCPC能够针对任何静态结构化图形模型，同时定理2提供了一种通过最大边际似然来学习模型参数的局部分解方法。DCPC还为神经科学中的典型皮层微回路[Bastos等人, 2012, 2020; Campagnola等人, 2022]假设提供了计算视角。实验表明，皮层微回路中的深层板层代表感官意象，而表层板层代表原始刺激信息[Bergmann等人, 2024]；在预测编码范式中的实验特别表明，深层代表“预测”，而浅层代表“预测误差”。这种回路可能为大脑提供了快速、可扩展、通用的贝叶斯推断能力。图1比较了DCPC的计算结构与之前预测编码模型的计算结构。以下部分详细介绍了这项工作的贡献：
• 第3节定义了分而治之预测编码算法，并展示了如何将其作为变分推断算法使用；
• 第4节探讨了大脑在何种假设下可能实现DCPC，并证明了两个有助于生物合理性的定理；
• 第5节通过与预测编码文献中最近的生成模型和推断算法进行直接比较，从实验上验证了DCPC。
第2节将回顾第3节算法的背景：预测编码旨在解决的问题，以及本文从中汲取灵感的一系列近期工作。
2 背景
通常，边际密度和后验密度没有闭合形式，因此学习和推断算法将联合分布视为潜在变量上的未归一化密度的闭合形式；其积分则给出用于近似的归一化常数：
通常，边际密度和后验密度没有闭合形式，因此学习和推断算法将联合分布视为潜在变量上的未归一化密度的闭合形式；其积分则给出用于近似的归一化常数：
这篇论文基于Kuntz等人[2023]的粒子梯度下降（PGD）算法，该算法的工作方式如下：在每次迭代 t 中，PGD以学习率 η 和独立的高斯噪声扩散粒子云穿过目标对数密度；然后通过平均粒子来估计对数似然的梯度，并更新参数 θ。更新规则如下：
上述方程针对的是整个图形模型的联合密度。当先验分布pθ(z)分解为许多独立的条件密度时，要实现高推断性能，通常也需要将推断网络或算法分解为条件形式[Webb等人, 2018]。如在投影梯度下降（PGD）和摊销推断[Dasgupta等人, 2020; Peters等人, 2024]中那样估计整个对数联合密度的梯度，也需要进行非局部的反向传播。为了仅使用局部计算为高维结构化模型提供一个通用的推断算法，第3节将应用方程3对联合密度中的单个随机变量进行采样，通过顺序蒙特卡洛方法组合坐标更新，并通过方程4局部估计模型参数的梯度。
3 分而治之预测编码
上一节提供了一个基于梯度更新和预测编码工作定义构建蒙特卡洛算法的数学工具箱。本节将结合这些工具来推广上述预测编码的概念，从而得出新颖的分而治之预测编码（DCPC）算法。给定一个因果图形模型，DCPC将使用K个样本的总体q(z)来近似后验分布，同时学习解释数据的参数θ。这将需要推导出局部坐标更新，并根据预测误差对其进行参数化。
由此产生的提议现在针对的是完全条件密度（方程5），同时满足了定义1对于纯粹局部提议计算的非正式要求，同时也将采样问题“分而治之”，转化为低维坐标更新。
4 生物合理性
文献中的不同作品考虑了不同的生物合理性标准。本文遵循非尖峰预测编码文献，认为如果一个算法在概率图模型中仅执行空间局部计算[Whittington和Bogacz，2017]，而不需要全局控制计算，则认为该算法在生物学上是合理的。然而，虽然在标准文献中，局部性要么是直接在目标函数中定义的[Rao和Ballard，1999]，要么是从联合密度的均场近似中得出的[Friston，2005]，但要证明DCPC参数的更新仅需要局部信息并不那么简单。为此，在这一部分中，我们首先正式展示了DCPC实现了潜变量z的分散推断（定理1），然后也通过局部信息更新参数θ（定理2）。
证明。见附录B中的推论4.1。
生物神经元经常通过突触来表示它们膜电位的变化[Mainen和Sejnowski，1995，Lundstrom等人，2008，Forkosh，2022]，有的甚至被测试并发现可以信号潜在信号的对数的时间导数[Adler和Alon，2018，Borba等人，2021]。理论家还提出了模型[Chavlis和Poirazi，2021，Moldwin等人，2021]，根据这些模型，单个神经元可以在内部计算梯度。简而言之，如果神经元可以代表概率密度，正如许多理论提议和实验表明它们可以的那样，那么它们很可能也可以计算DCPC中使用的预测误差。定理2将证明，鉴于上述“分解”，DCPC的模型学习只需要局部预测误差。
我们在上述实验中实际实现的DCPC（分而治之预测编码）利用了定理2，通过在生成模型的向前祖先采样过程中将样本从自动微分图中分离出来，以节省内存。
最后，DCPC通过重要性重采样操作，从局部坐标更新过渡到联合目标密度，这要求实现过程同步传输新提议粒子群体的数值密度或对数密度。虽然与皮质振荡的相位锁定可能使这在生物学上成为可能，但重采样随后需要归一化权重。幸运的是，除法归一化在整个大脑中普遍存在[Carandini和Heeger, 2012]，以及实现softmax函数（例如，用于归一化和重采样重要性权重）的“胜者全得”电路类型在皮质柱的浅层和深层之间的交叉谈话中也普遍存在[Liu, 1999; Douglas和Martin, 2004]。
5 实验
分而治之预测编码并不是第一个将采样纳入推断过程的预测编码算法，当然也不是第一个用于结构化图形模型的变分推断算法。因此，本节将DCPC的性能与预测编码文献中的模型以及标准深度生成模型进行了比较。除非另有说明，否则每个实验都保持生成模型、数据集和超参数不变。
我们已经将DCPC实现为深度概率编程语言Pyro中的变分提议或“指导”程序[Bingham et al., 2019]；这样做使我们能够在涉及神经网络的图形模型中高效地计算自由能量和预测误差。由于下面的实验涉及从未知分布的数据集D ∼ p(D)中对观察x ∼ B进行小批量子采样，我们用变分序贯蒙特卡洛梯度估计器的子采样形式（见Welling和Teh [2011]的推导）替换了方程9 [Naesseth et al., 2018]。
在所有实验中，我们使用Adam优化器[Kingma和Ba, 2014]来优化自由能，并确保在每次Pyro的sample()操作后调用detach()函数，以实现定理2和方程10中描述的纯粹局部梯度计算。下面的第一个实验考虑了三个简单数据集上的层次高斯模型。该模型由观测值上方的两个潜在编码组成。
Oliviers等人[2024]提出的带有预测编码的深度潜在高斯模型将预测编码与神经采样假设结合在一个模型中：蒙特卡洛预测编码（MCPC）。他们的推断算法在功能上反向传播了对数似然度的得分函数，并在此过程中应用Langevin提议从后验联合密度中采样潜在变量。他们使用深度潜在高斯模型[Rezende等人, 2014]（DLGM）在MNIST数据集上评估了MCPC的性能。他们的模型的条件密度由非线性操作后跟线性变换组成，以参数化每个高斯条件的均值，并学习协方差。图2显示，DLGM结构已经要求DCPC尊重层次依赖性。
我们在基本重建和生成任务上测试了DCPC的性能，通过使用它来训练这个确切的生成模型，只将似然从离散伯努利改为连续伯努利[Loaiza-Ganem和Cunningham，2019]。训练后我们评估了使用离散伯努利似然的情况。表2显示，无论是在惊讶度（负对数证据，使用离散伯努利似然）还是均方重建误差方面，DCPC都具有更好的平均性能和更低的性能标准差，后者低一个数量级。所有实验都使用了学习率 η = 0.1 和 K = 4 粒子。
图3展示了将这个实验扩展到EMNIST [Cohen等人，2017]、Fashion MNIST [Xiao等人，2017]以及原始的MNIST数据集，上面一行是真实图像，下面一行是从DCPC推断出的潜在代码的重建图像。真实图像来自于每个数据集10%的验证集分割，在该验证集上，DCPC只推断粒子。
上述数据集通常不会对新的推理算法构成挑战。因此，接下来的实验将尝试学习彩色图像的表示，就像在广泛使用的变分自编码器[Kingma和Welling, 2013]框架中那样，但不同的是，我们将不使用编码器网络或摊销推理。
基于表征学习的图像生成 Zahid等人[2024]最近也设计并评估了Langevin预测编码（LPC），它与MCPC和DCPC都有所不同，涉及到使用表示学习的图像生成。
虽然MCPC通过层级模型向上传递预测误差，但LPC则计算生成模型中所有潜在变量的对数联合梯度作为其预测误差。这意味着，出于生物合理性的考虑以及他们想要实现预测编码推理摊销的目标，他们只能使用从Higgins等人[2017]改编的单级解码器。我们使用了他们从Cheng等人[2020]、Ho等人[2020]中提取的离散高斯似然度进行评估，并按照Rybkin等人[2021]的方法学习方差。
我们采用Zahid等人[2024]中提到的Fréchet Inception Distance（FID）[Seitzer, 2020]来比较DCPC和LPC。在比较中，我们保持先验、神经网络架构、θ的学习率以及用于训练参数θ和潜在变量z的梯度评估次数不变。Zahid等人[2024]评估了多种场景，并报告说，如果以训练轮次来计算，他们的训练可以快速收敛，但他们在推理步骤中累积了θ的梯度。我们与他们报告的在先验初始化的潜在变量上应用300个推理步骤后的15轮训练结果进行了比较，这相当于每批对θ进行了15×300=4500次梯度步骤，并复制了他们64的批大小。由于算法1仅在外部循环中更新θ，我们设置S=30，并以η=0.001和K=16运行DCPC进行150轮训练。表3显示，在同类生成质量上，DCPC优于LPC，尽管仅就推理质量而言，它还没有达到与其他模型架构相媲美的程度。
图4显示了验证集（左）中的重建图像和后验预测生成模型（右）中的样本。重建图像存在模糊现象，这在变分自编码器中经常发生，但DCPC训练允许网络捕捉到背景颜色、发色、面部朝向以及其他视觉属性。图4a显示了验证集上的重建图像，而图4b显示了预测分布中的样本。
Kuntz等人[2023]也报告了在32×32较低分辨率下，就FID分数而言在CelebA数据集上进行的一项实验。由于他们提供了源代码和精确的数学描述，我们能够与PGD进行精确的一对一比较。表3中评估DCPC的一行采用了PGD在32×32分辨率下的示例神经网络架构（具有相似的粒子数和学习率），这证明了DCPC在FID方面有了显著改善，同时降低了FID方差。
必要的计算资源：最初的深度生成模型（DLGM）在类似MNIST数据集上的实验是在一台配备128GB RAM和8GB显存的NVIDIA Quadro P4000的桌面工作站上进行的。在CelebA数据集上的实验是在一台NVIDIA DGX上进行的，该DGX配备了八个（8）NVIDIA A100，每个A100有80GB的显存。后者的计算基础设施也用于在结构化时间序列建模上对几个不同数据集进行未发表的实验。
6 相关工作
Pinchetti等人[2022]首次将预测编码扩展到高斯生成模型之外，将得到的算法应用于通过变分推断训练变分自编码器，以及通过最大似然训练变换器架构。DCPC则将预测编码扩展到目标任意概率图模型，遵循Salvatori等人[2022]扩展到任意确定性计算图。DCPC在快速交替更新随机变量和模型参数方面，遵循了Salvatori等人[2024]的增量预测编码，给出了一种增量EM算法[Neal和Hinton，1998]。最后，Zahid等人[2024]和Olivier等人[2024]也认识到了预测编码的预测误差与Langevin动力学中用于连续随机变量的分数函数之间的类比。
在神经生物学上可行的回路如何实现概率推断方面，有大量工作。Shi和Griffiths[2009]的经典工作提供了一种通过重要性抽样实现层次推断的生物学上可行的实现；DCPC以重要性抽样为基础，通过预测误差参数化建议分布。Fang等人[2022]最近研究了基于Langevin动力学的抽样推断的神经学上可行算法，尽管仅针对稀疏编码的高斯生成模型。Golkar等人[2022]为了生物学上的可行性，对高斯生成模型施加了白化约束。最后，Dong和Wu[2023]和Zahid等人[2024]都提出了在生物学上可行的抽样算法中使用动量来减少梯度噪声的机制；前者有趣地将动量项类比为神经元适应。总之，其他工作已经在图像生成任务中实现了预测编码模型，一个著名的例子是神经生成编码框架Ororbia和Kifer[2022]，Ororbia和Mali[2022]。
7 结论
本文提出了分而治之的预测编码（DCPC），这是一种通过重要性抽样有效且可扩展地近似吉布斯采样器的算法；DCPC使用局部预测误差参数化模型的完整条件密度的高效建议。第4节展示了蒙特卡洛抽样如何实现一种“预期配置”[Song等人，2024]，首先从联合后验密度中推断样本（定理1），然后更新生成模型而不进行全局反向传播（定理2）。第5节的实验表明，DCPC在理论上考虑的简单生成模型上，与计算神经科学中的蒙特卡洛预测编码相比，具有优越的性能；DCPC还在完全局部计算的约束下，优于Kuntz等人[2023]的粒子梯度下降算法。DCPC的Langevin建议允许与哈密顿蒙特卡洛[Brubaker等人，2012]中应用的约束样本空间相同的扩展；我们的Pyro实现包括了通过Pyro对HMC的现有支持来扩展这一功能。
DCPC提供了多种前进的方向。特别是，本文采用了朴素的Langevin建议，而Dong和Wu[2023]，Zahid等人[2024]应用了基于动量的预处理来利用目标的几何结构。Yin和Ao[2006]证明了这种一般类型的梯度流也可以通过破坏Metropolis-Hastings算法所需的详细平衡条件来提供更有效的采样器，激励选择SMC而不是MCMC来纠正建议偏差。附录C推导了将DCPC扩展到离散随机变量的数学背景。未来的工作可以遵循Marino等人[2018]，Taniguchi等人[2022]，使用神经网络从粒子和预测误差迭代映射到建议参数。
7.1 局限性
DCPC的主要局限性是其训练时间较长，对学习率的敏感性比最先进的端到端训练的摊销变分推断更大。
这种局限性在受神经科学启发的学习算法文献中经常出现，以及在没有参数形式的基于粒子的算法文献中也很常见。
这项工作没有特定于DCPC的单一伦理问题，而是伴随着在生物学上可行的学习与贝叶斯推断进步的更广泛影响和责任。
进一步的实验和结果
交替图像生成/表示学习 如第2节所述，本文基于粒子梯度下降（PGD）算法；Kuntz等人[2023]通过在CelebA上训练一个生成器网络来展示算法的性能。他们的网络采用了固定标准差为0.01的高斯似然，并且在100个周期内对确切的10,000个抽样数据点评估了对数联合目标。然后，论文在图像修复任务上评估了均方误差，并在数据图像上评估了Frechet Inception Distance。
当应用于上述目标密度时，DCPC相当于带有重采样步骤的PGD。表4显示了使用DCPC训练和评估上述相同模型的结果。由于PGD在CelebA的10,000图像子样本上训练了100个周期，批量大小为128，我们使用整个数据集进行了100个周期的训练，批量大小为128。
我们怀疑所提供的用于计算对数联合的代码在平均粒子或批量项目上的方式与我们评估DCPC的方式不同（例如，我们在没有进一步除以任何形状维度的情况下调用mean()），这解释了对数联合之间明显的数量级差异。
应审稿人的要求，我们在正文中用图1替换了图5，展示了如何将DCPC映射到层状微电路结构上。
b .重要抽样和梯度估计证明
Titsias在2023年提出了针对目标密度的Fisher信息进行自适应的预处理Langevin动力学的最优估计器。定义2给出了该Fisher信息的最基本估计器，它的定义是基于我们用作预测误差的分数函数。
以下两个命题来自于Wu等人[2020]、Stites等人[2021]和Zimmermann等人[2021]的先前工作。想要了解基础理论的读者可以参考Naesseth等人[2015]或Chopin和Papaspiliopoulos[2020]的研究。
C 扩展到离散样本空间
最近关于离散分布的高效抽样工作集中在近似密度比率，例如方程18中的比率，通过由误差向量参数化的级数展开来实现。当底层的离散密度由可微能量函数指数化组成时，如Grathwohl等人[2021]所述，这些误差向量采取了梯度的形式，有限级数展开一直是泰勒级数。当不是这种情况时，Xiang等人[2023]展示了它们采取了有限差分和牛顿级数的形式。
https://arxiv.org/pdf/2408.05834
