METAL

Sakana AI 公开无需反向传播的学习方法 PC-ALM

Sakana AI 于 9 月 14 日公开了 PC-ALM,这种方法不用反向传播,各层只与相邻层交换信号,就能训练 1,000 层的神经网络。在宽度为 32 的 MNIST 实验中,准确率与反向传播的差距约为 2 个百分点。

Sakana AI 公开无需反向传播的学习方法 PC-ALM

图片:@SakanaAILabs (X)(视频截图)

摘要

  • Sakana AI 于 9 月 14 日在博客和 X 上公开了 PC-ALM,仅靠层局部动力学训练出 1,000 层的残差 MLP。
  • 核心是在每一层加上一个累积预测误差的对偶神经元,让每一层都变成一个 PI 反馈控制器。
  • 在 32 层的 Fashion-MNIST 实验中,测试准确率分别为反向传播 78.66%、传统 PC 68.13%、PC-ALM 77.75%。
Augmented Lagrangian Predictive Coding: training 1000-layer networks without backpropagation

Sakana AI 于 9 月 14 日在 X 和自家研究博客上公开了无需反向传播的学习方法 PC-ALM(Augmented Lagrangian Predictive Coding,增广拉格朗日预测编码)。它不使用深度学习的基本训练算法反向传播,而是仅靠每一层只与相邻层交换信号的局部动力学,训练出 1,000 层深的神经网络。据该公司介绍,用宽度为 32 的残差 MLP 训练 MNIST 时,PC-ALM 的准确率在直到 1,000 层的整个区间内都与反向传播相差约 2 个百分点以内。论文于 5 月 29 日发布在 arXiv 上,代码以 JAX 实现的形式公开在 GitHub

反向传播按严格顺序运行前向传播、反向传播和权重更新三个阶段。靠前各层的神经元必须保持自己的激活值,等待误差信号到达。博客指出,大脑中没有已知的机制能在整个网络范围内强制执行这种时序,因此神经科学早已得出结论:大脑无法精确实现反向传播。大脑在不使用反向传播的情况下如何解决跨多层的信用分配(credit assignment)问题,至今仍是神经科学的未解难题之一。

预测编码(predictive coding,PC)是研究时间最长的替代方案。它源自亥姆霍兹的无意识知觉理论,Rao 和 Ballard 在 1999 年将其形式化为视觉皮层的模型。每一层自行预测输入信号,只把预测失败的部分,也就是预测误差,传给下一层。问题在于,这种方式在又深又窄的网络中信号会衰减。博客指出,从输出端进入的监督信号要经过每一层与邻层妥协的链条,远在到达靠前各层之前就已经变弱。博客列举的既有成功案例,是用每层 512 个神经元的 128 层残差 MLP 训练 MNIST 和 Fashion-MNIST。

PC-ALM 在此基础上给每一层加上一个对偶神经元(dual neuron),在数学上就是一个拉格朗日乘子。这个乘子在每一步累积该层的预测误差。当前误差成为比例项,累积误差成为积分项,每一层就像一个 PI 反馈控制器那样运作。其原理和汽车的定速巡航一样,它同时根据当前的速度差和一路累积的速度差来踩油门。出发点是 Yann LeCun 1988 年的一篇论文,其中观察到带约束条件的神经网络的拉格朗日乘子在平衡态下等于反向传播的梯度。Sakana AI 把这一构造与 PC 的能量函数合并到增广拉格朗日函数上,在其上运行层局部动力学。论文证明,在无偏置项的线性网络中,对偶神经元会精确收敛到反向传播的信号,非线性网络则通过实验加以验证。

实验数据作为复现基准记录在 GitHub 仓库的 README 中。METAL 核对过的这张表显示,用宽度 32、深度 32 的 ReLU 残差 MLP 在 Fashion-MNIST 上训练一个 epoch 时,测试准确率分别为反向传播 78.66%、传统 PC 68.13%、PC-ALM 77.75%。衡量权重梯度与反向传播梯度方向一致程度的余弦值分别为反向传播 1.000、PC 0.604、PC-ALM 0.909。Fashion-MNIST 的完整网格共运行 675 次。论文把宽度和深度各取 8、16、32、64、128 组合,再与恒等、tanh、ReLU 三种激活函数交叉比较,并指出当推理步数预算设为深度的两倍时,PC-ALM 在整个区间内都达到了与反向传播相同的性能。训练以批大小 64 进行一个 epoch,数据集为 MNIST 和 Fashion-MNIST 两个。1,000 层实验是博客新增的结果,跑了五个 epoch,推理步数为层数的两倍,即 2,000 步。该公司还写道,用 ResNet-18 训练 CIFAR-10 和 Tiny ImageNet 的实验中,PC-ALM 也优于 PC。

作者们称之为出人意料的,是信号扩散的形态。在 PC 中,监督信号像热量一样扩散并随深度衰减;而在 PC-ALM 中,它像波前一样从输出端快速向输入端推进,均匀地铺满整个网络。作者把这称为弹道式(ballistic)传播,以区别于扩散式传播。单个神经元在推理过程中会呈现阻尼振荡。论文用公式写明了条件:对偶更新步长 α 为 0 时精确退化为传统 PC,增大 α 会出现复特征值并产生振荡,增得过大则会失稳。

从社会学家的视角看,这项研究也是连接两个学术共同体的桥梁。作者表示,他们的出发点是一个观察:神经科学一侧的 NeuroAI 共同体和优化一侧的分布式优化共同体都在处理局部性,彼此之间却几乎没有对话。第二个观察是,PC 的能量式与分布式优化中常用的增广拉格朗日函数的增广项相似得可疑,正是这一点把两个领域连在了一起。代价用数字写得清清楚楚。论文指出,PC-ALM 的激活值内存是 PC 的两倍,而且每个小批量都要运行与深度成正比的推理步数。作为交换,与通过加宽网络来逼近反向传播的 PC 方式相比,作者计算得出,把与反向传播梯度的余弦值推到 0.9 所需的计算量少了一个数量级以上。该公司把这项研究的用武之地指向神经形态硬件:在动力学仿真比 GPU 更便宜的芯片上,各层只与邻层通信的学习方法有望带来节能的深度学习。实验仍停留在 MNIST、Fashion-MNIST 和残差 MLP 上,尚未应用于 Transformer 或语言模型;作者也在论文附录中写明,能量梯度仍然使用同一层权重转置矩阵的权重传输(weight transport)问题,这一方法并未解决。METAL 曾报道 Sakana AI 介绍英国皇家学会世界模型专刊一事,追问大脑如何学习的研究正在这家公司延续。

论文作者、Sakana AI 的 Jeffrey Seely 和 Julian Gould 在博客中写道:"据我们所知,这是首个被证明能成功训练多达 1000 层网络的层局部方法。"该公司在 X 帖子中表示:"归根结底,我们的动机是理解像大脑这样的分布式物理系统,如何仅凭局部耦合和局部动力学来计算信用信号。"这条帖子发布几小时内浏览量就超过 6 万,所附的 45 秒无声视频用图示展示了反向传播、PC 和 PC-ALM 之间的差别。作者把跨时间的信用分配、自监督损失,以及更大的网络和更难的任务列为下一步课题。

如今的深度学习完全依赖反向传播这一种方法。各层只与邻层对话的网络在 1,000 层深度上追到了与反向传播相差 2 个百分点以内,这是类脑学习方法能够走出玩具规模的第一个证据。下一个值得关注的数字,是这种方法在神经形态芯片上究竟能少用多少能量。

评论