Efficient Sparse Selective-Update RNNs for Long-Range Sequence Modeling

Abstract: Real-world sequential signals, such as audio or video, contain critical information that is often embedded within long periods of silence or noise. While recurrent neural networks (RNNs) are designed to process such data efficiently, they often suffer from ``memory decay'' due to a rigid update schedule: they typically update their internal state at every time step, even when the input is static. This constant activity forces the model to overwrite its own memory and makes it hard for the learning signal to reach back to distant past events. Here we show that we can overcome this limitation using Selective-Update RNNs (suRNNs), a non-linear architecture that learns to preserve its memory when the input is redundant. By using a neuron-level binary switch that only opens for informative events, suRNNs decouple the recurrent updates from the raw sequence length. This mechanism allows the model to maintain an exact, unchanged memory of the past during low-information intervals, creating a direct path for gradients to flow across time. Our experiments on the Long Range Arena, WikiText, and other synthetic benchmarks show that suRNNs match or exceed the accuracy of much more complex models such as Transformers, while remaining significantly more efficient for long-term storage. By allowing each neuron to learn its own update timescale, our approach resolves the mismatch between how long a sequence is and how much information it actually contains. By providing a principled approach to managing temporal information density, this work establishes a new direction for achieving Transformer-level performance within the highly efficient framework of recurrent modeling.

• The paper introduces selective-update RNNs, which use neuron-level binary gates to either update or exactly carry hidden states, shortening effective gradient paths from T steps to roughly p̄T update steps.
• The paper reports strong causal results, including a 70.43 Long Range Arena average, 84.92% on Pathfinder, 99.5% on Selective Copy, and 18.29 WikiText-103 perplexity at 100 million parameters.
• The paper shows that sparse execution can reduce suGRU step latency by 5.3× at 83% gate sparsity, while noting that gains depend on mask-aware kernels and that full backpropagation through long sequences remains costly.

Overview

This paper introduces Selective-Update RNNs (suRNNs), a recurrent architecture that replaces continuous gating with a per-neuron binary gate controlling whether each hidden unit updates or exactly carries its state at each time step (2603.02226). The motivation is a mismatch between sequence length and information density: conventional RNNs apply a dense, time-agnostic transition at every step, forcing constant overwriting of memory and degrading long-range credit assignment. By decoupling effective update frequency from raw sequence length, suRNNs aim to preserve memory exactly during low-information intervals while retaining $\mathcal{O}(1)$ streaming inference cost.

Architecture and mechanism

The core modification reparameterizes the standard recurrence $h_t = f_\theta(h_{t-1}, x_t)$ as

$$h_t = (I - D_t)\,h_{t-1} + D_t\,f_\theta(h_{t-1}, x_t),$$

where $D_t = \operatorname{diag}(g_t)$ and $g_t \in \{0,1\}^H$ is a neuron-wise binary gate. When $g_{t,i}=0$, neuron $i$ is carried identically; when $g_{t,i}=1$, it undergoes the usual nonlinear update. The one-step Jacobian takes the residual form $J_t = I + D_t(J_t^{(f)} - I)$, analogous to ResNet skip connections but along the temporal axis, so identity rows provide exact gradient carry routes.

Gates are generated by a rhythmic module: per-neuron learned amplitudes, phases, and biases applied to $K$ shared sinusoidal frequencies initialized on a logarithmic grid, thresholded by a Heaviside function trained via the straight-through estimator. The authors emphasize that this generator is a design choice; data-driven alternatives (MLPs, SSMs) can substitute without altering the theoretical properties of selective update.

The implementation strategy for scaling is notable: rather than custom stepwise control flow, the cuDNN-fused suGRU augments GRU inputs with a "skip drive" $h_t = f_\theta(h_{t-1}, x_t)$0 whose columns are hard-wired to the update-gate preactivation with a negative constant $h_t = f_\theta(h_{t-1}, x_t)$1, saturating $h_t = f_\theta(h_{t-1}, x_t)$2 toward zero on carry steps. This yields near-vanilla GRU throughput with asymptotic complexity unchanged, plus an always-on subset option guaranteeing live gradient paths.

Theory: shortened gradient paths

The central analytical result (Proposition 1) states that under $h_t = f_\theta(h_{t-1}, x_t)$3 on update steps, the row-wise sensitivity satisfies

$h_t = f_\theta(h_{t-1}, x_t)$4

where $h_t = f_\theta(h_{t-1}, x_t)$5 counts only timesteps on which neuron $h_t = f_\theta(h_{t-1}, x_t)$6 actually updated. With average update rate $h_t = f_\theta(h_{t-1}, x_t)$7, the effective multiplicative depth scales as $h_t = f_\theta(h_{t-1}, x_t)$8 rather than $h_t = f_\theta(h_{t-1}, x_t)$9. Unlike LSTM/GRU continuous gates, which merely rescale the spectral radius while still applying fractional updates every step, binary gating reduces the number of non-identity Jacobian factors by a factor of $$h_t = (I - D_t)\,h_{t-1} + D_t\,f_\theta(h_{t-1}, x_t),$$0. The paper also expands the unrolled product into a sum over subsets of update times, interpreting each gate realization as selecting a sparse sub-RNN — an ensemble view paralleling Veit et al.'s analysis of ResNets.

