PaperlayerView exact revision on arXiv

Source-linked AI summary

Resurrecting Recurrent Neural Networks for Long Sequences

Antonio Orvieto, Samuel L Smith, Albert Gu, Anushan Fernando, Caglar Gulcehre, Razvan Pascanu, Soham De

arXiv:2303.06349v1cs.LG

TL;DR

RNNs are fast at inference but difficult to optimize and slow to train, whereas deep SSMs perform well on long sequences with efficient training and inference. The paper analyzes successive RNN modifications—linearization, diagonalization, stable parameterization, initialization, and normalization—and introduces the LRU. The resulting model matches deep SSM performance on LRA while matching their computational efficiency.

  • Problem

    RNNs are hard to optimize for long-range dependencies and slow to train, while the source of deep SSMs’ performance advantage over superficially similar RNNs remains unclear.

  • Method

    The paper uses signal-propagation analysis and ablations of linearization, diagonalization, stable exponential parameterization, initialization, and normalization to design the LRU.

  • Results

    LRUs match deep SSMs on all Long Range Arena tasks and match their training efficiency.

  • Takeaways & Limitations

    Strong long-range reasoning and efficient training do not require SSM discretization or structured transition matrices when RNNs are properly designed.

  • Takeaways & Limitations

    The S4 discretization result assumes the continuous-time input remains constant between timestamps, making this a modeling assumption.

Abstract

from arXiv · show

Recurrent Neural Networks (RNNs) offer fast inference on long sequences but are hard to optimize and slow to train. Deep state-space models (SSMs) have recently been shown to perform remarkably well on long sequence modeling tasks, and have the added benefits of fast parallelizable training and RNN-like fast inference. However, while SSMs are superficially similar to RNNs, there are important differences that make it unclear where their performance boost over RNNs comes from. In this paper, we show that careful design of deep RNNs using standard signal propagation arguments can recover the impressive performance of deep SSMs on long-range reasoning tasks, while also matching their training speed. To achieve this, we analyze and ablate a series of changes to standard RNNs including linearizing and diagonalizing the recurrence, using better parameterizations and initializations, and ensuring proper normalization of the forward pass. Our results provide new insights on the origins of the impressive performance of deep SSMs, while also introducing an RNN block called the Linear Recurrent Unit that matches both their performance on the Long Range Arena benchmark and their computational efficiency.

1. Introduction

RNNs are difficult to train and scale for long-range dependencies, while deep SSMs combine strong long-sequence performance with efficient training and inference. The paper shows that carefully designed RNNs can recover SSM-like performance and efficiency on LRA.

  • Motivation: RNNs suffer from vanishing and exploding gradients, making long-range dependencies difficult to learn, and their sequential computation slows optimization.These limitations make deep RNNs hard to scale despite their theoretical expressivity.
  • Motivation: Deep SSMs address long-sequence modeling with parallelizable training and RNN-like fast inference, avoiding attention’s quadratic bottleneck.Their performance on the Long Range Arena motivates investigating which design choices distinguish them from ordinary RNNs.
  • Research question: The paper asks whether deep RNNs can match the performance and efficiency of deep continuous-time SSMs on long-range reasoning tasks.The study directly compares recurrent designs within architectures modeled on S4-like deep sequence models.
  • Contribution: The resulting Linear Recurrent Unit recovers deep SSM performance and efficiency through small changes to a vanilla deep RNN.The paper emphasizes understanding and simplifying effective architectures rather than surpassing S4-based models.
  • Design steps: Removing recurrence nonlinearities substantially improves RNN accuracy, while nonlinear MLP or GLU blocks retain the ability to model complex sequence-to-sequence maps.Vanilla linear RNNs still lag S4 on some tasks, motivating further modifications.
  • Design steps: Diagonalizing the linear recurrence enables parallel scans without reducing expressivity or changing initialization features, improving training efficiency.The diagonal form also improves accuracy on some LRA tasks.
  • Design steps: Stable exponential parameterization and forward-pass normalization improve stability and long-range reasoning, allowing LRUs to match deep SSMs across LRA tasks.The paper links the normalization strategy to the discretization structure in S4.

