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

文章插图
【基于牛顿求根法，新算法实现并行训练和评估RNN，带来超10倍增速】其中 J_pf 是 f 在其第 p 个参数上的雅可比矩阵。根据上式，通过选择

文章插图
可让 δy^(i+1) 的一阶项为 0 。
这表明，根据上式选择矩阵 G_p，能以最快的速度收敛到解附近。这还表明，3 式和 5 式中的迭代相当于在巴拿赫空间（Banach space）中实现牛顿法，因此能提供二次收敛性。
3 式中的迭代过程涉及到评估函数 f、其雅可比矩阵和矩阵乘法，这些运算可以使用现代加速器（如 GPU 和 TPU）来并行化处理。如果能以并行方式求解线性方程，那么整个迭代过程都可利用并行计算。在深度学习背景中，将非线性微分方程视为定点迭代问题来求解还有另一个优势，即可以将前一步骤的解（如果能放入内存）用作下一训练步骤的起始猜测。如果起始猜测更好，则能减少寻找非线性微分方程的解所需的迭代步骤。
实际实现
为了将 3 式的 DEER 框架用于具体问题，需要遵循一些步骤。
第一步是将问题改写成 1 式，定义变量 y、线性算子 L [?] 和非线性函数 f (?) 。
第二步是实现研究者所说的位移器函数（shifter function）。这个位移器函数是以 y (r) 的整体离散值为输入，返回经过位移的 y 值的列表，即 y (r − s_p)，其中 p = {1, ..., P} 。这个位移器函数可能需要一些附加信息，比如起始或边界条件。这个位移器函数的输出将会是非线性函数的输入。
下一步（通常也是最难的一步）是根据矩阵列表 G_p 和在某些点离散的向量值 h 实现逆算子

文章插图
。这个逆算子可能也需要有关边界条件的信息。
只要能提供算法 1 中的需求，就可以将 DEER 框架应用于任意微分或差分方程。

文章插图
并行化常微分方程（ODE）
ODE 的形式通常是 dy/dt = f (y (t), x (t), θ)，其中初始条件 y (0) 是已给定的。上面的 ODE 形式如果用 1 式表示，则有 r = t、L = d/dt、P = 1 和 s_1 = 0 。这意味着 ODE 中的算子
相当于在给定初始条件 y (0) 时求解下面的线性方程。

文章插图
假设 G (t) 和 z (t) 是 t = t_i 和 t = t_{i+1} 之间的常量，分别为 G_i 和 z_i，则可以将 y_{i+1}=y_(t_i+1) 和 y_i = y (t_i) 之间的关系写成：

文章插图
其中 ?_i = t_{i+1} − t_i，I 是单位矩阵，exp (?) 是矩阵指数。9 式可以使用并行前缀扫猫算法进行评估。具体来说，首先可以为每个离散时间点 t_i 定义一对变量
，初始值 c_0=(I|y_0) 以及一个关联算子

文章插图
给定上面的初始值 c_0 和关联算子，可以并行方式运行关联扫描以获取上述算子的累积值。解 y_i 可从这个并行扫描算子的结果的第二个元素获取。
并行化 RNN
循环神经网络（RNN）可以看作是一种离散版的 ODE 。令索引 x 处的输入信号为 x_i，前一状态为 y_{i-1}，则当前状态可以写成 y_i = f (y_{i-1}, x_i , θ) 。
这个形式可以捕获常见的 RNN 单元，比如 LSTM 和 GRU 。而如果用 1 式来写这个形式，则有 r = i、L [y] = y、P = 1 和 s_1 = 1 。这意味着给定起始状态 y_0，可以通过求解下式来计算逆线性算子：