Learning dynamics in linear recurrent neural networks

Learning dynamics in linear recurrent neural networks

Speakers: Alexandra Proca, Clémentine Dominé, Murray Shanahan, Pedro Mediano

Conference: ICML 2025

YouTube: https://slideslive.com/39044097

Overview

This talk, presented by Alexandra Proca at ICML 2025, delves into the intricate mechanisms of learning within linear recurrent neural networks (RNNs). The research addresses a critical gap in the understanding of how RNNs, vital for tasks involving temporal dependencies in both machine learning and neuroscience, acquire their functional structures during training. While much prior work has focused on analyzing the properties of trained RNNs, this presentation shifts the focus to the dynamic learning process itself, exploring how network parameters evolve when exposed to temporally structured data.

The work builds upon the analytical tractability offered by deep linear networks to extend this investigative approach to recurrent architectures. By providing a novel, flexible framework, Proca and her co-authors illuminate how task dynamics—the temporal structure of the input data—and the inherent recurrence of these networks fundamentally shape various properties, including learning speed, generalization, stability, and the nature of learned representations. This research is pivotal for advancing both the theoretical understanding of neural computation in the brain and for guiding the design and training of more robust and efficient RNN models in artificial intelligence.

The significance of this study is amplified by the resurgence of RNNs for long-range sequence modeling and their widespread use as cognitive models in neuroscience. By dissecting the learning dynamics in a simplified yet insightful linear setting, the authors offer foundational insights into complex phenomena observed in their non-linear counterparts. This analytical framework promises to bridge theoretical neuroscience and deep learning, providing a deeper understanding of how temporal processing capabilities emerge and adapt.

Background

▶ Watch: Introduction: RNNs, learning dynamics, and linear models (0:00)

Recurrent Neural Networks (RNNs) stand as cornerstone architectures for processing sequential data, playing a crucial role across diverse fields. In machine learning, their ability to model temporal dependencies has driven significant advancements, particularly with the recent resurgence in their popularity for long-range sequence modeling. Concurrently, in neuroscience, RNNs serve as prevalent cognitive models of neural dynamics, often replicating observed population dynamics in animal brains when trained on abstract versions of the same tasks. This dual importance has spurred extensive theoretical inquiry from both deep learning theorists and neuroscientists, all striving to unravel the mysteries of neural computation.

Prior research by Laura Driscoll and colleagues, for instance, demonstrated that RNNs trained on a multitude of tasks develop distinct "dynamical building blocks" that are then flexibly composed to address different computational challenges. However, a pervasive limitation in much of this work is its focus on the end state of training, analyzing the properties of the final solutions discovered by the networks while largely overlooking the dynamic learning process itself. Consequently, the mechanisms by which the underlying functional structures of RNNs emerge from training on temporally structured data remain largely unexplored.

A separate, influential line of research, notably pioneered by Andrew Saxe and colleagues, has championed the use of deep linear networks to study learning dynamics. Despite their functional simplicity, these networks exhibit surprisingly complex non-linear learning dynamics, yet crucially, they remain analytically tractable. This unique combination has established deep linear networks as an invaluable framework for theoretical investigation. Saxe's work has been instrumental in proposing theories for semantic development and cognitive flexibility, advancing the theoretical understanding of feature learning, and elucidating its dependence on weight initialization. Furthermore, this analytical approach has even been extended to certain non-linear cases.

However, a significant constraint of this highly successful body of work is its predominant focus on feedforward networks. While providing profound insights into static input-output mappings, it largely neglects the complexities introduced by dynamic task settings and the rich computational power inherent in recurrent architectures. Given that recurrence is a fundamental component of neural computation and essential for processing time-varying information, there has been a pressing need to extend the analytical tools developed for deep linear networks to the recurrent domain. This talk directly addresses this need, aiming to characterize learning more broadly by accounting for the critical impact of dynamic tasks and the unique computational capabilities of recurrent networks.

Key Findings

▶ Watch: Understanding independent 'connectivity modes' in RNNs (4:00)

The research presented by Alexandra Proca and co-authors introduces a groundbreaking framework for analyzing learning dynamics in linear recurrent neural networks (RNNs), yielding several profound insights into how these networks learn and generalize.

