Research map / Learning theory and algorithms / Combinatorial and search optimization
Towards Scalable and Stable Parallelization of Nonlinear RNNs
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes quasi-DEER (diagonal Jacobian approximation) and ELK (Levenberg-Marquardt trust regions solved via parallel Kalman filtering) to make parallel evaluation of nonlinear RNNs scalable and numerically stable.
Problem
DEER evaluates nonlinear RNNs in parallel via Newton's method but has cubic work and quadratic memory in state size and is numerically unstable because undamped Newton can diverge.
Approach
Replace the dense Jacobian with its diagonal (quasi-Newton) to get O(TD) memory and O(TD) work per iteration, and stabilize Newton by adding a Levenberg-Marquardt trust region whose linear-Gaussian MAP solve is a parallel Kalman filter/smoother (ELK), with the combination quasi-ELK.
Key results
- Quasi-DEER matches DEER runtime while using up to an order of magnitude less memory, and both are up to 20x faster than sequential evaluation on an untrained GRU.
- On the trained autoregressive GRU, quasi-ELK takes 221 ms versus 1,225 ms for DEER, though sequential evaluation (96 ms) is faster.
- In C. elegans GRU training, each quasi-DEER training step is 2.5x faster than DEER despite about 2x more Newton iterations, with similar validation accuracy.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace dense Jacobians with diagonal Jacobian approximations | DEER on untrained GRU evaluation · Untrained GRU evaluation | memory usage: as much as an order of magnitude less memory | Helped |
| Replace dense Jacobians with diagonal Jacobian approximations | DEER on untrained GRU evaluation · Untrained GRU evaluation | runtime vs sequential evaluation: up to twenty times faster | Helped |
| Replace dense Jacobians with diagonal Jacobian approximations | DEER for C. elegans classification training · C. elegans time-series classification | training iteration runtime: 2.5 times faster | Helped |
| Replace dense Jacobians with diagonal Jacobian approximations, including in the backward pass | DEER for C. elegans classification training · C. elegans time-series classification | validation accuracy | No clear effect |
| Add trust-region damping via Kalman filtering | DEER on a trained autoregressive GRU · Trained autoregressive GRU | convergence stability | Helped |
| Add trust-region damping via Kalman filtering | DEER on a trained autoregressive GRU · Trained autoregressive GRU | convergence speed | Helped |
| Combine diagonal Jacobian approximation with ELK trust-region damping | ELK on a trained autoregressive GRU · Trained autoregressive GRU | runtime: 221 milliseconds | Helped |
| Use DEER instead of sequential generation | Trained autoregressive GRU · Trained autoregressive GRU | runtime: 96 milliseconds; 1,225 milliseconds | Hurt |
| Use ELK-style parallel methods instead of sequential generation | ELK and quasi-ELK · Trained autoregressive GRU | runtime: 221 milliseconds; 96 milliseconds | Hurt |
| Use the diagonal Jacobian approximation rather than the exact Jacobian (ablation) | DEER · General nonlinear RNN evaluation | Newton iterations to convergence: approximately 2 times more Newton iterations | Hurt |
Surprising findings
- On the trained autoregressive GRU, all parallel methods including DEER are slower than sequential generation.
- DEER was numerically unstable on an A100 for the C. elegans training experiment while quasi-DEER was not.
Limitations
- Quasi-Newton methods lose local quadratic convergence and need more iterations.
- ELK has cubic complexity in state dimension; quasi-ELK adds a trust-region hyperparameter that must be swept.
- On the autoregressive GRU task, all parallel methods are slower than sequential evaluation.
Experimental setup
Datasets: C. elegans eigenworms time series synthetic noisy sine wave
Benchmarks: untrained GRU evaluation timing/memory C. elegans eigenworms classification autoregressive GRU sine wave generation
Metrics: wall-clock time peak memory mean absolute discrepancy validation accuracy Newton iterations
Baselines: sequential evaluation DEER
Builds on: DEER Levenberg-Marquardt parallel associative scan iterated extended Kalman smoother
Closest papers by meaning
- Parallel-in-Time Training of Recurrent Neural Networks for Dynamical Systems Reconstruction (NeurIPS 2026 · similarity 0.86)
- Predictability Enables Parallelization of Nonlinear State Space Models (NeurIPS 2025 · similarity 0.86)
- Fixed-Point RNNs: Interpolating from Diagonal to Dense (NeurIPS 2025 · similarity 0.84)
- Learning Rate Transfer for Hybrid Transformer-SSM Architectures (NeurIPS 2026 · similarity 0.83)
- FlashRNN: I/O-Aware Optimization of Traditional RNNs on modern hardware (ICLR 2025 · similarity 0.83)
- Accelerating Sinkhorn algorithm with sparse Newton iterations (ICLR 2024 · similarity 0.82)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.