2. Preliminaries

The paper contrasts vanilla RNN and S4-like recurrent blocks, then evaluates a shared deep architecture while progressively replacing SSM cores with RNN variants. Linear SSM recurrences support parallel scans, but discretization and parameter sharing distinguish them from ordinary linear RNNs.

  • RNN Layer: A recurrent layer maps an input sequence to an output sequence through a hidden state and learnable transition, input, output, and skip parameters.The hidden state starts at zero, and the recurrence may include a tanh or sigmoid nonlinearity.
  • S4-like recurrent layer: S4-like layers use linear state and input computations with mostly complex-valued parameters and a diagonal transition matrix initialized structurally.Their transition initialization is motivated by HiPPO theory.
  • S4-like recurrent layer: S4-like continuous-time systems are discretized for training and inference using Zero-Order-Hold or Bilinear methods at stepsize Δ.With piecewise-constant inputs, Zero-Order-Hold integration is exact, and diagonal structure makes discretization element-wise.
  • Computational differences: Linear recurrences can be parallelized with scans, whereas nonlinear RNN computation remains sequential.S4-like layers additionally use discretization-prescribed, parameter-shared matrices in which Δ affects both A and B.
  • Experimental setup: The experiments progressively replace SSM layers with RNN layers in a shared deep architecture to identify the modifications needed for S4-like LRA performance.The study uses six layers, residual connections, normalization, three runs, and reports means and standard errors.
  • Experimental setup: The Long Range Arena tests long-range sequence modeling, and is used because deep SSMs perform well on its tasks while Transformers perform poorly on most of them.The sequential CIFAR-10 task uses colored rather than grayscale images.

3. Designing Performant Deep RNNs

The paper progressively modifies vanilla deep RNNs to improve expressivity, stability, long-range modeling, and computational efficiency. These changes culminate in the Linear Recurrent Unit, which matches deep SSM performance and efficiency across LRA tasks.

  • 3.1. Linear RNN layers are performant: Linear recurrences improve test accuracy on most LRA tasks and remain expressive when combined with nonlinear MLP or GLU blocks.The deeper architecture can model nonlinear sequence-to-sequence maps without recurrent nonlinearities.
  • 3.2. Learning in the diagonalized space: Complex diagonal recurrent matrices preserve expressivity and initialization features while enabling highly parallelizable recurrence unrolling.Diagonal linear RNNs were approximately 8 times faster to train than dense ReLU RNNs on sCIFAR, matching S4D and S5 speed.
  • 3.2. Learning in the diagonalized space: Diagonalizing the recurrence improves accuracy on tasks such as ListOps and sCIFAR while substantially reducing training and inference time across LRA.The resulting RNN is as fast to train as deep SSMs such as S4D and S5.
  • 3.3. Benefits of stable exponential parameterization: Stable exponential parameterization enables stability during training and supports eigenvalue initializations near the unit disk for longer-range interactions.On Pathfinder, this change raises performance above 93%.
  • 3.4. Additional considerations for long-range reasoning tasks: Gamma normalization of hidden activations consistently improves tasks that benefit from eigenvalues initialized close to the unit disk, including sCIFAR and Pathfinder.The normalization addresses training-loss blow-up observed when eigenvalue magnitudes approach one.
  • 3.4. Additional considerations for long-range reasoning tasks: Combining linear recurrence, complex diagonalization, stable exponential parameterization, tuned initialization, and normalization yields the LRU.The LRU matches deep SSM performance and efficiency across all LRA tasks.

4. Insights on S4 and Variants

