不依赖梯度的并行化神经网络训练新框架
A projection-based framework for gradient-free and parallel learning
- 用投影操作替代梯度下降，通过满足局部约束来训练网络
- 在多个标准数据集上成功训练MLP/CNN/RNN，支持非可微操作
- 适合追求高并行性、需处理不可导模块的研究者
我们提出一种基于投影的神经网络训练方法，突破传统梯度下降范式。该数学优化框架将训练问题重构为大规模可行性问题：寻找满足由基本运算导出的局部约束的网络参数与状态。训练过程通过迭代投影到这些约束实现，每个投影操作均为局部计算，可并行分布在全网。我们构建了PJAX——一个基于JAX的软件框架，可组合基本运算的投影算子，自动推导可行性问题的求解算子（类比自动微分）。该框架天然支持GPU/TPU加速，提供类似NumPy的接口且可扩展。我们在多个标准基准上使用PJAX训练了MLP、CNN和RNN等多样化架构，验证了其功能性和通用性。结果表明，该方法是梯度训练的有力替代方案，在并行性与处理非可微操作方面具有显著优势。
原文摘要 · Abstract (English)
We present a feasibility-seeking approach to neural network training. This mathematical optimization framework is distinct from conventional gradient-based loss minimization and uses projection operators and iterative projection algorithms. We reformulate training as a large-scale feasibility problem: finding network parameters and states that satisfy local constraints derived from its elementary operations. Training then involves projecting onto these constraints, a local operation that can be parallelized across the network. We introduce PJAX, a JAX-based software framework that enables this paradigm. PJAX composes projection operators for elementary operations, automatically deriving the solution operators for the feasibility problems (akin to autodiff for derivatives). It inherently supports GPU/TPU acceleration, provides a familiar NumPy-like API, and is extensible. We train diverse architectures (MLPs, CNNs, RNNs) on standard benchmarks using PJAX, demonstrating its functionality and generality. Our results show that this approach is a compelling alternative to gradient-based training, with clear advantages in parallelism and the ability to handle non-differentiable operations.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。