Empirically, on the Copying-Memory task with delay $$h_t = (I - D_t)\,h_{t-1} + D_t\,f_\theta(h_{t-1}, x_t),$$1, gated models converge faster and to lower loss than ungated baselines, and gradient norms remain bounded ($$h_t = (I - D_t)\,h_{t-1} + D_t\,f_\theta(h_{t-1}, x_t),$$2–$$h_t = (I - D_t)\,h_{t-1} + D_t\,f_\theta(h_{t-1}, x_t),$$3) and nearly parallel across increasing delays, consistent with the theory. A caveat the authors themselves note: because copy-task delays are filled with zeros during training, a standard GRU can exploit delay statistics as timing cues; replacing zero delay with noise at test time drops GRU accuracy from 27.81% to ~10%, while suGRU degrades from 85.31% to 48.56% — better, but still substantial degradation.

Benchmark results

Long Range Arena: Under strict uni-directional, token-by-token streaming, suGRU achieves 70.43 average, with 84.92% on Pathfinder — far above RWKV-v4 (58.42%) and typical causal RNNs, though below non-causal S4 (86.09 average). An ablation shows binary gating decisively outperforms a sigmoid variant (59.44 average), supporting the claim that discretization, not merely added gating, drives the gains.

Model	Streaming	LRA Avg
Transformer	No	53.66
S4	No	86.09
suGRU (binary)	Yes	70.43
suGRU (sigmoid)	Yes	59.44

Selective Copy ($$h_t = (I - D_t)\,h_{t-1} + D_t\,f_\theta(h_{t-1}, x_t),$$4): suGRU with an S4-style backbone reaches 97.2% with two layers and 99.5% with three, versus 16.3% for the ungated GRU counterpart — among the strongest results reported alongside S6-based models.

WikiText-103: parameter-matched suGRU attains test perplexity 19.20 versus 18.44 for a same-size Transformer; scaling to 100M parameters yields 18.29, and a hybrid interleaving suGRU with self-attention reaches 18.03. These are competitive but do not surpass the matched Transformer; the hybrid result indicates selective-update recurrence composes well with attention.

Pixel-level tasks: suGRU matches the best reported sMNIST accuracy (99.53%), reaches 98.46% on psMNIST, and substantially outperforms prior recurrent models on sCIFAR (87.26% vs. 84.65% for LSSL). Applied to spiking networks, su-SNN sets a reported state of 97.33% on psMNIST, exceeding specialized spiking baselines such as Rhythm-SNNs.

Mackey–Glass forecasting: suGRU is nearly horizon-invariant from 10 to 1000 prediction steps over length-5000 chaotic sequences, where LSTM/GRU degrade sharply beyond ~100 steps.

Efficiency: a mask-aware C implementation of stepwise suGRU at 83% gate sparsity reduces end-to-end step latency from 466 ms to 88 ms (5.3×), with update-gate and candidate computations reduced to ~8–9% of dense GRU cost. This speedup requires conditional execution not exposed by default PyTorch kernels; otherwise selective-update recurrence matches dense worst-case cost.

Gate-schedule sensitivity

Ablations on psMNIST show that correct update timing is essential: fixed random rhythms remain competitive (97.73%), but fixed random binarized gates, input-threshold gates ($$h_t = (I - D_t)\,h_{t-1} + D_t\,f_\theta(h_{t-1}, x_t),$$5), and naive learnable schedules collapse to chance or fail to optimize entirely. Learned sinusoidal schedules outperform fixed ones, indicating that end-to-end learning of heterogeneous per-neuron timescales — not merely the presence of sparsity — accounts for much of the benefit.

Limitations and open questions

Several constraints are acknowledged or evident. First, despite the theoretical path-length reduction, training still relies on full BPTT unrolling, which remains a bottleneck at extreme lengths; event-driven backpropagation or sparse checkpointing are proposed but not evaluated. Second, the mask-aware latency gains depend on implementation support for conditional execution, and the asymptotic complexity table's favorable entries assume such a mask-aware implementation. Third, suRNN trails non-causal SSMs on LRA overall (70.43 vs. 86.09), so the causal-streaming advantage does not close the gap on bidirectional spatial tasks. Fourth, the retention guarantee for long carry runs depends on a bounded margin assumption on the original update-gate preactivation, requiring normalization for a constant $$h_t = (I - D_t)\,h_{t-1} + D_t\,f_\theta(h_{t-1}, x_t),$$6 to suffice. Finally, whether context-aware or input-driven gate generators (rather than rhythmic schedules) would further improve performance on natural data remains untested.

Conclusion

suRNNs demonstrate that neuron-level binary selective updates — implemented efficiently within fused cuDNN kernels — structurally shorten BPTT credit-assignment paths and yield strong streaming performance on long-range benchmarks, including near-SOTA causal results on Pathfinder, Selective Copy, and psMNIST, and competitive WikiText-103 perplexity. The work establishes that exact identity-through-time carries, rather than spectral-radius engineering alone, are a viable route to Transformer-competitive recurrent modeling under strict causality, while leaving open the questions of scalable sparse backpropagation and richer gate generators noted above.