The paper argues that diagonal SSM performance comes primarily from linear recurrences, complex diagonal exponential matrices, and discretization-induced normalization and initialization, rather than HiPPO structure, parameter sharing, or continuous-time interpretation.

  • Exponential parameterization helps through magnitude-phase decoupling and diagonalized learning, not necessarily through accurate continuous-time integration.The authors also report benefits from stabilizing the recurrence with a nonlinearity.
  • Simple uniform initialization on a slice of the unit disk with proper normalization solves PathX, the hardest Long Range Arena task.Uniform initialization on the disk also approaches complex SSM performance on the remaining LRA tasks.
  • Linking eigenvalue phase and magnitude is unnecessary for good performance on Long Range Arena except PathX.A magnitude-independent initialization on the complete ring recovers accuracy on the other LRA tasks.
  • Discretization’s Δ multiplier effectively normalizes the recurrence and avoids blow-ups, similarly to the LRU’s γ normalization factor.This comparison follows from the recurrence’s long-run scaling under the first-order S4D approximation.
  • Decoupling the discretization parameter’s effects on the transition and input projection does not decrease S4 performance in the authors’ ablations.This suggests parameter sharing induced by ODE discretization is not necessary for S4-level performance.
  • Linear recurrences and complex diagonal exponential matrices, combined with normalization and initialization, account for diagonal SSM success.The authors distinguish these factors from discretization artifacts that do not necessarily contribute to performance.

5. Conclusion

The paper introduces the Linear Recurrent Unit as a deep RNN layer designed through initialization and forward-pass analysis, and reports that it matches modern deep SSMs on all LRA tasks.

  • The LRU combines linearization, diagonalization, stable exponential parameterization, and normalization to improve long-range reasoning.Its design does not rely on latent continuous-time discretization or structured transition matrices.
  • The final LRU model matches the performance of modern deep state-space models such as S4 or S5 on all Long Range Arena tasks.

A. Simplified Implementation of the Linear Recurrent Unit

The implementation section describes the LRU’s linear recurrent computation, its efficient parallel-scan execution, and its relationship to prior RNN and SSM architectures.

  • The LRU state follows x_k = Λx_k−1 + exp(γ_log) ⊙ (Bu_k), and outputs are computed as y_k = Cx_k + Du_k.The recurrence is efficiently parallelized using an associative scan.
  • The LRU uses Glorot-initialized input and output projections, a normalization factor, and an elementwise skip projection D.
  • The implementation uses JAX and NumPy, with a parallel scan for the diagonal linear recurrence and separate state and output projections.
  • Λ is initialized with complex values uniformly distributed on a ring between r_min and r_max, with phase restricted to [0, max_phase].
  • The related-work discussion contrasts LRU with gated RNNs, unitary RNNs, diagonal SSMs, S4, S5, and DLR across efficiency, parameterization, and long-range modeling.

C.1. Training speedups

The LRU achieves training speeds similar to S5 across tasks, while the comparison table reports steps-per-second measurements against tanh RNN, S4D, and S5 baselines.

  • The LRU has similar training speeds to S5 on all tasks.The comparison uses closely matched optimal model sizes for LRU and S5.
  • Table 4 reports training speed in steps/sec on an A100 GPU and includes LRU speedups over the tanh RNN for each task.

C.2. Effect of stability and normalization

Stability and normalization allow recurrent eigenvalues to be initialized closer to the unit disk without performance degradation, improving long-range reasoning and training stability.

  • Under the baseline exponential parameterization, performance degrades when r_max exceeds 0.9.
  • With enforced stability, optimal performance shifts toward higher r_max values, reaching its best setting at r_min = 0.0 and r_max = 0.99.
  • Adding γ normalization removes degradation at r_max = 0.999 and yields the best reported performance of 89.0%.The best setting uses r_min = 0.9 and r_max = 0.999.
  • These changes make training more stable and support better performance on tasks requiring very long-range dependencies.
  • The expanded results attribute the accuracy improvement to ring initialization near the circle boundary and normalization, which closes the gap with deep SSMs.

D.1. Architecture

The experiments use a six-block residual architecture in which recurrent layers are followed by mixing layers, with task-specific directional variants and parameterizations compared against SSM baselines.

  • Architecture: All experiments use networks with a depth of 6 residual blocks, each containing an identity skip connection and a residual path.
  • Architecture: Each RNN or SSM block places a recurrent layer before a GLU mixing layer, generally with dropout.PathX omits one additional linear transform and uses GLU without dropout.
  • Architecture: PathFinder and PathX use bidirectional models, whereas the remaining experiments use unidirectional models.
  • Architecture: Table 7 compares parameterizations of linear diagonal complex RNNs, including exponential representation and tuned ring initialization.
  • Architecture: Table 8 evaluates normalization with stable exponential parameterization and reports closely matching deep SSM results.PathX uses phases in [0, π/10], while other tasks use [0, 2π].