Firstly, the talk reveals that the learning speed in RNNs is considerably more complex than in their feedforward counterparts. While feedforward networks typically prioritize learning components associated with the largest singular values, RNNs, due to their temporal nature, exhibit a dependency on both the scale and the temporal ordering of singular values. Specifically, the study finds that singular values appearing later in the input trajectory are learned faster, demonstrating a recency bias in how RNNs prioritize computations.

Secondly, the work rigorously demonstrates how task dynamics fundamentally impact critical properties such as extrapolation ability and network stability. The authors define specific classes of task dynamics that, even when perfectly learnable, can lead to solutions that fail to extrapolate to sequence lengths different from those encountered during training. Similarly, task dynamics characterized by decreasing singular values through a trajectory are shown to drive the network's recurrent modes to values greater than one, leading to unstable training regimes and potential exploding gradients or diverging hidden layer activity.

Thirdly, the research explores how RNNs learn when presented with tasks that are not perfectly learnable, revealing computational tradeoffs. By varying the balance between recurrent and feedforward computational components within the task dynamics, the authors identify a phase transition. If the feedforward computation dominates, the network actively prunes the recurrent mode, leading to low-rank solutions where the input-output mode handles the computation. Conversely, as the recurrent component's contribution to the error increases, the network rapidly transitions to a regime where it approximates the recurrent computation while effectively ignoring the feedforward part. This highlights how RNNs adapt their internal structure based on the dominant computational demands of the task.

Finally, the study provides compelling evidence that recurrence facilitates rich feature learning. By analyzing the Neural Tangent Kernel (NTK) for finite-width linear RNNs, the researchers observe that recurrence promotes greater "kernel movement" from initialization to the end of training across a wide range of settings and initializations. This suggests that RNNs inherently possess a bias towards learning structured, task-relevant representations (a rich feature learning regime) rather than simply performing high-dimensional projections of the input (a lazy learning regime). This finding has significant implications for understanding the representational capabilities of recurrent architectures.

Technical Deep Dive

▶ Watch: Learning speed: singular values in RNNs vs. feedforward (4:40)

The core of this research lies in developing a novel, analytically tractable framework to study learning dynamics in linear recurrent neural networks (RNNs). The model under consideration is a standard linear RNN that receives an input at each time step and updates its hidden state. For simplicity, the initial derivation focuses on a loss computed over a single output at the final time step, though the paper generalizes this to the autoregressive case where an output is produced at each time step.

A key innovation is the representation of data and task statistics. For each input at time T and the final output, the authors construct a data correlation matrix Sigma YXT. Assuming whitened input, the entire task is fully specified by the sequence of these data correlation matrices within the loss function. These matrices are further represented through their singular values, enabling the framework to account for complex task dynamics. By making a crucial assumption that the left and right singular vectors (U and V) remain constant across all time steps—an assumption justified by prior work suggesting early alignment of eigenvectors in training—the task's dynamics are elegantly contained within the sequence of singular value matrices, denoted as S. This S sequence captures how the input through time correlates to the output.

The most significant technical simplification arises from the derivation of a form where the learning dynamics become diagonalized. This is achieved through a conceptual "rotation" or "change of basis" of the network. In this transformed space, the complex, fully connected network is simplified into a diagonalized network where learning in each dimension is entirely decoupled. This decoupling means that the gradient descent dynamics for each dimension are specific and independent, driven solely by the singular value trajectory in that particular dimension. The gradients of the network parameters are thus reduced to a set of scalar ordinary differential equations (ODEs), each governing a single "connectivity mode" along these independent chains or axes.

Within this framework, the authors decompose the functions performed by RNNs into two primary components: an input-output mode that performs a constant scaling, and a recurrent mode that learns a time-dependent function. By separately deriving solutions for the learning dynamics of these modes, they confirm that each indeed learns its respective component. The singular values themselves are decomposed into a constant term, delta, and a function F parameterized by lambda and the trajectory time step, allowing for a detailed analysis of their impact.

A critical observation regarding learning speed emerged from analyzing the gradient flow equations. Unlike feedforward networks where the largest singular values are learned first, RNNs exhibit a more nuanced behavior. The research identifies an exponential term, specifically on the recurrent mode B, which, when initialized to be less than one, downscales the gradient contribution of earlier time steps. This mathematical property directly explains the observed recency bias, where singular values occurring later in the trajectory are learned faster and have a greater impact on learning speed.

