面向可扩展稳定性：非线性RNN的并行化

Transformer和线性状态空间模型可在现代硬件上并行评估，但非线性RNN的评估似乎是固有的顺序问题。最近，Lim等人（2024年）开发了一种称为DEER（Deep Equilibrium Solver for Recurrent Networks）的方法，通过将状态表示为固定点问题的解来实现非线性RNN的并行评估。他们推导了求解固定点问题的并行Newton方法形式，并相对于顺序评估实现了显著的加速。然而，DEER的计算复杂度随状态大小呈三次方增长，算法可能 suffer from numerical instability（数值不稳定）。我们通过两个新方法来解决这些限制。为了降低计算复杂度，我们应用准Newton近似，证明其收敛性相当于Newton，占用更少内存且更快。为了稳定DEER，我们利用Levenberg-Marquardt算法与Kalman平滑之间的联系，这我们称为ELK。这种联系允许我们在保持性能的同时，使用高效并行Kalman平滑算法来稳定Newton方法。在多个实验中，我们展示了这些创新使非线性RNN能够在更大规模和更高稳定性下进行并行评估。