Technical Summary: Signal Propagation in Long-Range Linear Recurrences

Problem Statement

Understanding signal propagation in neural networks is fundamental to addressing the vanishing/exploding signal problem, which hinders the training of deep architectures. While existing theoretical frameworks rely heavily on the infinite-width limit ($n \to \infty$), where mean-field approximations yield tractable signal statistics, these theories often fail to capture the behavior of realistic networks operating at finite width.

This gap is particularly acute for recurrent models (such as Linear Recurrent Units and State-Space Models) designed for long-range sequence tasks. In these models, the same weight matrix $W$ is applied repeatedly over time steps $t$. The central question addressed is: How long does the infinite-width approximation remain accurate when the recurrent depth $t$ grows jointly with the width $n$? Specifically, at what depth scale do finite-width effects accumulate sufficiently to invalidate standard initialization schemes (e.g., Glorot) and infinite-width predictions?

Methodology

The paper analyzes linear recurrent models initialized with i.i.d. proper complex Gaussian entries of variance $1/n$. The models considered are:

1. Linear Recurrent Unit (LRU): $h^{(t)} = Wh^{(t-1)} + x^{(t)}$
2. Linear Recurrent Network (RNN): $h^{(t)} = Wh^{(t-1)}$ (without sequential inputs, $h^{(t)} = W^t x^{(0)}$)

The authors define signal energies as the normalized expected squared norms of the hidden states:
$$S_{n,t} = \frac{1}{n}\mathbb{E}[\|h^{(t)}_{LRU}\|^2], \quad Q_{n,k} = \frac{1}{n}\mathbb{E}[\|h^{(k)}_{RNN}\|^2]$$

Theoretical Derivation

To derive exact finite-width formulas, the authors employ a combination of random matrix theory and combinatorial analysis:

1. Trace Moment Reduction: The signal energies are reduced to the evaluation of the finite-width trace moment $\frac{1}{n}\mathbb{E}[\text{tr}((W^k)^* W^k)]$.
2. Wick Expansion: Using Wick's theorem for complex Gaussian variables, the trace moment is expanded into a sum over permutations $\sigma \in S_k$.
3. Combinatorial Counting: The contribution of each permutation is determined by the number of "free indices" (unconstrained equivalence classes) after applying Kronecker delta constraints. This count is shown to correspond exactly to the number of cycles in a specific commutator of permutations, $[\tau^{-1}, \tilde{\sigma}^{-1}]$, where $\tau$ is a long cycle.
4. Exact Formulas: Leveraging known results on the cycle counts of random permutations (specifically relating to Hultman numbers and Dubach's work), the authors derive exact closed-form expressions for $S_{n,t}$ and $Q_{n,k}$ involving rising and falling factorials.

Key Contributions

1. Exact Finite-Width Formulas

The paper provides the first exact closed-form formulas for signal propagation in finite-width linear recurrences with complex Gaussian initialization. These formulas (Theorem 2.1) are valid for arbitrary joint scalings of depth and width, moving beyond the standard infinite-width limit.

2. Identification of Three Scaling Regimes

Using the exact formulas, the authors characterize three distinct regimes governing signal propagation based on the relationship between depth ($t$ or $k$) and width ($n$):

• Subcritical Regime ($t, k = o(\sqrt{n})$):

  • The infinite-width approximation remains accurate.
  • Signal energies behave as $Q_{n,k} \to 1$ and $S_{n,t} \sim t+1$.
  • Standard Glorot initialization remains stable.
  • Recurrent and feedforward signal propagation are asymptotically indistinguishable in this regime.

• Critical Regime ($t, k \sim c\sqrt{n}$):

  • Non-negligible deviations from infinite-width predictions emerge.
  • A nontrivial joint scaling limit appears.
  • For RNNs: $Q_{n,k} \to \frac{2}{c^2}\sinh(\frac{c^2}{2})$.
  • For LRUs: $S_{n,t} \sim \sqrt{n} \int_0^{\tilde{c}} \frac{2}{x^2}\sinh(\frac{x^2}{2}) dx$, which grows exponentially with the scaling parameter.
  • Significance: This identifies $\sqrt{n}$ as the precise threshold where Glorot initialization becomes unstable in long-range linear recurrences.

• Supercritical Regime ($t, k \gg \sqrt{n}$):

  • Finite-width effects dominate, leading to signal explosion.
  • In the sublinear range ($\sqrt{n} \ll t \ll n$), signal growth becomes exponential in depth: $S_{n,t} \sim \frac{n^2}{t^3} \exp(\frac{t^2}{2n})$.
  • No finite joint scaling limit exists; signals diverge rapidly.

3. Distinction Between Recurrent and Feedforward Architectures

The paper demonstrates that finite-width effects accumulate much faster in recurrent models than in feedforward networks.

• In feedforward networks, finite-width corrections become significant when depth scales linearly with width ($L \sim cn$).
• In recurrent networks, these effects become visible at the square-root scale ($t \sim c\sqrt{n}$).
  This explains why recurrent models are empirically more prone to instability than feedforward models, a distinction invisible to standard infinite-width theories.

Results and Validation

• Theoretical Limits: The derived formulas precisely quantify the breakdown of the infinite-width approximation. They show that standard initialization schemes, while stable in the infinite-width sense, fail to stabilize signals once the recurrent depth exceeds the $\sqrt{n}$ scale.
• Real vs. Complex Weights: While the main results are derived for complex weights (common in modern SSMs), the paper establishes that real Gaussian weights exhibit even larger signal energies ($Q^R_{n,k} \ge Q_{n,k}$), implying that the instability thresholds derived for complex weights also apply (and are conservative) for real weights.
• Numerical Validation: Monte Carlo experiments confirm the theoretical predictions. The empirical signal energies align closely with the derived critical profiles (e.g., the $\sinh$ function in the critical regime) across various widths and depths.

Significance and Claims

The paper claims to provide a sharp answer to the question of how long the infinite-width approximation lasts in recurrent models. Its primary contributions are:

1. Pinpointing the Breakdown Scale: It establishes that the infinite-width theory breaks down at $t \sim \sqrt{n}$, a scale significantly smaller than the $t \sim n$ scale relevant for feedforward networks.
2. Explaining Empirical Instability: It offers a theoretical explanation for why standard initialization (Glorot) is unstable in modern long-range linear recurrences (like LRUs and SSMs), necessitating the rescaling strategies observed in practice.
3. Qualitative Difference: It highlights a fundamental qualitative difference in signal propagation between recurrent and feedforward architectures, driven by the repeated application of the same weight matrix.

The authors note that their results are specific to linear recurrences with random initialization. They acknowledge that extending these findings to nonlinear activations, structured non-Gaussian initializations (common in state-of-the-art models), and the training dynamics (feature learning regimes) remains an open direction for future work.