论文

LBI:通过潜在有界接口的并行扫描反向传播

LBI: Parallel Scan Backpropagation via Latent Bounded Interfaces

AI 基础设施分布式训练基础设施

摘要

反向传播本质上是跨深度的顺序,创建了一个 $O(K)$ 深度的依赖链,成为并行训练的瓶颈。虽然并行扫描公式理论上可以将此深度减少到 $O(\log K)$,但由于在整个隐藏状态上组成全秩 $d\times d$ 雅可比行列式的成本为 $O(d^3)$,因此对于现代架构来说,它们在计算上是令人望而却步的。我们引入潜在有界接口(LBI),这是一种算法公式,通过将区域间通信限制为低维潜在接口 $ m_k \in \mathbb{R}^{r}$,其中 $r \ll d$,使基于扫描的反向传播变得易于处理。这将伴随递归减少为 $r \times r$ 雅克比矩阵上的后缀扫描,将每次组合成本从 $O(d^3)$ 削减到 $O(r^3)$,同时在有界界面模型下保留精确的梯度。我们证明,LBI 在 47--61M 块参数下保持了四种架构(Mamba-2、Mamba-3、Transformer 和 Mamba--Transformer 混合)的模型质量。维度 $r=16$ 的接口足以将训练质量保持在密集基线的 0.16--0.35 交叉熵内。由此产生的框架为区域并行训练提供了算法基础,将跨设备反向通信减少到对 $K$ 固定大小矩阵的单次扫描,对于我们的实验配置来说,该矩阵大约为 56 KB。