Gespeichert in:
Bibliographische Detailangaben
Hauptverfasser: Gonzalez, Xavier, Warrington, Andrew, Smith, Jimmy T. H., Linderman, Scott W.
Format: Preprint
Veröffentlicht: 2024
Schlagworte:
Online-Zugang:https://arxiv.org/abs/2407.19115
Tags: Tag hinzufügen
Keine Tags, Fügen Sie den ersten Tag hinzu!
_version_ 1866929677055033344
author Gonzalez, Xavier
Warrington, Andrew
Smith, Jimmy T. H.
Linderman, Scott W.
author_facet Gonzalez, Xavier
Warrington, Andrew
Smith, Jimmy T. H.
Linderman, Scott W.
contents Transformers and linear state space models can be evaluated in parallel on modern hardware, but evaluating nonlinear RNNs appears to be an inherently sequential problem. Recently, however, Lim et al. '24 developed an approach called DEER, which evaluates nonlinear RNNs in parallel by posing the states as the solution to a fixed-point problem. They derived a parallel form of Newton's method to solve the fixed-point problem and achieved significant speedups over sequential evaluation. However, the computational complexity of DEER is cubic in the state size, and the algorithm can suffer from numerical instability. We address these limitations with two novel contributions. To reduce the computational complexity, we apply quasi-Newton approximations and show they converge comparably to Newton, use less memory, and are faster. To stabilize DEER, we leverage a connection between the Levenberg-Marquardt algorithm and Kalman smoothing, which we call ELK. This connection allows us to stabilize Newton's method while using efficient parallelized Kalman smoothing algorithms to retain performance. Through several experiments, we show that these innovations allow for parallel evaluation of nonlinear RNNs at larger scales and with greater stability.
format Preprint
id arxiv_https___arxiv_org_abs_2407_19115
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle Towards Scalable and Stable Parallelization of Nonlinear RNNs
Gonzalez, Xavier
Warrington, Andrew
Smith, Jimmy T. H.
Linderman, Scott W.
Machine Learning
I.2.6
Transformers and linear state space models can be evaluated in parallel on modern hardware, but evaluating nonlinear RNNs appears to be an inherently sequential problem. Recently, however, Lim et al. '24 developed an approach called DEER, which evaluates nonlinear RNNs in parallel by posing the states as the solution to a fixed-point problem. They derived a parallel form of Newton's method to solve the fixed-point problem and achieved significant speedups over sequential evaluation. However, the computational complexity of DEER is cubic in the state size, and the algorithm can suffer from numerical instability. We address these limitations with two novel contributions. To reduce the computational complexity, we apply quasi-Newton approximations and show they converge comparably to Newton, use less memory, and are faster. To stabilize DEER, we leverage a connection between the Levenberg-Marquardt algorithm and Kalman smoothing, which we call ELK. This connection allows us to stabilize Newton's method while using efficient parallelized Kalman smoothing algorithms to retain performance. Through several experiments, we show that these innovations allow for parallel evaluation of nonlinear RNNs at larger scales and with greater stability.
title Towards Scalable and Stable Parallelization of Nonlinear RNNs
topic Machine Learning
I.2.6
url https://arxiv.org/abs/2407.19115