The study also delves into extrapolation and stability issues. By defining two classes of task dynamics with zero-loss solutions, the authors illustrate how specific dynamics (e.g., exponential dynamics) lead to solutions that are inherently dependent on trajectory length and will not extrapolate. For stability, the analysis of solutions reveals that task dynamics with strongly decreasing singular values through the trajectory length drive the network's recurrent modes to increase over one. This instability, leading to exploding gradients or diverging hidden layer activity, is a direct consequence of the solutions the network is driven to by these specific task dynamics, particularly those with strong early dependencies.

In scenarios where tasks are not perfectly learnable, the energy function is rewritten as a sum of two terms. When the RNN cannot perfectly fit the data, the second term acts as an effective regularizer, incentivizing connectivity modes to remain close to zero. This insight motivates the study of computational tradeoffs by varying the recurrent and feedforward components of the task dynamics. The observed phase transition demonstrates that if the energy is dominated by feedforward computation, the network prunes the recurrent mode, resulting in low-rank connectivity. Conversely, if the recurrent component's error contribution increases, the network rapidly approximates the recurrent computation. The critical point of this transition is dependent on the ratio of feedforward to recurrent computation.

Finally, to understand feature learning, the authors employ the Neural Tangent Kernel (NTK). They derive the NTK for finite-width linear RNNs and measure the kernel distance—the change in activations—from network initialization to the end of training. By varying trajectory length and initialization scale, they find that recurrence consistently facilitates greater kernel movement across many settings. This kernel movement is a hallmark of the rich feature learning regime, where networks learn structured, task-relevant representations, contrasting with the lazy learning regime where NTK remains constant. This suggests that recurrence inherently biases RNNs towards rich feature learning. The framework is also extended to handle rotational dynamics through an eigen decomposition, allowing for the study of more complex task dynamics.

Experimental Setup & Results

▶ Watch: How task dynamics affect RNN extrapolation ability (6:20)

The talk primarily presents a theoretical and analytical framework, deriving mathematical insights into linear RNN learning dynamics rather than presenting traditional experimental results on large-scale datasets. The "experiments" in this context refer to the analytical derivations and numerical simulations of the derived scalar ODEs and energy functions, demonstrating the properties of the connectivity modes under various task dynamics.

The authors construct specific scenarios by manipulating the sequence of singular value matrices (S), which define the task dynamics. For instance, they define two classes of task dynamics where the dynamic part of the singular values is set by a specific function. These are proven to be the only task dynamics with zero-loss solutions in their linear RNN model, providing a controlled environment to study extrapolation ability. By analyzing the derived solutions, they show that one class, characterized by exponential dynamics, leads to solutions inherently dependent on trajectory length, preventing extrapolation to different sequence lengths.

Regarding learning speed, numerical plots illustrate the error through training for different connectivity modes exposed to singular value trajectories. These plots visually confirm that modes trained on singular values that are large towards the end of the trajectory are learned faster, directly supporting the "recency bias" finding.

For stability, the analytical study of solutions reveals that task dynamics with decreasing singular values throughout the trajectory drive the network's recurrent modes to values greater than one. This derived instability indicates that tasks with strong early dependencies will lead to unstable training regimes.

The investigation into computational tradeoffs involves varying the recurrent (S1) and feedforward (ST) components of the task dynamics and observing the convergence of the network's connectivity modes. The results analytically demonstrate the existence of a phase transition: when feedforward computation dominates, the network prunes the recurrent mode; conversely, as the recurrent mode's error contribution increases, the network rapidly approximates recurrent computation. The critical point of this transition is shown to depend on the ratio of feedforward to recurrent computation.

Finally, to study feature learning, the Neural Tangent Kernel (NTK) is derived for finite-width linear RNNs. The "results" here involve measuring the kernel distance of the network from its initialization to the end of training under various conditions (varying trajectory length and initialization scale). The finding that recurrence facilitates greater kernel movement across many settings and initializations is a direct analytical result from this derivation, suggesting a bias towards rich learning.

The talk also mentions extensions of the framework to more complex scenarios, including sensory integration tasks for validation and to rotational dynamics using an eigen decomposition for broader applicability of the theory. While not traditional experimental results, these analytical findings provide concrete, quantifiable insights into the behavior of linear RNNs under diverse conditions.

Practical Implications

▶ Watch: Task dynamics' role in RNN stability problems (7:00)

The theoretical insights derived from this study on learning dynamics in linear RNNs carry significant practical implications for various stakeholders in the machine learning and neuroscience communities.