D.2. General experimental details

The paper combines controlled experimental settings with spectral and Koopman-theoretic analysis to explain why linear recurrent blocks can support nonlinear long-range modeling.

  • Experimental details: Training uses AdamW with linear learning-rate warmup followed by cosine annealing, starting and ending at 10^-7.
  • Experimental details: RNN/SSM parameters and applicable normalization or discretization parameters receive smaller learning rates than the base rate.
  • Experimental details: Experiments use six Long Range Arena tasks, replacing grayscale sCIFAR images with colored sCIFAR images.
  • Spectral and Koopman analysis: Interleaving linear RNNs with MLPs provides a flexible recipe for approximating nonlinear transition maps.
  • Spectral and Koopman analysis: ReLU causes spectral leakage by copying information across frequency components, including toward higher frequencies.
  • Spectral and Koopman analysis: Koopman theory motivates representing regular nonlinear dynamical systems with linear recurrent evolution after nonlinear reparameterization by an MLP.

E.2. Optimization of recurrent blocks

The optimization experiments show that removing recurrent nonlinearities and using exponential eigenvalue parameterization improves learning, while diagonalization preserves expressivity and supports efficient computation.

  • Recurrent nonlinearities: Linear RNNs outperform tanh RNNs even on the simple convolution task, supporting linear recurrence in deep architectures with interleaved MLPs.
  • Recurrent nonlinearities: For a length-100 convolution-learning task, the linear RNN converges faster than the tanh RNN across all tested learning rates.
  • Exponential parameterization: Exponential parameterization λ = exp(−ν + iθ) improves learning by decoupling phase and magnitude effects on recurrent powers.
  • Exponential parameterization: The advantage is strongest when the model must change phase while preserving initialization magnitude, because exponential-parameter gradients align with phase.
  • Diagonalization: Diagonalizing a real recurrent matrix yields real eigenvalues and conjugate complex-eigenvalue pairs, which can be represented using structured real 2 × 2 blocks.

F.1. Proof of Lemma 3.2

The lemma constructs a complex-valued sample uniformly on an annulus by independently sampling its phase and magnitude. The magnitude is obtained through the annulus’s radial distribution and inverse sampling.

  • The phase is sampled as θ = 2πu2, with u2 uniform on [0, 1].
  • The resulting exp(-ν + iθ) is uniformly distributed on the complex ring between radii rmin and rmax.
  • The target distribution is sampled by treating phase and magnitude independently.Phase is generated by scaling a uniform variable, while the magnitude is sampled separately.
  • The magnitude distribution is determined by the area of the annulus between radii rmin and rmax.Its radial cumulative distribution satisfies F_r(rmin) = 0 and F_r(rmax) = 1.
  • Writing r = e^-ν enables inverse sampling of ν from a uniform variable.The inverse sampling theorem gives ν = F^-1(u) for uniform u on [0, 1].

F.2. Proof of Proposition 3.3

The proposition analyzes the asymptotic squared state norm of a diagonal recurrence whose eigenvalues are sampled uniformly from a complex annulus. Its validity is checked numerically for constant and random inputs, with greater accuracy observed for random input.

  • Proposition 3.3 gives the limiting squared state norm for diagonal eigenvalues sampled uniformly between radii rmin and rmax under constant or white-noise input.
  • The proof first treats constant input, setting Bu_k = ũ for every timestep.
  • The gain formula is numerically validated with N = 500 and sequence length L = 10k over 10 runs, and is more accurate for random input.The figure compares the prediction with boxplot statistics for constant and random inputs.
  • For diagonal Λ, the eigenvalues are sampled with independent radius and phase within the annulus T(rmin, rmax).
  • The white-noise case simplifies because cross terms vanish when m ≠ n, allowing the proof to reduce to a single sum.
Loading 2303.06349v1…