Parallelizable memory recurrent units

TL;DR

The paper addresses the challenge of achieving persistent memory in a model that remains parallelizable over sequence length. It introduces memory recurrent units (MRUs) that leverage multistability to store information indefinitely while maintaining parallelizable computations, and presents a concrete instantiation, the bistable memory recurrent unit (BMRU), based on hysteresis-inspired dynamics. The BMRU is designed to be compatible with the parallel scan algorithm and can operate alongside state-space models (SSMs) to combine transient linear dynamics with stable memory, as demonstrated on tasks requiring long-term dependencies. Empirical results on Copy-first-input, permuted sequential MNIST, and Pathfinder show BMRU’s strengths in memory retention with shallow networks, and the work outlines promising directions for hybrid models and further optimization. Overall, MRUs offer a new paradigm that unifies persistent memory with parallelizable RNN computation, potentially enriching scalable sequence modeling in practice.

Abstract

With the emergence of massively parallel processing units, parallelization has become a desirable property for new sequence models. The ability to parallelize the processing of sequences with respect to the sequence length during training is one of the main factors behind the uprising of the Transformer architecture. However, Transformers lack efficiency at sequence generation, as they need to reprocess all past timesteps at every generation step. Recently, state-space models (SSMs) emerged as a more efficient alternative. These new kinds of recurrent neural networks (RNNs) keep the efficient update of the RNNs while gaining parallelization by getting rid of nonlinear dynamics (or recurrence). SSMs can reach state-of-the art performance through the efficient training of potentially very large networks, but still suffer from limited representation capabilities. In particular, SSMs cannot exhibit persistent memory, or the capacity of retaining information for an infinite duration, because of their monostability. In this paper, we introduce a new family of RNNs, the memory recurrent units (MRUs), that combine the persistent memory capabilities of nonlinear RNNs with the parallelizable computations of SSMs. These units leverage multistability as a source of persistent memory, while getting rid of transient dynamics for efficient computations. We then derive a specific implementation as proof-of-concept: the bistable memory recurrent unit (BMRU). This new RNN is compatible with the parallel scan algorithm. We show that BMRU achieves good results in tasks with long-term dependencies, and can be combined with state-space models to create hybrid networks that are parallelizable and have transient dynamics as well as persistent memory.

Figures (12)

• Figure 1: Monostability vs bistability in a RNN with internal clock. The figure shows internal state trajectories of the RNN unit described by \ref{['eq:bistRNN']} for different initial conditions $\tilde{h}_t[0] = h_{t-1}$. (left) Evolution of the system when $\beta = -1.5$. (right) Evolution of the system when $\beta = 1.5$.
• Figure 2: Convergence properties of the RNN unit described by \ref{['eq:bistRNN_clock']} for different values of input $x_t$ and either $\beta=-1.5$ (A) or $\beta=1.5$ (B) (left) Internal state trajectories corresponding to different initial conditions $\tilde{h}_t[0]=h_{t-1}$ and 4 different input values $x_t$. (right) Solutions of the steady-state equation \ref{['eq:hyst']} for $\beta=-1.5$ (A) and $\beta=1.5$ (B). The red arrows show convergence trajectories from different initial conditions $\tilde{h}_t[0]$.
• Figure 3: Comparison between the implicit function and its approximation. This figure compares the solutions $(h_t, x_t)$ of the implicit function defined by \ref{['eq:hyst']} (in blue) with the approximation defined by \ref{['eq:approx']} (in red) where $\alpha = 1$. Solid lines correspond to stable points, and dashed lines to unstable points.
• Figure 4: Example simulation of the cell defined by \ref{['eq:final_approx']} for different values of $x_t$ and $\beta$. The top graph defines two illustrative variations of $x_t$ and $\beta$, as well as the bistable region defined by $\beta$. The bottom graph shows the effect of variations in $x_t$ and $\beta$ on the evolution of the state $h_t$. The gray areas highlight the timesteps at which the memory is updated, which occurs when $x_t$ is outside of the bistable region.
• Figure 5: Surrogate gradient used in BMRU. (left) Comparison between the Heaviside function and the function defined by \ref{['eq:surr_func']} used to approximate the gradient for different values of $\alpha_{surr}$. (right) Impact of $\alpha_{surr}$ on the surrogate gradient defined by \ref{['eq:surrogate']}.

Theorems & Definitions (6)

• Theorem 1
• proof
• Theorem 2
• proof
• Theorem 3
• proof