For practitioners and model builders, understanding the profound impact of task dynamics on RNN behavior is paramount. The finding that the temporal structure of the data dictates learning speed, generalization, and stability suggests that careful consideration of the input sequence statistics is as crucial as architectural design. For instance, if a task exhibits strong dependencies early in a sequence (e.g., decreasing singular values), model builders should be aware of the increased risk of unstable training regimes and potential exploding gradients. This might necessitate stronger regularization techniques or alternative architectural choices designed to mitigate such instabilities.

The challenges related to extrapolation ability highlight a critical limitation for deploying RNNs in real-world scenarios where sequence lengths might vary from training data. If task dynamics inherently lead to solutions dependent on specific trajectory lengths, models trained on short sequences may perform poorly on longer ones, and vice versa. This implies a need for training methodologies or architectural inductive biases that explicitly promote length generalization, or at least a thorough evaluation of generalization capabilities across a range of sequence lengths.

The discovery of a phase transition and the concept of low-rank solutions when certain computational components dominate offer avenues for designing more efficient RNNs. If a task primarily requires feedforward computation, the network may naturally prune its recurrent modes. This suggests that for specific task types, simpler, more computationally efficient architectures with reduced recurrent connectivity might suffice, leading to faster training and inference. Understanding these tradeoffs can guide the development of sparse or dynamically pruned recurrent networks.

The "recency bias" observed in learning speed, where later singular values are learned faster, implies that RNNs might prioritize recent information. While this can be beneficial for tasks requiring immediate responses to recent inputs, it could pose challenges for tasks demanding robust integration of very long-term dependencies. Practitioners might need to adapt loss functions, employ attention mechanisms, or modify training schedules to counteract this bias when long-term memory is critical.

Furthermore, the evidence that recurrence facilitates rich feature learning is encouraging for tasks requiring complex, structured representations. This suggests that RNNs, even in their linear form, inherently move beyond simple lazy learning to develop more meaningful internal representations. This bias towards "rich learning" could be a significant advantage when building models for tasks like natural language understanding or complex sequential decision-making, where the quality of learned features is crucial for performance and interpretability. It also suggests that for tasks demanding rich features, leveraging recurrence might be more effective than simply increasing network width in a feedforward setting.

However, it is crucial to acknowledge the limitations of this work. The analytical tractability is largely enabled by focusing on linear RNNs. While providing fundamental insights, the non-linearities present in most modern RNNs (e.g., LSTMs, GRUs, Transformers) introduce additional complexities that may alter some of these dynamics. The assumption of constant singular vectors throughout most of training, though supported by early alignment, also simplifies the dynamics of basis learning. Future work would need to investigate how these findings translate to and are modified by non-linear activations and dynamic basis changes. Nonetheless, this research provides a robust theoretical foundation for understanding the intricate learning processes within recurrent architectures, offering guiding principles for more complex models.

Key Takeaways

• Learning dynamics in RNNs are distinct and complex: Unlike feedforward networks, RNN learning speed is influenced by both the scale and temporal ordering of singular values, not just their magnitude.
• Task dynamics critically shape RNN properties: The temporal structure of input data fundamentally impacts learning speed, generalization (extrapolation), and network stability, potentially leading to poor extrapolation or unstable training regimes.
• Recency bias in learning: RNNs exhibit a bias towards learning features associated with singular values that occur later in the input trajectory, prioritizing more recent information.
• Computational tradeoffs and phase transitions: Depending on the balance of feedforward versus recurrent computational demands, RNNs can undergo a phase transition, leading to the pruning of specific connectivity modes and the emergence of low-rank solutions.
• Recurrence facilitates rich feature learning: Analysis using the Neural Tangent Kernel suggests that recurrence biases networks towards learning structured, task-relevant representations (rich learning) rather than high-dimensional, less structured projections (lazy learning).
• Novel analytical framework: The research introduces a flexible and analytically tractable framework for studying linear RNNs, enabling a deeper understanding of how task dynamics and recurrence interact to shape the learning process.

About the Speaker(s)

The talk was presented by Alexandra Proca, who, along with her co-authors Clémentine Dominé, Murray Shanahan, and Pedro Mediano, conducted this research on the learning dynamics of linear recurrent neural networks. Clémentine Dominé was also referenced during the Q&A session for her related work on lifting certain assumptions in similar analytical frameworks. The team's collective expertise spans the theoretical underpinnings of machine learning and neuroscience, focusing on understanding complex neural computations.