🤖 machine learning

A projection-based framework for gradient-free and parallel learning

本文介绍了 PJAX，这是一个基于 JAX 的框架，它利用迭代投影算子将神经网络训练重新表述为一个可并行化且无需梯度的可行性问题，为传统的基于梯度的优化方法提供了极具吸引力的替代方案，其在处理不可微操作和实现大规模并行化方面具有优势。

原作者： Andreas Bergmeister, Manish Krishan Lal, Stefanie Jegelka, Suvrit Sra

发布于 2026-04-30

📖 1 分钟阅读☕ 轻松阅读

原作者： Andreas Bergmeister, Manish Krishan Lal, Stefanie Jegelka, Suvrit Sra

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象你正在尝试解决一个巨大而复杂的拼图，比如一个巨大的三维拼图或魔方，但你并不知道最终的图案是什么。

旧方法（基于梯度的学习）
目前，大多数人工智能模型使用一种称为“反向传播”的方法进行学习。这就像一名徒步者试图在雾蒙蒙的山谷中找到谷底（最佳解决方案）。徒步者感受脚下的坡度（梯度），然后向山下迈进一步。他们就这样一步一步地继续，直到到达一个低点。

问题所在： 有时，徒步者会困在一个并非真正谷底的小凹陷处（局部极小值）。有时路径过于陡峭或平坦，导致徒步者迷路或移动过慢。此外，为了知道哪边是“下”，徒步者必须从谷底向山顶发送信号，这既缓慢，又需要一条非常特定且对称的路径。

新方法（基于投影的学习）
本文的作者提出了一种完全不同的策略。他们不试图寻找山谷的底部，而是将训练视为一个可行性问题。

想象你有一个房间，里面布满了墙壁，每面墙都有一条特定的规则。

墙 A 说：“红色积木必须紧挨着蓝色积木。”
墙 B 说：“绿色积木必须在红色积木的上方。”
墙 C 说：“总重量必须等于 50 公斤。”

你的目标不是滑下山坡，而是找到一种单一的积木排列方式，使得每一面墙的规则同时得到满足。

它是如何工作的：“投影”隐喻
作者将他们的方法称为“基于投影”。以下是他们的做法：

分解问题： 他们将巨大的拼图（神经网络）分解为微小的、简单的部分，称为“基本函数”（如简单的数学运算：数字相加、相乘，或判断一个数字是否为正数）。
局部修正： 他们不看整个拼图，而是只看一面墙（一条规则）。如果积木不符合该墙的规则，他们就将积木“投影”到墙上。想象用光照在积木上；它们投射在墙上的影子就是针对该特定规则的“正确”位置。
并行优势： 这是神奇之处。因为每面墙只关心其直接的邻居，你可以同时修正墙 A、墙 B 和墙 C。你不必等墙 A 修好后再开始修墙 B。这就像有一支 100 人的团队同时修复房子的不同部分，而不是一个人先修屋顶，再修厨房，然后修浴室，逐个进行。
重复： 他们一遍又一遍地这样做。每次，他们都会微调积木以更好地符合局部规则。最终，积木会稳定在一个位置，同时满足所有规则。这就是你训练好的 AI。

为什么这很酷（根据论文）

无需“坡度”： 你不需要计算“坡度”（梯度）。这意味着你可以使用那些“凹凸不平”或断裂的规则（不可微分），比如一个要么开要么关的开关。旧方法处理这些很吃力；而新方法能轻松应对。
生物学合理性： 在大脑中，神经元不会从思考的末端向始端发送全局的“误差信号”。它们只是根据直接邻居的行为进行调整。这种新方法模仿了这种局部的、邻居对邻居的调整。
速度： 因为所有人并行工作，它在现代计算机芯片（GPU/TPU）上可以快得多，这些芯片专为同时执行多项任务而设计。

权衡：“内存”成本
论文承认存在一个陷阱。为了做到这一点，计算机必须在每一步记住拼图中每一个“边”的位置。

类比： 在旧方法中，你只需记住徒步者当前的位置。而在新方法中，你必须记住房间里每一块积木的位置，以及它们之间的每一个连接，对于你团队中的每一个人都是如此。
结果： 这会消耗更多的计算机内存（RAM）。作者不得不缩小一些测试模型，以便将它们放入计算机内存中，而旧方法则能更轻松地处理更大的模型。

结果
作者构建了一个名为PJAX（Projection JAX）的软件工具来测试这一方法。他们在不同类型的拼图上进行了尝试：

简单模式（MLPs）
图像识别（CNNs）
语言预测（RNNs）

他们发现，虽然“旧方法”（使用 Adam 或 SGD 优化器）在许多情况下仍然是原始速度和最终准确率的冠军，但这种新的“投影”方法表现得出乎意料地好。它是一种可行的替代方案，具有以下特点：

无需梯度即可学习。
能够处理令其他方法困惑的“凹凸不平”的规则。
在并行硬件上学习非常高效，特别是在语言建模等任务中，旧方法在处理“梯度消失”（忘记句子的开头）方面存在困难。

总结
论文指出：“停止试图滑下山坡去寻找答案。相反，将问题视为一组局部规则。局部且同时地修正每条规则，最终，整个系统将各就各位。”这是一种训练 AI 的新方法，它更具并行性，对不同类型的数学运算更灵活，但目前需要更多的内存。

以下是论文《基于投影的无梯度并行学习框架》的详细技术总结。

1. 问题陈述

本文针对深度学习训练主导范式的根本局限性：即通过反向传播进行的基于梯度的优化。尽管该方法非常有效，但面临几个关键挑战：

