Sakana AI 用逐层本地方法训练千层网络

Sakana AI 的研究者提出 Augmented Lagrangian Predictive Coding（PC-ALM），这是预测编码（predictive coding，PC）的一个变体：每次权重更新都保持在本层内部，同时恢复出与反向传播（backpropagation）对齐的信用分配信号。研究团队称，在 MNIST 上训练最深至 1000 层的残差 MLP，准确率与反向传播的差距约为 2 个百分点。

方法

反向传播是一种全局算法：前向传播、反向传播、权重更新依次锁定，后一步必须等前一步完成。目前尚未发现大脑中存在这种全网范围的相位锁定机制，因此本地学习方案一直受到研究关注。PC 把每个隐藏激活当作优化变量，惩罚它与来自下层的预测之间的偏差。问题在于监督信号只从输出端进入，在又深又窄的网络中，信号传到输入端之前就已衰减。

PC-ALM 从训练的约束优化视角出发，保留 PC 的惩罚项，同时为每层约束附加一个拉格朗日乘子。推理过程交替执行两个本地步骤：对激活做一次原始梯度步，再做一次累积本层预测误差的对偶步。团队把这套机制解读为每层一个 PI 控制器，预测误差是比例项，乘子是积分项。把乘子设为 0，方法就完全退回 PC。

实验结果

在 Fashion-MNIST 和 MNIST 上，团队对宽度与深度在 8 到 128 之间的残差 MLP 做了网格扫描。在 T = 2L 的推理预算下，PC-ALM 在所有宽度、深度和激活函数（identity、tanh、ReLU）上都追平了反向传播，而 PC 在深而窄的配置中明显下滑。代码库的参考配置（宽度 32、深度 32、ReLU、Fashion-MNIST）测试准确率为：反向传播 78.66%，PC 68.13%，PC-ALM 77.75%；梯度与反向传播的余弦相似度从 0.604 升到 0.909。

团队表示，在试过的所有基准上 PC-ALM 都优于 PC，包括 CIFAR-10 和 Tiny ImageNet 上的 ResNet-18。PC-ALM 是训练方法而非模型，目前只在小规模图像基准上验证过。MIT 许可的 JAX 参考实现可在 CPU 上运行，并能复现论文中的宽度与深度网格结果。详情见 MarkTechPost 的报道。