基于牛顿求根法，新算法实现并行训练和评估RNN，带来超10倍增速 _算法

过去十年来，深度学习领域发展迅速，其一大主要推动力便是并行化。通过 GPU 和 TPU 等专用硬件加速器，深度学习中广泛使用的矩阵乘法可以得到快速评估，从而可以快速执行试错型的深度学习研究。
尽管并行化已经在深度学习研究中得到了广泛的使用，但循环神经网络（RNN）和神经常微分方程（NeuralODE）等序列模型却尚未能完全受益于此，因为它们本身需要对序列长度执行序列式的评估。
序列评估已经变成了训练序列式深度学习模型的瓶颈。这一瓶颈可能会使人们关注的研究方向偏离序列模型。

文章插图
举个例子，注意力机制和 transformer 在近些年中超过 RNN 成为了语言建模的主导技术，部分原因就是它们能以并行的方式训练。连续归一化流（CNF）过去常使用的模型是 NeuralODE，现在却转向了训练过程不涉及到模拟 ODE 的新方向。
近期有一些尝试复兴序列 RNN 的研究工作，但它们的重心都是线性循环层 —— 可使用前缀扫描（prefix scan）来进行并行化地评估，非线性循环层在其序列长度上依然无法并行化。
近日，英国 machine Discovery 公司和牛津大学的一篇论文提出了一种新算法，可将 RNN 和 NeuralODE 等非线性序列模型的评估和训练工作并行化，并且他们宣称这一算法还不会在「合理的数值精度」内改变模型的输出。

文章插图
论文地址：https://arxiv.org/pdf/2309.12252.pdf
那么他们是怎么做到这一点的呢？
据介绍，他们引入了一种用于求解非线性微分方程的通用框架，其做法是将这些方程重新表述为二次收敛的定点迭代问题，这相当于牛顿求根法。定点迭代涉及到可并行运算和一个可并行地评估的逆线性算子，即使是对于 RNN 和 ODE 这样的序列模型也可以。
由于是二次收敛，所以定点迭代的数量可以相当小，尤其是当初始起点接近收敛的解时。在训练序列模型方面，这是一个相当吸引人的功能。由于模型参数通常是渐进式更新的，所以之前训练步骤的结果可以被用作初始起点。
最重要的是，研究者表示，新提出的算法无需序列模型具备某种特定结构，这样一来，用户不必改变模型的架构也能收获并行化的好处。
DEER 框架：将非线性微分方程视为定点迭代
DEER 框架具有二次收敛性，并且与牛顿法存在关联。这一框架可以应用于一维微分方程（即 ODE），也可用于更高维的微分方程（即偏微分方程 / PDE）。该框架还可以应用于离散差分方程以达到相同的收敛速度，这一特性可以应用于 RNN 。
使用该框架，用户可以设计一种用于评估 RNN 和 ODE 的并行算法，并且不会对结果产生明显的影响。
DEER 框架
令我们感兴趣的输出信号为 y (r)，其由 n 个在 d 维空间的信号构成，其坐标表示为 r 。输出信号 y (r) 可能依赖于输入信号 x (r)，其关系是某个非线性的延迟微分方程（DE）：

文章插图
其中 L [?] 是 DE 的线性算子，f 是非线性函数，其依赖于 P 个不同位置的 y 值、外部输入 x 和参数 θ 的。这是一个通用形式，足以表示各种连续微分方程，比如 ODE（当 L [?] = d/dt 且 r = t）、偏微分方程（PDE）、甚至用于 RNN 的离散差分方程。
现在，在左侧和右侧添加一项
，其中 Gp (r) 是一个依赖于位置 r 的 n×n 矩阵。G_p 的值会在后面决定。现在 1 式就变成了：