可微性要求：它要求所有网络组件和损失函数必须是（次）可微的，从而排除了不可微操作（例如量化、离散逻辑）。
顺序瓶颈：反向传播本质上是顺序执行的，限制了并行性并增加了更新延迟，尤其是在深层架构中。
生物学不合理性：对称权重传输机制（将误差梯度通过用于前向传播的完全相同的权重反向发送）被认为在生物学上是不现实的。
优化病理：梯度方法通常难以应对梯度消失/爆炸、局部极小值和鞍点等问题。

作者提出了一种范式转变：将神经网络训练重新表述为可行性问题，而非损失最小化问题。目标不再是导航损失景观，而是寻找网络参数和中间状态，使其同时满足从网络结构和输入 - 输出映射导出的一组局部约束。

2. 方法论

核心方法论依赖于应用于细粒度计算图的迭代投影算法。

A. 可行性表述

训练被定义为寻找一个状态向量 $z$ （表示计算图中所有边的值），该向量位于多个约束集 $C_v$ 的交集中：
$\text{寻找 } z \in \bigcap_{v \in V} C_v$
其中 $V$ 代表计算图中的节点。约束针对每种节点类型局部定义：

常数输入节点：强制 outgoing 边携带输入值。
参数节点：强制来自同一参数的所有 outgoing 边达成共识（确保单个权重值在所有地方被使用）。
隐藏函数节点：强制边的输出等于将基本函数（例如 ReLU、点积）应用于其输入的结果。
目标/损失节点：强制网络输出相对于目标最小化损失（或满足边界约束）。

B. 投影算子

解决方案使用迭代算法找到，例如交替投影 (AP)、循环投影 (CP) 和 Douglas-Rachford (DR)。

局部投影：对于每个节点，算法将当前状态投影到由该节点定义的约束集上。
基本函数：本文推导了针对基本操作图（例如点积、和的 ReLU、最大池化、量化）进行投影的闭式或高效数值解。
不可微性：由于投影依赖于距离最小化而非导数，该框架自然地处理不可微函数（例如量化、阶跃函数）。

C. 通过二分划实现并行化

一个关键的理论贡献是定理 2，它利用了计算图的二分性质。

图节点可以划分为两个不相交的集合 $A$ 和 $B$ ，使得同一集合内的任意两个节点都不相邻。
集合 $A$ 中所有节点的投影可以独立且并行地计算，因为它们操作的是不相交的边变量集。集合 $B$ 同理。
这使得整个更新步骤可以在网络上完全并行化，消除了反向传播的顺序依赖性。

D. PJAX 框架

作者引入了 PJAX (Projection JAX)，这是一个构建在 JAX 之上的软件框架。

组合性：类似于 autodiff 组合导数，PJAX 组合投影算子。
API：提供类似 NumPy/JAX 的 API 用于定义模型。
硬件：继承 JAX 的 JIT 编译以及对 GPU/TPU 的支持。
自动化：自动构建计算图并协调迭代投影步骤。

3. 主要贡献

理论重构：利用边变量和局部约束，将神经网络训练严谨地表述为大规模可行性问题，证明了可行解对应于经验风险最小化器。
可并行算法：证明了通过利用计算图的双分结构，投影更新可以完全并行化，这提供了相对于顺序反向传播的独特优势。
投影算子：针对关键基本函数（点积、ReLU、最大值、量化）和损失函数（交叉熵、边界损失）的投影算子进行了数学推导。
PJAX 框架：首个通用的、无梯度的训练框架，支持任意架构（MLP、CNN、RNN）和硬件加速。
实证验证：通过广泛的实验证明了该方法在多样任务和架构中的可行性。

4. 实验结果

作者在标准基准测试（MNIST、CIFAR-10、HIGGS、Shakespeare）上使用 MLP、CNN 和 RNN 评估了 PJAX，并与 SGD、Adam、反馈对齐 (FA) 和前向 - 前向 (FF) 进行了比较。

性能：
- MLP：投影方法（尤其是 Douglas-Rachford）达到了与 SGD 相当的准确率，通常具有显著更快的单步时间（在浅层网络中每步快达 10 倍）。
- RNN：投影方法通过避免随时间反向传播 (BPTT) 显示出结构优势，规避了梯度消失/爆炸问题，并在步数上比 SGD 收敛得更快。
- CNN：与基于梯度方法的性能差距较大，且由于共享参数带来的高内存开销（见下文），单步时间较慢。
深度与跳跃连接：没有跳跃连接的深层网络在基于投影的训练中难以进行信息传播。然而，引入跳跃连接（缩短从早期层到输出的路径）显著提高了性能，验证了在此范式中调整架构的必要性。
效率：
- 每步时间：由于大规模并行性，投影方法 (AP/CP) 的每步时间通常低于 SGD，尽管达到收敛的总挂钟时间有所不同。
- 内存：该方法需要为图中的每条边存储状态变量。这导致 $O(|E|)$ 的内存复杂度，高于反向传播的 $O(|V|)$ ，特别是在具有大量参数共享的架构（CNN、RNN）中。

5. 意义与未来方向

生物学合理性：该方法通过使用局部更新并避免全局误差反向传播，更符合生物学习原则。
处理不可微性：它为训练包含离散组件的网络（例如量化神经网络、逻辑门）打开了大门，而无需梯度近似。
可扩展性：更新的并行性质使其非常适合现代硬件（GPU/TPU），有可能克服非常深或非常宽网络中反向传播的顺序瓶颈。
局限性：主要局限性是内存使用量，它随批量大小、序列长度和参数共享而扩展。
未来工作：作者建议探索自适应阻尼、混合方案（混合投影与梯度）、节省内存的技术（低秩边状态）以及定制架构以减少参数 - 数据交互。

总之，本文提出了一个基于梯度训练的有吸引力的替代方案，证明了通过并行投影寻求可行性是一种可行、通用且高效的深度学习范式，特别适用于不可微和循环架构。