arXiv:2609.34034v1 [cs.NE] 27 Sep 2026
ADPTN ET: A DAPTIVE WITH P RESCRIPTIVE T IMESCALES N ON -L INEAR SSM FOR S EQUENCE M ODELLING
Matei-Ioan Stan * International Centre for Neuromorphic Systems The University of Manchester Manchester, United Kingdom [email protected]
Oliver Rhodes International Centre for Neuromorphic Systems The University of Manchester Manchester, United Kingdom [email protected]
September 29, 2026
A BSTRACT A central aim of neuromorphic computing is to provide a viable alternative to highly energy-intensive Transformer-based AI. However, efficient alternatives struggle to capture the set of qualities that have secured the Transformer’s status as the de facto standard in sequence modelling. Any realistic contender must be data-adaptive, able to capture long-range dependencies, and GPU-parallelisable, but also non-linearly recurrent to enable complex reasoning. Based on evidence suggesting the auditory cortex operates on fixed timescales, this work proposes the ADaptive with Prescriptive Timescales Network (ADPTNet) as a potential solution to achieving all four properties simultaneously. ADPTNet is built around local topological conjugates, obtained by a novel combination of linear attention and Riemannian optimisation, applied to static global dynamics. This enables non-linear yet predictable long-term behaviour. Dynamical systems theory proofs provide theoretical guarantees for the parametric control of ADPTNet’s timescales (its Lyapunov spectrum). ADPTNet improves performance on Selective Copying over Hawk, the existing method balancing long-range memory and adaptability, while also improving state tracking over linear SSMs like Mamba. On sequential CIFAR10, ADPTNet matches linear SSM accuracy and outperforms existing selective models (incl. the Transformer), using fewer parameters. We also introduce a neuromorphic SpikingADPTNet, which achieves a new state-of-the-art accuracy on the Spiking Speech Commands dataset (83.56% ± 0.15). Finally, ADPTNet’s constant timescales enable two efficient, Jacobian-free extensions to the DEER parallel simulation algorithm (Conv and Forward DEER) that retain the same average convergence. Conv DEER adds no computational overhead beyond the network’s forward pass and enables nonlinear RNN parallelisation via iterated convolutions for the first time. Keywords Spiking Neural Networks · State Space Models · Sequence Modelling · Selectivity · State Tracking · Dynamical Systems · Long Range Dependencies
1
Introduction
In recent years, sequence-based tasks have proven a hotbed of progress in deep learning research. This is best exemplified by the billion-dollar industry that has emerged around Large-Language Models (LLMs), essentially everimproving token-sequence compression and decoding neural network algorithms. More specifically, this massive effort to build state-of-the-art LLMs has largely relied on developing variations of what has effectively become the standard architecture, Transformers [177]. Some research on the remarkable success of Transformers in Natural Language Processing (NLP) applications has focused on how their core component, self-attention, overcomes vanishing and exploding gradients that affect Recurrent Neural Networks (RNNs) [93, 142], enabling the use of longer context windows. Another strong argument for the Transformer’s ubiquitous adoption has been the formulation of self-attention around massive and parallelisable matrix multiplications, which are well-positioned to take advantage of modern GPUs and thus are currently "winning" the
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
hardware lottery [96, 94]. A final research direction for explaining the Transformer’s performance, which has attracted increasing interest over time, has been to formalise the role of the architecture’s perceived "ad aptability". This latter property of self-attention layers has been described in literature using terms at various levels of abstraction. Some have defined adaptability in relation to the flexibility granted by direct modelling of token-to-token interactions [116]. Others have focused on how self-attention implicitly defines the weights used for linear sequence-wise mixing as functions conditioned on the input tokens, thereby making it a data-controlled operation [126, 145]. Links have also been drawn between attention key-value pairings and associative Hopfield Networks [148]. Furthermore, associative storage and recall are thought to be key components of in-context learning, which refers to the emergent ability of Transformer-based LLMs to effectively predict unseen inputs based on information provided through prompts at inference-time [7, 138]. The dominance of Transformer-based architectures, however, is not unchallenged, with a major current challenge being the infamous energy-intensity of training and deploying state-of-the-art LLM systems [162]. From an algorithmic perspective, this is partly due to quadratic-in-sequence-length compute and memory scaling of self-attention. In this regard, recurrent architectures have seen a resurgence of interest based on comparatively favourable sub-quadratic scaling. Most notably, linear Transformers [170] and State Space Models (SSMs) [74] have emerged as overlapping research streams focused on matching or outperforming Transformers in associative recall and long-context sequence modelling, respectively. Consequently, contributions from both fields have provided efficiency gains in industrial LLM development [20, 171, 25]. However, fundamentally, they still rely on dense floating-point Multiply-Accumulate (MAC) operations on GPUs, much like Transformers, and thus inherit a similar energy-consumption paradigm [8]. One might assume that this taxing energy consumption is an inherent drawback of intelligent systems, were it not for the evident counterexample of the human brain. The cortex vastly outperforms current frontier models, including in key areas of deep learning research, such as reasoning and few-shot generalisation, while using a fraction of the energy [82]. Neuromorphic computing aims to bridge this gap by developing algorithms and computational substrates that emulate the brain’s efficacy and efficiency. Perhaps most famously, Spiking Neural Networks (SNNs), the third generation of neural network models [123] explicitly aim to reduce energy consumption by avoiding MAC operations through sparse, binary communication between biologically plausible neurons. SNNs are also intended for deployment on neuromorphic hardware, which can leverage spike-train sparsity and potentially avoid the von Neumann bottleneck of traditional hardware, including GPUs [114]. Consequently, neuromorphic systems have been shown to reduce energy consumption by orders of magnitude compared to mainstream counterparts [39, 149]. While SNNs have a proven track record of reducing energy consumption, to date, it remains mostly confined to smallscale machine learning applications. As neuromorphic algorithms scale, they typically begin to sacrifice biological inspiration, creating an inherent tension in their definition. For instance, scaling SNNs in depth typically requires loosening the definition of spikes (e.g., "graded" integer spikes [55]), or replacing them altogether (e.g., ternary-weight linear layers with continuous activations [201, 165]), becoming harder and harder to distinguish from low-bit quantised RNNs. It is, therefore, worthwhile to ask how cortical computational principles can still contribute to scaling up SNNs in the age of massive foundation models and alleviate this tension. Realistically, a precondition for neuromorphic systems to catch up to the widespread adoption of current LLMs is to first match the Transformer’s favourable qualities. For the purposes of this study, as mentioned above, these target qualities include long-range sequential dependency modelling, amenability to fast parallel simulation, and a notion of "adaptability". The present study attempts to address all three properties, drawing on recent developments in understanding temporal processing in the auditory cortex. The brain performs computations using recurrent, non-linear dynamics [42, 161, 178], determined by a mix that includes, but is not limited to, individual neuronal dynamics [65], synaptic connectivity [124, 22], and neuromodulation [135]. The auditory cortex is no exception, as it essentially implements non-linear mappings from auditory sensory inputs to higher-order neural representations [108]. Crucially, however, recent evidence suggests that across both primary (i.e., acoustic processing) and non-primary (i.e., speech or music processing) areas, inputs are processed within almost fixed-duration temporal windows [151, 137]. In other words, the dynamical systems within the auditory cortex have a non-linear fading memory [21], that leaks information over time at a roughly constant, inherent rate, irrespective of the current state of neural activity (e.g., individual neuron membrane voltages) or the incoming inputs. This is not a universal property of non-linear dynamical systems, as one cannot generally predict trajectories in phase space a priori [168], including whether input information is forgotten at all. Traditional non-linear SNNs, much like RNNs, do not generally provide guarantees about time scales across phase space either [168, 163]. The linear sub-threshold behaviour in popular neuron models, such as the Leaky Integrate-and-Fire (LIF) model [65], can be parametrised to have pre-determined and constant memory properties in individual neurons (Section 2.1). Furthermore, one can also parametrise recurrent synaptic connectivity matrices to guarantee certain linear
2
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
memory properties [88]. Nevertheless, once non-linear spiking is introduced, the analytical guarantees granted by linear dynamics no longer apply, making direct initialisation and parametrisation of network timescales difficult [47]. This is a well-known shortcoming of non-linear recurrent architectures, and is the root cause of the recently mentioned vanishing and exploding gradient problems. If one were to further increase the complexity of the network architecture, for example, by adding neuromodulation (i.e., changing recurrent weights at inference based on context [4]), it is safe to assume that timescale parametrisation would become even more difficult. These challenges in reliably controlling the timescales of non-linear dynamics have fuelled a recent wave of interest in SNNs with purely linear recurrence, taking advantage of progress made in effective initialisation and parametrisation schemes for linear SSMs in deep learning literature [165, 131, 158]. However, as interest in linearity has increased in both neuromorphic and deep learning research, evidence has also been mounting for previously unknown advantages of non-linear recurrence. Recurrent computation has been shown to improve reasoning capabilities while using fewer parameters compared to traditional vanilla parallel blocks such as Transformers and SSMs [202, 103]. Non-linear RNNs layers have also been theoretically proven to model algorithmic state-tracking tasks and finite-state automata (e.g., tracking the state of a chess game), unlike Transformer or SSM layers [129, 172]. This is a crucial observation, because state-tracking is devised as a synthetic test for complex reasoning capabilities, which have become a major focal point in LLM research [128, 103]. Much like non-linear dynamics, synaptic neuromodulation has also attracted recent attention in the RNN and SNN literature. It typically takes the form of LSTM-like gating [92, 102], or subnetworks that control the gain scaling of recurrent weights [30, 4], and has been shown to potentially help emulate Transformer "adaptability" in both linear [71] and non-linear [132] recurrent architectures.
Figure 1: Trade-Offs in RNNs and SSMs This figure shows the four qualities required for state-of-the-art and versatile sequence modelling. Three of them are staples of Transformer performance: adaptability (ii), long-range sequence modelling (iii), and parallelisability (iv). The fourth, non-linear recurrence (i), is thought to enable performance unattainable by vanilla Transformers. A competitive RNN or SSM should possess all four. However, as the red arrows highlight, they are fundamentally at odds with each other. Contributions To summarise, the overarching goal of this work is to progress neuromorphic computing toward a viable energy-efficient alternative to current large-scale AI systems. That implicitly means directly competing with the Transformer. This can be measured by performance on the current taxonomy of sequence modelling tasks, probing four specific qualities needed to match and even outperform the Transformer (Fig. 1): (i) Non-linear recurrent dynamics are needed for reasoning, and their effectiveness is measured by state tracking ability. This is where architectures have a concrete lever to improve on the Transformer; (ii) Adaptability, measured by selectivity, is essential for in-context learning and language modelling more broadly; (iii) Long-range memory is necessary for processing long conversations, documents, or high-frequency/long duration signals in general. One cannot deploy a Transformer alternative if it cannot process the same context lengths; (iv) Parallelisability is crucial for GPU-optimised, fast, and scalable training. Any architecture that cannot scale to the same large-scale training datasets as the Transformer immediately faces a fundamental hurdle to widespread adoption. 3
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
The main challenge is that pairing together all four qualities requires addressing a series of inherent trade-offs. Non-linear dynamics (i) typically require sequential simulation, preventing parallelism (iv). Furthermore, non-linear recurrence (i) and adaptability (ii) usually hinder long-range modelling performance (iii), since one cannot directly use state-of-the-art stable initialisation and parametrisation schemes from linear time-invariant SSMs. Thus, a clear research gap remains in tackling all trade-offs simultaneously, especially for neuromorphic and energy-efficient architectures. To address this, this study provides the following contributions: 1. A novel recurrent architecture, the ADaptive with Prescriptive Timescales Network (ADPTNet), is proposed. ADPTNet draws inspiration from the auditory cortex to process inputs using input-invariant timescales while remaining adaptive through neuromodulation. The main building block of the proposed dynamics is the topological conjugation operation, which uses data-dependent similarity transforms constrained to the orthogonal manifold to modulate recurrent dynamics. This is the first study to extend linear Transformer and DeltaNet online update rules to Riemannian manifold optimisation. Furthermore, it is also the first to combine them with Dynamical Systems and Chaos Theory. Accordingly, this work provides theoretical guarantees on parametric control of the Lyapunov spectrum of ADPTNet dynamics as a function of its recurrent eigenspectrum. The theory is then validated empirically, showing how to use SSM initialisation and parametrisation techniques in non-linear ADPTNet without losing performance on long-context sequence modelling (iii). At the same time, the study also shows how the same topological conjugation is sufficiently non-linear (i) to enable state tracking, while also injecting selectivity in the network (ii) to the same level as established selective and recurrent architectures [40], tackling two of the three trade-offs mentioned above. 2. Two novel efficient parallelisation methods co-designed with ADPTNet and based on DEER [118] are proposed. Conv-DEER is a Jacobian-free DEER variation that achieves similar average convergence to Quasi-DEER [68]. It leverages ADPTNet’s constant, predictable timescales to enable parallelisation via convolution, for the first time for non-linear RNNs. Within each iteration, Conv-DEER adds no computational overhead beyond the forward pass and scan/convolution, and uses only the network’s static recurrent eigenspectrum. Forward-DEER is also proposed, reaching even closer convergence behaviour to Quasi-DEER, while still not computing any derivatives or Jacobians, relying solely on the forward dynamics of the network. 3. ADPTNet is used as a backbone for an SNN, yielding state-of-the-art accuracy on the Spiking Speech Commands classification task [31], outperforming current state-of-the-art methods by over a percentage point (83.56% ± 0.15).
2
Background
2.1
Spiking Neural Networks
As previously mentioned, the most common backbone for SNNs is the LIF neuron model [53, 65]. In continuous time, sub-threshold LIF membrane voltage u ∈ R dynamics are equivalent to an RC circuit (resistance R, capacitance C) with a time constant of τ = RC (Eq. 1) 1 . In practice, Euler discretisation is used for simulation (Eq. 2) and the resting membrane voltage (urest ) is set to 0, yielding an exponential moving average with a decay rate (β) parametrised by the step size (∆t) and τ (Eq. 4). As mentioned in Section 1, the memory horizon of neurons in the sub-threshold regime can be directly defined by β. However, once the membrane voltage crosses the firing threshold (θ), a spike is fired, and the voltage is reset (Eq. 5 and 3), anticipating the difficulties in predicting overall timescales (Section 1). In the context of building SNNs with vector states u ∈ RN , an additional recurrent feedback connection with trainable synaptic weights (Wrec ) can also be added (Eq. 6), sometimes referred to as Recurrent SNNs (RSNNs) [16]. Finally, because spiking (s) is a non-differentiable Heaviside step function, vanilla backpropagation cannot be directly applied. One could address this by applying a local, brain-inspired learning rule such as Spike-Timing Dependent Plasticity (STDP) [17]. However, the community has converged on surrogate gradients as a de facto standard owing to their performance being closest to that of standard backpropagation [136]. While the binary threshold function is applied in the forward pass, the derivative of a differentiable surrogate function σ (e.g., sigmoid, arctan, etc.) (Eq. 7) is applied backwards, allowing backpropagation to otherwise be applied as usual. τ
du(t) = −(u(t) − urest ) + I(t)R, τ = RC dt
(1)
u[t + 1] = βu[t] + (1 − β)I[t + 1]
(2)
1
Note that Equation 1 uses I to denote input current. This notation will later be overloaded to mean the identity matrix, as is convention in linear algebra.
4
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
u[t + 1] = u[t + 1] − θs[t] −∆t
β=e τ s[t] =
2.2
1, u[t] ≥ θ 0, u[t] < θ
(3) (4)
(5)
u[t + 1] = βu[t] + (1 − β)(I[t] + Wrec s[t]) − θs[t]
(6)
ds[t] dσ[t] ≈ du[t] du[t]
(7)
State Space Models
SSMs have a long history of usage in dynamical systems and control theory [32], along with well-established applications such as neural activity modelling [119]. Much like a vanilla RNN (Eq. 8), SSMs are generically a function of a recurrent state x and incoming input u (Eq. 9), with output y typically obtained using a readout matrix and an often left-out skip connection (Eq. 10). However, for the purposes of this work, the SSMs considered do not include a non-linearity (σ) in the recurrence. Depending on whether model parameters also vary over time, SSMs can be Linear Time-Varying (LTV) or Linear Time-Invariant (LTI) (Eq. 11). f (xt+1 ) = σ(Wrec xt + Win ut+1 )
(8)
xt+1 = At+1 xt + Bt+1 ut+1
(9)
yt+1 = Ct+1 xt+1 + Dt+1 ut+1
(10)
xt+1 = Axt + But+1 yt+1 = Cxt+1 + Dut+1
(11)
In the context of deep learning, SSMs were first introduced in the Legendre Memory Unit (LMU) [179], in which a linear memory network is coupled to the main RNN. Because of its linearity, the memory unit can be initialised to project the input signal onto a Legendre orthogonal polynomial basis, thereby granting theoretical guarantees for compressing a fixed-width sliding window of the past. Building on the LMU, the S4 architecture [74] relinquished the non-linear RNN in favour of a purely linear recurrent unit that also extended the initialisation schemes to other orthogonal polynomial bases (i.e., Laguerre, Chebyshev, etc.), and compression strategies beyond fixed windows, such as exponentially decaying [77]. S4 was a landmark achievement for SSMs, as it was the first to be proven to outperform Transformers on long-range dependency modelling tasks. Follow-up studies have established that the exponentially decaying memory inductive bias is sufficient to achieve this performance [117, 122]. Accordingly, the current consensus is that A is typically a complex-valued diagonal matrix, with its eigenvalues distributed close to the unit circle [76, 140]. One may notice a strong resemblance to sub-threshold LIF dynamics (Section 2.1), or their complex-valued variation, Resonate-and-Fire (RF) neurons [139]. Similar to LIF neurons, SSMs are typically first formulated in continuous time and then discretised for simulation. For notational simplicity, the SSM descriptions included here are in discrete time. All SSMs described so far in this section have been LTI architectures. However, from the perspective of Language Models, LTI SSMs are known to lag behind Transformers [145, 59] and lack the latter’s adaptability (Section 1). To account for this, LTV SSMs have also been developed, first introduced as Mamba [71]. Here, At , Bt , Ct are functions of the input ut , typically the output of a Multi-layer Perceptron (MLP), and thus do not create non-linear dependencies between time steps. Mamba has made progress in closing the gap with Transformers for language modelling and has been widely adopted in hybrid architectures alongside them [20]. Mamba’s success on language tasks has, however, come at the cost of long-range dependency modelling performance, losing the theoretical guarantees of earlier LTI SSMs [197]. One may also notice that the input-dependent parameters in Mamba are reminiscent of neuromodulated RNNs in 5
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
neuromorphic literature [30, 4], showcasing the apparent trade-off between adaptability and long-range dependencies referenced in Section 1. A key to the success of SSMs has been the GPU-parallelisable nature of linear recurrences for training, while retaining efficient O(1) sequential deployment. The output of an LTI system (y) is equivalent to the convolution of an input signal with a global sequence-long kernel defined by the powers of the recurrent matrix [27, 75] (Eq. 12). Furthermore, these convolutions can be efficiently computed using element-wise multiplication in the Fourier domain, a near-linear operation in sequence length (O(N log(N ))) compared to the Transformer’s quadratic scaling, when using Fast Fourier Transforms (FFTs) [29]. Mamba-like LTV architectures are not equivalent to convolutions, and instead rely on parallel associative scans for parallelisation [19]. (12)
κ = (CB, CAB, CA2 B, ..., CAL B)
(13)
y = u ⊛ κ = F F T −1 (F F T (u) ∗ F F T (κ))
(14)
Lyapunov Exponents Unstable LTI Vector Field
Unstable LTI Vector Field
4
Unstable Critical Point
3
3
2
2
2
1
1
1
0
0
0
1
1
1
2
2
2
3
3
4
4
3
2
1
0 x
1
2
3
(a) Unstable LTI System
4
4
Non-Linear System Vector Field
4
Stable Critical Points
3
y
y
4
y
2.3
y[t] = CAt Bu[0] + ... + CABu[t − 1] + CBu[t]
Critical Points
B
A
3 4
3
2
1
0 x
1
2
(b) Stable LTI System
3
4
4
4
3
2
1
0 x
1
2
3
4
(c) Non-Linear System
Figure 2: Difference in predictability between LTI and Non-Linear Systems. Subfigures 2a and 2b show the vector fields of unstable/stable 2-dimensional LTI systems (ẋ = A11 x + A12 y, ẏ = A21 x + A22 y), where A has positive and negative eigenvalues, respectively. Subfigure 2c shows the vector field of a non-linear dynamical system: ẋ = −2x − 3xy, ẏ = 3y − y 2 [87]. The system has two critical points: A is a stable point (sink) that attracts trajectories in its vicinity, while B is a saddle point that attracts and repels trajectories depending on their location. As highlighted in Sections 2.1 and 2.2, the long-term dynamics of LTI systems are fully determined by the eigenvalues of recurrent matrices. In continuous time, if all the eigenvalues have negative real parts, then the system exponentially relaxes back to its resting state in the absence of an external driver (Figure 2b). Conversely, if there are eigenvalues with positive real parts, inputs may persist indefinitely in memory, causing the system to diverge exponentially from its resting state over time (Figure 2a). Alternatively, in discrete settings, stability is tied to eigenspectra bounded within the unit circle. The eigenvalues also describe the rate of convergence/divergence from the resting state. LTV and non-linear systems do not generally benefit from a priori predictors of long-term behaviour. Consider the phase plane in Figure 2c. Depending on the state of the system, trajectories may either converge to point B or converge/diverge around the saddle point A. Moreover, the rate of convergence/divergence also differs by region. While one can approximate locally with Jacobian linearisation in the vicinity of critical points, long-term trajectories cannot be predicted. Instead, one can describe the behaviour of already observed trajectories using Lyapunov exponents. For a given discrete non-linear system xt+1 = f (xt ), Lyapunov exponents (λk ) are the averaged singular values (σk ) of a system’s long-term Jacobian (G) (Eq. 15) for a sampled trajectory (Eq. 16) [64, 47]. In particular, the Largest Lyapunov Exponent (LLE or λ1 ) indicates whether the system behaviour is stable or chaotic. First, if the LLE is negative, a trajectory starting from a perturbed state xt + ∆ eventually converges with its unperturbed counterpart. The rate of convergence is lower the closer the LLE is to zero, i.e., the system has a longer memory horizon. Second, if the LLE is positive, perturbations result in exponentially diverging trajectories and chaotic dynamics. 6
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
G1:N =
1:N Y
Jt , Jt =
t
λk = limN →∞
dxt dxt−1
(15)
1 log(σ(GN )k ) N
(16)
Lyapunov exponents have been used not only to formalise vanishing and exploding gradient problems, but also to mitigate them. Most notably, Gradient Flossing [47, 49] has been proposed as a method to directly control network timescales through optimisation over the Lyapunov spectrum. For example, one can use the LLE magnitude as a loss function to encourage a slowly decaying fading memory, as sought after in SSMs, and use it in conjunction with traditional Backpropagation-Through-Time (BPTT) [186]. However, separate Gradient Flossing epochs are typically required, and computing Jacobians and their singular values is computationally and memory-intensive, slowing training and reducing scalability. Moreover, Lyapunov exponents are descriptors of particular trajectories, which implies that input sequences outside the training set may still produce undesirable dynamics. Finally, as network complexity increases and sequences grow longer, it becomes increasingly difficult to train for precisely targeted Lyapunov exponents, which also tend to drift during training anyway. The solution proposed here aims in part to provide a more scalable and reliable alternative to Gradient Flossing. 2.4
Parallel Simulation Algorithms
As highlighted in Section 2.3, non-linear dynamics cannot be predicted as easily as linear time-invariant systems. Consequently, computing outputs from non-linear RNNs and SNNs traditionally requires step-by-step sequential simulation, which cannot fully leverage GPUs and thus limits scalability compared to parallel architectures such as SSMs and Transformers [104, 140]. An emerging solution to bridging this gap has been to reframe RNN simulation through Newton’s Fixed-Point Method, iteratively refining the entire trajectory in a single step [118, 68, 33]. Given a generic discrete non-linear recurrence rule (Eq. 17) and a state-guess tensor for all time steps (s ∈ RT ×N ), one can define a residual tensor (Eq. 18) which becomes 0 if and only if the state-guess is the target trajectory. Each Newton iteration minimises the Euclidean norm of the residual (Eq. 19), by finding the optimal update step (Eq. 20). The resulting sequence-wide iteration step (Eq. 21) is now known as the DEER algorithm [118]. While similar Newton-iteration-based parallel-in-time simulation methods have been discussed for decades [61], DEER is notable for including a GPU-efficient way to compute the large Jacobian and its product with the residual required for the state update (Eq. 20). Because the residual Jacobian matrix J is block lower bidiagonal, its inverse is populated, using the same terminology from Eq. 15, by long-term Jacobians (Eq. 22), and its product with the residual matrix is equivalent to an associative parallel scan. At any Newton iteration k, DEER is guaranteed to have converged on the first k time steps of the simulation [68]. Therefore, it takes at most as many steps as a step-by-step simulation. However, since DEER iterations are considerably more computationally intensive than applying a single RNN step, speed-ups are only possible if DEER converges in substantially fewer steps. As hinted by the presence of long-term Jacobians Gi:j in J , there is a direct connection between the LLE of the recurrent systems simulated (Section 2.3) and the convergence rate of DEER [69]. For instance, positive LLEs, i.e., chaotic dynamics, typically cannot be effectively parallelised with DEER. Intuitively, errors in earlier steps propagate across time and can affect all subsequent states, increasing the number of fixed-point iterations and approaching the total number of sequential steps. For the same reason, even with a negative LLE, the longer the system’s memory horizon, the more future states are affected, and hence, more iterations are required for convergence. xt+1 = f (xt ), xi ∈ RN
(17)
r(s) = [f (s0 ) − s1 , f (s1) − s2, ..., f (sT −1 ) − sT ], r(s) ∈ RT ×N
(18)
L(s) =
1 ||r(s)||2 2
∆s = min(∥r(s + h)∥22 ) = min(∥(r(s) + h
h
7
dr(s) 2 h∥2 ) = − ds
(19)
dr(s) ds
−1 r(s)
(20)
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
(i+1)
s
I −1 −J2 0 dr(s) J −1 = = . ds .. 0
(i)
=s
+ ∆s
0 I −J3
0 0 I .. . ...
0
(i)
(i)
=s
−
dr(s(i) ) ds(i)
−1
r(s(i) )
(21)
... ... ...
−1 I 0 0 G2:2 0 = G2:3 . .. .. .
0 I G2:2
0 0 I .. .
−JT
I
G2:T
G2:T −1
G2:T −2
... 0 . . . 0 . . . 0 .. . ...
(22)
I
DEER suffers from two main drawbacks. Computing and storing full Jacobians Jt for each time step t in each batch during training limits scalability. Furthermore, the long-term Jacobians Gi:j in J can become numerically unstable even if the system is asymptotically stable, depending on the presence of volatile transitory dynamics. To address the former, Gonzalez et al. [68] showed that Quasi-Newton methods, which replace the full Jt Jacobians with diagonal approximations, have similar convergence properties to the original DEER algorithm, reducing compute and memory overheads. To tackle numerical instability, the study also proposes a Kalman filter-based approach that dampens all Jt eigenvalues by a fixed factor, thereby preventing the Gi:j terms from exploding. This work aims to provide more efficient heuristics for both challenges (see Section 4.5). 2.5
Orthogonal Matrices
The Orthogonal Group (O(n)) comprises matrices Q with the property that QQT = I. From a neural network perspective, orthogonal matrices are of interest because, multiplied by any vector v ∈ Rn , they conserve Euclidean norms and angles. In other words, because they have eigenvalues restricted to {±1} and singular values in {1}, a system that multiplies an input vector by a sequence of orthogonal matrices (e.g., QT QT −1 ...Q1 v) has a memory that neither decays nor explodes. Such a dynamical system, where all recurrent Jacobians Jt are in On , would have a Lyapunov Spectrum (Section 2.3) equal to 0, a property referred to as criticality or the edge of chaos [50]. One could be tempted to argue that, for example, an SSM with an orthogonal recurrent matrix A (Eq. 9) would have favourable memory properties. However, since no information dissipates and new input signals u continue to arrive, the state could grow asymptotically, leading to numerical instabilities [140]. Therefore, orthogonal recurrent matrices have been mostly found in non-linear RNNs (Eq. 23) [6, 86, 188]. Because of the added element-wise non-linear function σ, the recurrent Jacobian of the network Jt becomes a product between a diagonal and orthogonal matrix (Eq. 24). If one uses a common activation function, such as a sigmoid, tanh or ReLU, the spectral norm (i.e., maximum singular value) of the diagonal derivative term Dt is bounded by 1. In turn, because the recurrent weight is orthogonal, the spectral norm of every recurrent Jt and, implicitly, long-term G1:T are also ≤ 1 (Eq. 25). Therefore, the Lyapunov spectrum is negative, and the system is stable with fading memory and no exploding gradients (Section 2.3 and Arjovsky et al. [6]). Here, it is worth emphasising that orthogonal RNNs cannot guarantee that all timescales of the network are slowly decaying, long-term memory. As Eq. 25 suggests, there is no lower bound on the norm of each Jt . For example, one could have all ReLU activations set to zero at once and thus all information forgotten. Orthogonal parametrisation does not solve such "dead neuron" pathologies [43]. Furthermore, the power of initialisation schemes for SSMs lies in setting the distribution of all recurrent eigenvalues near the unit circle, thereby enabling multiple long-term timescales [140]. Even if there are no time steps where Dt is 0, then the inequality in Eq. 25 only shows the upper bound of slow timescales being 1, but does not help with setting the rate of decay for all other timescales besides the slowest. Regularisation methods like Gradient Flossing [47] can be added to mitigate this, but they cannot offer any guarantees about the system’s specific timescales either (Section 2.3). The solution proposed in this work uses orthogonal matrices in a way that also guarantees a lower bound on the Lyapunov Spectrum. xt+1 = σ(Qxt + Win ut+1 )
Jt =
(23)
dσ Q = Dt Q dxt−1
(24)
1:T 1:T 1:T 1:T Y Y Y Y ∥G1:T ∥ = ∥ Jt ∥ = ∥ Dt Q∥ ≤ ∥Dt Q∥ = ∥Dt ∥ t
t
t
8
t
(25)
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
Training neural networks that use orthogonal parametrisations takes advantage of the fact that the On also carries a Riemannian manifold structure, which enables smooth optimisation [3]. Over the general linear group GL(n), one can define a dynamical system whose state M ∈ Rn×n evolves in the steepest descent direction that minimises a scalar objective function f (Eq. 26). Discretisation with a step size η yields the traditional gradient descent algorithm (Eq. 27). Importantly, the update step automatically produces a valid new matrix Mt in GL(n).
Figure 3: Riemannian Gradient Descent over the Orthogonal Manifold Given an initial orthonormal matrix Q and a Euclidean gradient ∇f (Q), Riemannian Gradient Descent is performed by first projecting the Euclidean gradient to the tangent space to O(n) at Q, i.e., TQ [3]. The updated Q from the tangent space is then projected back to the manifold using a retraction function R . That is not the case in Riemannian optimisation over orthogonal matrices, as taking a Euclidean gradient step could break the orthogonality constraint (Q − η∇f (Q))(Q − η∇f (Q))T ̸= I. To keep trajectories constrained to On , any point on it has to satisfy QQT = I, and therefore any infinitesimal direction of travel from that point has to respect Eq. 28. In practice, this is obtained using a projection onto the tangent space to On at Q (TQ On ) (Eq. 29), which is essentially a mapping to a skew-symmetric matrix (a matrix X with the property that X T + X = 0). Once a valid directional derivative is computed, one can integrate over it to obtain the final modified orthogonal matrix. In discrete gradient descent, this is done by projecting the updated matrix back onto the manifold using a retraction (R(Q, ψQ (∇f (Q))). The exact solution is the matrix exponential (Eq. 30) [3]. This is evidently computationally expensive, so it is often substituted with a Taylor approximation in the form of the Cayley map (Eq. 31) [1]. dM = −∇f (M ) dt
(26)
Mt = Mt−1 − η∇f (Mt−1 )
(27)
d(QQT ) dI dQ T = ⇒ Q +Q dt dt dt ψQ (∇f (Q)) =
dQ dt
T =0
(28)
1 (∇f (Q)QT − Q∇f (Q)T ) 2
(29)
Rexp (Q, η · ψQ (∇f (Q)) = exp(η · ψQ (∇f (Q))Q
(30)
η η Raprox (Q, η · ψQ (∇f (Q)) = (I − ψQ (∇f (Q))(I + ψQ (∇f (Q))Q 2 2
(31)
3
Related Work
3.1
Neuromodulation in Neural Networks
The term "neuromodulation" is used in this work to refer to neural architectures where parameters vary during training and inference as a function of the network input. The focus here is mostly on neuromodulated RNNs, where both the recurrent weights and the state of the network evolve as functions of the previous weights and state, and incoming inputs (Eq. 32). As previously highlighted in Sections 1 and 2.2, this typically takes the form of gating or hypernetworks [78]. 9
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
Wt+1 = f (Wt , xt , ut+1 ) xt+1 = g(Wt xt , ut+1 )
(32)
First popularised by LSTMs and GRUs [92, 28], gating is now considered any operation based on Hadamard products that scales activations in an input-dependent manner. As argued in Gu and Dao [71], the most basic example of it is the Gated Linear Unit activation [37] (Eq. 33), comprising two linear layers, where one also passes through a sigmoid function and element-wise scales the output of the other. In recurrent models, the goal of gating is most commonly to enable dynamic forgetting, in which the current state and/or input determine which information decays more quickly [102, 71]. Hypernetworks go a step further by generating all model weights using a smaller subnetwork, rather than performing only element-wise scaling. In their original conception, the sub-networks only took fixed qualities, such as depth, as input to produce parameters. However, more recently, they have also become data-controlled [26]. Modern SSMs have been blurring the distinctions between the two research directions, with recurrent weight matrices more and more commonly becoming diagonal, thereby effectively reducing hypernetworks down to forget gates [71, 201]. GLU (x) = (WA x + bA ) ⊙ σ(WB x + bB )
(33)
All forms of neuromodulation described so far are heuristic in nature: the goal is to explore the general effect of datacontrolled parameters on model performance, for example, on tasks requiring selective forgetting [101], or continual learning [180]. More recently, however, increased attention has been given to endowing neuromodulation with more principled structure, with the deep learning and neuromorphic research communities adopting slightly diverging goals and understandings. In mainstream deep learning, structured neuromodulation has been widely explored in associative recall (AR). AR tasks test a model’s ability to store key-value pairs in working memory and then retrieve the correct value when prompted with a query [7]. Because all keys, values, and queries are provided at inference time, the network is said to "learn" the associations "in-context". While content-addressable associative memory has been a popular subject of study since at least the advent of Hopfield Networks [95], this modern framing emphasises performing AR online and at scale. To generalise, LLMs have to be able to dynamically answer prompts such as "The apple is red. What colour is the apple?" without necessarily relying on prior knowledge. The Transformer [177] is explicitly formulated in these terms. Its core operation, self-attention, linearly projects the input sequence (X ∈ RT ×d ) to keys (K = XWk ), queries (Q = XWq ), and values (V = XWv ), where Wk , Wq , Wv ∈ Rd×dk . The product of the keys and queries yields the attention matrix, which then mixes all values across the sequence (Eq. 34). Because the attention matrix is a function of the input, some have argued that the Transformer is a hypernetwork [156]. Furthermore, given the explicit associative memory formulation, it can also be considered an example of structured neuromodulation, using the terminology of this study. QK T O = SelfAttention(Q, K, V ) = Softmax( √ )V dk
(34)
The QK T ∈ RT ×T product in self-attention (Eq. 34) incurs a quadratic computational and memory complexity in sequence length T , making efficient alternatives to Transformers highly desirable. Linearised Transformers [170, 144, 199] aim to reduce this overhead by not computing a T × T matrix and replacing the row-wise Softmax over the attention matrix with non-linearities over K and Q individually [107]. In the original Linear Transformer (Eq. 35) [107], for example, the attention matrix is a sum over dk × dk outer products of individual key-value pairs (kt , vt ). Softmax normalisation is only applied between keys, and the query is transformed using a generic element-wise non-linearity (e.g., sigmoid). These changes reduce the computational and memory requirements to linear scaling in sequence length. However, interest in linearised Transformers of this original form waned under pressure from developments in hardwareoptimised vanilla self-attention implementations [35] and the advent of SSMs (see Section 2.2). One could argue that this ebb ended with the landmark connection between linearised Transformers and the delta rule [187] in Schlag et al. [153]. The structure of the Linear Transformer, and related architectures, is amenable to a recurrent formulation (Eq. 36). However, unlike traditional RNNs, the recurrence is not over a vector state but rather over a so-called "fast weights" matrix W , which evolves over time during inference. As highlighted in Schlag et al. [153], without the Softmax normalisation in Eq. 36, this weight-dynamics view of Linear Transformers is remarkably similar to using online gradient descent over a recall error correction objective (Eq. 37). The difference is that while Linear Attention keeps adding key-value pairs to memory, the gradient descent step first retrieves and removes the current T 2 value associated with the key (via the Householder transformation I − kt+1 kt+1 ) before writing a new associated 2
The key kt+1 is typically normalised [193]
10
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
T value (kt+1 vt+1 ). The result is that only changes in the value vector vt+1 are stored, hence the name of the delta rule. This architecture, referred to as DeltaNet [193], is now an important pillar of efficient Transformer development, with significant research focused on In-Context Learning and Test-Time Training directly adopting its online optimisation framing [13]. Moreover, a priority has been to improve compression capabilities in theoretically grounded ways. For instance, Lattice [106] treats the columns of Wt as individual memory slots and only stores new information that is orthogonal to each of them. MesaNet [181] extends online optimisation from single gradient descent steps at each time step to a regression over all preceding time steps, producing an optimal readout matrix for the query.
ot = P
! t X ki (e ⊗ vi ) σ(qt )
1
t kt i=1 e
· σ(qt )
(35)
i=1
Wt+1 = Wt + ekt+1 ⊗ vt+1 zt+1 = zt + ekt+1 1 ot+1 = zt+1 ·σ(q Wt+1 σ(qt+1 ) t+1 ) η Wt+1 = Wt − ∇Wt (∥Wt kt+1 − vt+1 ∥2 ) 2 T T = Wt − η(Wt kt+1 kt+1 − kt+1 vt+1 )
(36)
(37)
T T = Wt (I − ηkt+1 kt+1 ) + ηkt+1 vt+1
It is worth emphasising that Linear Transformer and DeltaNet-inspired models do not fully fit the neuromodulated RNN template from Eq. 32. In these architectures, the output at each time step t is the readout product of Wt and qt (Eq. 36) and is only fed to the next layer, not the next time step. In other words, since the query only depends on the current input and is not recurrent, the weights Wt are the state of the system. This differs from a neuromodulated RNN, where recurrence also applies to a vector state, making Wt both a state and a modulator of the network’s dynamics. Because linearised Transformer recurrences rely mostly on addition (writing key-value outer products to memory) and, when present, multiplications are restricted to Householder reflections, the main concern is keeping the norm of Wt bounded. In contrast, a multiplicative interaction between recurrent states and weights in neuromodulated RNNs warrants careful consideration of long-term Jacobian stability when parametrising Wt (Section 2.3). In practical terms, this means that in a recurrent Transformer, the eigenvalues of Wt itself are typically irrelevant. In a neuromodulated RNN, however, they can play an important role in how well the model can perform (Section 2.5). This difference helps explain the approach neuromorphic research has taken regarding structured neuromodulation. In a significant body of neuroscience-inspired literature, structured synaptic neuromodulation has taken the form of constraining Wt in RNNs, rather than adopting an associative recall formalism as in deep learning research. For instance, relevant to this study, in Zador et al. [198], neuromodulated recurrent weights are constrained to linear interpolations of basis points on a rigid matrix manifold (e.g., a straight line or an ellipse). In Costacurta et al. [30], the recurrent weights Wt are constrained to rank-K matrices, where each low-rank component is scaled by an individual neuromodulatory signal. Although not within a single, comparable architecture, the fundamental concepts behind DeltaNet are also present in neuromorphic research. Evidently, associative learning algorithms such as the Delta and Hebbian rules originated in neuroscience [85, 105]. Moreover, SNNs benefit from additional biologically plausible online and local rules that can leverage temporal information, such as STDP [17] or e-prop [16], which have yet to see substantial adoption in mainstream deep learning. However, all these learning algorithms are not typically paired with Transformer-like Key-Value-Query projections and are generally intended as full backpropagation replacements rather than inference-time augmentations. In-context gradient descent is present to some extent in SNNs, though. For instance, Parameter-Free Attention [167] minimises, at inference, a linear-separability loss that implements lateral inhibition for input currents to SNN layers. 3.2
Parallelisable Architectures
As mentioned in Section 1, GPU-parallelisation is one of the core advantages that have fuelled the popularity of Transformers, and thus is a prerequisite for any efficient alternative to be competitive. Furthermore, SSMs have shown through equivalent parallel and iterative formulations that one need not trade off training and deployment efficiency (Section 2.2). In fact, parallel/recurrent duality has emerged as an essential characteristic of state-of-the-art efficient Transformer contenders. 11
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
Both Linear Transformers and DeltaNet variants require only linear operations between time steps, so, unsuspectingly, they also exhibit parallel/sequential duality (Section 3.1). However, naive full parallelisation can lead to similar quadratic-scaling issues to those of traditional self-attention. Recurrence rules such as Eq. 36 or 37 assume causality, and therefore cannot be directly mapped to an efficient matrix multiplication with linear scaling in sequence length (e.g., Q(K T V ), K T V ∈ Rdk ×dk ), because a binary causal mask M is required (Eq. 38) [192]. Instead, one can hedge between parallel and sequential forms through chunk-wise parallel simulation [193, 194]. Each segment of length C is processed in parallel at a quadratic in-chunk-size cost, with the final state WC∗k of each chunk k then recurrently fed into the next as its initial state. However, the overall complexity across the entire sequence length becomes sub-quadratic. O = (QK T ⊙ Mcausal )V
(38)
As highlighted in Section 2.4, there has also been a growing interest in parallelising non-linear dynamics. Much like the aforementioned DEER algorithm [118] and its variations [68, 33], many parallelisation methods for non-linear systems typically trade fully sequential, step-by-step execution for large, whole-sequence fixed-point iterations. Notably, Schöne et al. [155] introduces implicit SSMs, which iteratively apply an SSM layer to the entire sequence until convergence to a fixed point. While implicit SSMs have been shown to simulate certain non-linear dynamics, unlike DEER-based algorithms, they are not 1:1 interpretable equivalents to arbitrary non-linear RNNs. In implicit SSMs, non-linear dynamics are intrinsically encoded in fixed-point iterations, even if the simulation is sequential, i.e., each sequential simulation step still requires multiple fixed-point iterations for itself. While the ultimate goal of SNNs is generally considered deployment on specialised low-power hardware, the most popular training paradigm remains backpropagation (see Section 2.1). Hence, to improve scalability, GPU parallelism at training time has also gained interest within the neuromorphic community. The most direct, and perhaps most popular, strategy to achieve this has been to build SNNs based on mainstream parallel architectures. First introduced in Stan and Rhodes [165], SSM-based SNNs can be trained using the same convolutional approach as standard SSMs (see Section 2.2) [158, 10]. Analogously, spiking counterparts to vanilla [200] and linearised [195] Transformers have also been proposed, similarly inheriting parallelism from the baseline architectures. However, a direct spiking adaptation of DeltaNet has not yet been studied to the best of the authors’ knowledge. DEER-like algorithms have also found counterparts in SNN literature. As the name suggests, Fixed-Point Parallel Training (FPT) [57] follows the same principles as DEER (Section 2.4). However, unlike DEER, FPT does not require computing Jacobians J or long-term Jacobians G in J (see Eq. 22). Instead, FPT takes advantage of LIF sub-threshold dynamics, and propagates spikes "in the future" using powers of the membrane decay constant β (Eq. 39). To handle the spiking non-linearities, where DEER would compute Jt , FPT simply propagates the spikes themselves, or, during training (for numerical stability), a smooth approximation of the Heaviside step function s(u − θ) (e.g., sigmoid). Therefore, instead of using a residual tensor for the fixed-point iteration, FPT directly treats all membrane voltages across time u ∈ RT ×d as weighted sums of previous output spikes s(u − θ) and input currents i (Eq. 40). While not computing any Jacobians reduces memory and computational overheads, it makes the critical assumption that the network has no recurrent weights (see Eq. 6), and thus that all neurons are independent. This effectively limits expressivity compared to vanilla RNNs. 1 β β2 B= . ..
β T −1
0 1 β
0 0 1 .. .
... 0 . . . 0 . . . 0 .. .
β T −2
β T −3
...
(39)
1
u1 = Bi u2 = −θ(B − I)s(u1 − θ) + Bi .. . uk = −θ(B − I)s(uk−1 − θ) + Bi
(40)
It is worth noting that SNNs also benefit from an additional, idiosyncratic simulation paradigm. Event-based simulations reduce execution time by computing between-spike sub-threshold membrane evolution using closed-form solutions, thereby leveraging spiking sparsity to skip simulation steps [190]. Moreover, they rely on producing sequences of exact spike timings in continuous time rather than requiring temporal discretisation [48]. Event-based methods have also been 12
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
combined with DEER-like Newton fixed-point iterations to produce all spikes in parallel, providing further speed-ups [133]. Because only spike times are computed in event-based simulation, training uses these continuous time points rather than surrogate gradients to account for spike non-differentiability during backpropagation [190, 130]. However, in practice, this also means that backpropagation cannot effectively control the creation of new spikes, only shift existing ones forwards/backwards in time. Spikes can be deleted by moving them outside the task’s time window, thereby causing neurons to become silent. Once silent, neurons will remain silent (i.e., dead neurons) because non-spiking sub-threshold intervals do not receive any learning signals. [53, 185, 111]. Therefore, while event-based simulation holds tremendous potential, surrogate gradient methods are still generally preferred for achieving state-of-the-art accuracy [53]. In this context, event-based discretisation could be considered a middle ground. Traditional discretisation schemes, such as zero-order hold (ZOH) or Euler, generally employ fixed step sizes, imposing regular sampling on the data that can be processed by systems parametrised with them, such as SSMs. In contrast, neuromorphic sensors such as retina-inspired Dynamic Vision Sensor (DVS) Cameras [5] produce irregular, asynchronous, and sparse events at high temporal resolution. A frame-based approach with regular sampling, such as a traditional SSM, would waste resources processing a large number of empty input frames. Similar to event-based simulation, event-based discretisation can be applied to SSMs to reduce computational overhead by evolving their linear dynamics in closed form between input spikes [154]. However, these event-based SSMs operate only on precise input spike times, not on internal network spikes, unlike full event-based simulation. In other words, this method would not be used to speed up SNN execution in general. 3.3
Long Range Modelling in SNNs
Learning dependencies over long temporal horizons is difficult in non-linear recurrent architectures, owing to the inherent unpredictability of dynamical systems (see Sections 2.3 and 2.5). Being non-linear RNNs themselves (Section 2.1), LIF-based SNNs are no exception. Accordingly, mirroring the progress from vanilla RNNs to LSTMs, Long Short-Term Memory Recurrent SNNs (LSNNs) were a significant milestone in tackling long sequences [15, 196]. Not to be confused with the notion of adaptability from Section 1, LSNNs use adaptive LIF neurons (ALIF), where the firing threshold temporarily rises θ[t] after each spike emission (Eq. 41). Evolving with a slower timescale α than the membrane voltage decay rate β, the threshold becomes a form of linear long-term memory. Because higher thresholds increase firing sparsity, the neurons’ silence patterns implicitly encode their memory [152]. Bittar and Garner [18] further generalised adaptive LIF (adLIF) neurons to slowly and linearly evolving currents w instead of thresholds (Eq. 42). One may notice that both ALIF and adLIF could be considered ontological precursors to the LMU [179], being non-linear RNNs with auxiliary linear memory units, in this case using heuristic parametrisation schemes. u[t + 1] = βu[t] + (1 − β)(I[t + 1] + Wrec s[t]) − θ[t + 1]s[t] θ[t + 1] = αθ[ t] + (1 − α)s[t]
(41)
u[t + 1] = βu[t] + (1 − β)(I[t + 1] + w[t]) − θs[t] w[t + 1] = αw[t] + (1 − α)(u[t] + s[t])
(42)
Another dominant approach to increasing long-range processing capabilities in SNNs has been the use of axonal delays. In biological neural circuits, depending on factors including the cell type, conductance and size of the axon, action potentials can spend anywhere between below 1ms to over 100ms travelling to their post-synaptic destinations [183]. These heterogeneous delays have inspired computational theories regarding the expressive power of the spike patterns they produce [100, 99]. Continuous-time delayed dynamical systems have even been theoretically proven to have infinite-dimensional state spaces [52]. For the purposes of discretised SNNs, however, their main contribution to long-sequence processing is avoiding vanishing and exploding gradients. In very deep neural networks, one can add skip connections between layers (Outk (x) = Layerk (x) + x) to mechanistically enable better gradient propagation across depth [84]. In essence, axonal delays achieve a similar effect over the unrolled temporal computation of SNNs, adding skip connections between time points τ steps apart (Eq. 43). At a high level, long-term Jacobians G1:t (Section 2.3) ds contain terms of the general form Wrec du with exponents at most τt . Therefore, delays reduce their vanishing/exploding exponents by a constant factor, regardless of the Wrec initialisation and parametrisation [147]. That means that in the absence of additional gradient controls, delays must grow proportionally with the input sequence length to keep up the same long-term Jacobian properties. Therefore, because τ -long delays require equally large state buffers, axonal delays are difficult to scale for state-of-the-art long sequence tasks (e.g., over 10,000 steps). u[t + 1] = βu[t] + (1 − β)(I[t + 1] + Wrec s[t − τ ]) − θs[t] 13
(43)
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
Finally, leveraging modern SSM methods has also improved SNN performance on long sequences. As mentioned in Section 2.2, SSMs are essentially leaky integrators akin to LIF sub-threshold dynamics. This shared computational primitive has led to the development of SSM-based SNNs, first comprehensively explored in Stan and Rhodes [165], which have shown that linear SNNs with SSM initialisation and parametrisation schemes can also retain strong longrange modelling performance. Follow-up studies have also integrated SSM concepts into non-linear SNN architectures such as adLIF by reintroducing resets after spiking [54]. While these approaches benefit from increased sparsity from refractory periods and thus improved efficiency, recurrent non-linearities also mean that the linear memory properties of SSMs are no longer guaranteed, potentially diminishing their scalability with longer sequences.
4
Methods
The goal of ADPTNet is to possess all three of the qualities highlighted in Section 1 as key to the Transformer’s success: (a) adaptability, (b) long-range dependency modelling, and (c) parallelisability. This section details how the topological conjugate backbone of ADPTNet serves to attain them.
4.1
ADPTNet Recurrence Rule
dx In general, two dynamical systems with recurrence functions dx dt = f (x) and dt = g(x) are said to be topologically conjugate if there exists a bijection h(x) such that g(x) can be formulated as the composition g(x) = (h−1 ◦ f ◦ h)(x). dx Relevant here, given two linear systems dx dt = Ax and dt = Bx, if there exists a non-singular matrix C such that −1 B = C AC, then the systems are also said to be topologically conjugate. As linear topological conjugacy equates to similarity transforms, the two systems share intrinsic qualities such as eigenspectra but differ in their eigenvectors/modes (Fig. 4a). In this context, ADPTNet employs topological conjugates by chaining them together to change flow direction locally while also retaining globally coherent properties (Fig. 4b). Furthermore, the similarity transforms are parametrised to be a function of the system’s recurrent state, i.e. C −1 (x)AC(x), matching the neuromodulation template set out in Section 3.1. Considering a discrete-time setting, an RNN-like step function can be derived as Eq. 45.
−1 C AC1 x, t ≤ T1 dx 1−1 = C2 AC2 x, T1 < t < T2 −1 dt C3 AC3 x, T2 ≤ t
(44)
xt+1 = C −1 (xt )AC(xt )xt + But+1
(45)
14
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
Trajectories Comparison: A vs Topological Conjugate of A A Topological Conjugate of A Initial Point Origin
0.8 0.6 0.4
Stitching Together Topological Conjugates
600
1.5
500
1.0
400
0.5
Re(x1)
Time Step
0.2 x1
300 0.0
0.4 0.6 0.4
0.2
0.0
x0
0.2
0.4
0.6
0.8
0.0
1.0
100
0.6
Initial State Transition Between Conjugates Origin
0.5
200
0.2
C1 1 AC1 C2 1 AC2 C3 1 AC3
1.5
0
1.5
(a) Linear Topological Conjugates
1.0
0.5
0.0
Re(x0)
0.5
1.0
1.5
(b) Switching between Topological Conjugates
Figure 4: Topological Conjugates in Linear Dynamics. Subfigure 4a shows trajectories from two linear dynamical dx −1 systems, one governed by dx AC. Notably, they dt = Ax and the other by a topologically conjugate equation dt = C share the same oscillation frequency and decay back to the origin at the same rate, as they share the same eigenvalues. However, the direction of travel, i.e., eigenmodes, differs. Subfigure 4b shows a toy example of how one can chain together (or "stitch") topological conjugates, changing eigenvector direction every 450ms, while keeping certain dynamical properties such as oscillation frequency constant (Eq. 44). Any potential implementation fitting this proposed parametrisation needs to be computationally inexpensive and effectively enable direct, fine-grained control over global dynamics, i.e., guarantee long-range memory. By definition, controlling global memory properties is a matter of parametrising A. The most computationally cheap and interpretable solution is to constrain A to be a diagonal matrix containing the eigenvalues of the system (Λ). This also allows taking advantage of state-of-the-art diagonal SSM parametrisation, with the details of the specific methods employed here described in Section 4.2. The goal of defining C as a function of the recurrent state is to enable adaptability in the system’s dynamics. Therefore, one way to obtain C from x would be to follow current state-of-the-art practices in efficient Transformer alternatives (see Section 3.1), and base it around outer products of key-value projections (k ⊗ v). Concretely, taking inspiration from the Linear Transformer [107], C can statefully evolve over time as in Eq. 46. It is important to highlight that here k and v are functions of the recurrent state x, rather than incoming inputs as in Linear Transformers. While appealing, this stateful and unconstrained C approach poses several problems. Cnew = Cold + k ⊗ v, k = fK (x), v = fV (x)
(46)
A first drawback is that, without careful consideration, accumulating key-value pairs in C poses inherent challenges when computing its inverse C −1 . First, generic matrix inversion for GL(n) has a high computational cost, with O(n3 ) complexity. Moreover, matrix conditioning plays a significant role in the numerical stability of inversion, effectively constraining the floating-point precision admissible by the algorithm if not accounted for [181]. The solution proposed here is to constrain C to the Orthogonal Group O(n), which enables efficient and stable inversion by transposing. As detailed in Section 2.5, constraining matrices that evolve over time on the orthogonal manifold requires Riemannian Gradient Descent. In this case, the gradient ∇f (Q) 3 can be considered the k-v outer products, and thus the update rule for Q becomes Eq. 47. Qnew = Qold exp(S), S = kv T QT − Qvk T 3
To align with notational convention for orthogonal matrices, C is replaced from here on with Q
15
(47)
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
new A second drawback with the premise of stateful Q evolution as in Eq. 47, is that it necessitates large dQ dQold Jacobians for backpropagation. Describing the credit assignment between all n × n elements in Qnew with respect to all n × n entries in Qold imposes an infeasible O(n4 ) memory cost. In general, Linear Transformers can avoid this drawback since the recurrent weight matrix does not have to be materialised sequentially (see Section 3.1). Here, however, Q Riemannian evolution is a non-linear function and thus requires full instantiation. In the interest of scalability and keeping network memory capacity transparently and entirely tied to the recurrent state x, for the purposes of this work, the memory of Q itself is removed. From a theoretical perspective, this is equivalent to taking a Hebbian-like associative key-value step on the orthogonal manifold from the identity I at each recurrent step of the network (Eq. 48).
Q = Iexp(kv T I − I T vk T ) = exp(kv T − vk T ), I = eye(n)
(48)
Taking all into consideration, the ADPTNet recurrence rule amounts to Eq. 49 and 50. To implement the matrix exp function required for the retraction back to the orthogonal manifold, the most widely used method is the Cayley Map as in Eq. 514 and 52 5 , where S is the skew-symmetric map applied to the kv T outer product, and k and v are normalised linear projections of the state x (Eq. 53). The β parameter in Eq. 51 is effectively the discretisation step size for travelling over the orthogonal manifold in the direction pointed by S. Because in this case the starting point is always the identity I, β gains the additional function of directly tuning the off-diagonal interactions between the elements of x within the dynamics of the network (Fig. 5).
Effect of on Off-Diagonal Interactions = 0.1
=1
= 10 0.8 0.6 0.4 0.2 0.0 0.2 0.4
0.8 0.6 0.4 0.2 0.0
0.8 0.6 0.4 0.2 0.0 0.2 0.4 0.6
Figure 5: Effect of β on off-diagonal interactions Intuitively, as one takes a larger step away from the identity I on the orthogonal manifold (Eq. 48), the resulting matrix becomes less and less dominated by its main diagonal, strengthening interactions between the neurons in the network recurrent state x.
Mt = Qt ΛQTt , Qt ∈ O(n), Λ = diag(λ1 , . . . , λn )
(49)
xt+1 = Mt xt + But+1
(50)
Q = (I −
β −1 β S) (I + S) 2 2
(51)
QT = (I −
β β S)(I + S)−1 2 2
(52)
S = kv T − vk T , k =
Kx Vx ,v = ||Kx|| ||V x||
(53)
Having established the core mechanics of the network, it is important to emphasise how they differ from traditional RNNs. For simplicity, the example of a vanilla RNN with ReLU activations is used here, namely xt+1 = ReLU(Wrec xt ) + Win ut+1 . The baseline logic is that the recurrent weight matrix Wrec stretches and rotates the state x, which then determines the outputs of the ReLU activations. The activations, in turn, form a diagonal D that element-wise scales entries in x by ∈ {0, 1}, i.e. xt+1 = D(Wrec xt ) ⊙ (Wrec xt ) + Win u. The Lyapunov Spectrum (Section 2.3) of the 4 5
The notation for skew-symmetric matrices, S, becomes overloaded in this work, since it may also refer to a spiking activation. From now it is assumed that M , Q, S, k, and v are always functions of x, and thus the notation is simplified to omit this detail.
16
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
network cannot be known a priori because D entries change arbitrarily. In contrast, ADPTNet fixes and parametrises the entries of the diagonal scaling factors. To obtain non-linearity, the state x is rotated by QT to decay along different modes with the pre-determined decay rates. This distinction enables the claims in Lemma 1 and 2 to be made. Lemma 1. For any given initialisation of Λ where all 0 < λi < 1, the singular values (σi ) of the recurrent Jacobian matrix J of the ADPTNet recurrence lie within the bounds λi ± 8βλmax (Assuming K, V are orthogonal). Proof. As per the product rule, the recurrent Jacobian matrix J comprises the matrix M itself and its derivative with respect to the state dM dx (Eq. 54). The eigenspectrum of M is known by definition, Λ, and because by construction M is symmetric positive definite, they are also M ’s singular values (λi = σi (M )). If the norm of the second term is sufficiently small, it is then possible to use Weyl’s inequality to find bounds on the singular value spectrum of their sum, J (Eq. 55). J =M+
dM x dx
|σi (J) − σi (M )| <
dM x dx 2
(54)
(55)
To find bounds on the norm of dM dx x, it is necessary to reduce it with the chain rule to terms with known norms. Using the product rule again, dM x further breaks down into two terms for each of the two sides of the similarity transform dx (Eq. 56). For the purposes of this proof, it is assumed that Q is parametrised using the Cayley Map (Eq. 51), and thus dQ dx and its transpose take the form of Eq. 57 and 58. When chaining them together within derivatives of the similarity dM transform in dM dx x (Eq. 59), it is worth observing that the β hyperparameter directly scales the norm of dx x. dQ dQ dM = ΛQT + QΛ dx dx dx
(56)
dQ β β dS = β(I − S)−1 (I − S)−1 dx 2 dx 2
(57)
dQT β β dS = −β(I + S)−1 (I + S)−1 dx 2 dx 2
(58)
dQ β dS β ΛQT = β(I − S)−1 (I − S)−1 ΛQT dx 2 dx 2 β −1 dS β −1 β β = β(I − S) (I − S) Λ(I − S)(I + S)−1 2 dx 2 2 2
(59) (60)
To simplify the notation, the A and B 6 substitutions defined in Eq. 61 and 62 are introduced. β −1 S) 2
(61)
β −1 β S) Λ(I − S) 2 2
(62)
A = (I −
B = (I −
Using the shorthand notation, the similarity transform derivative terms can be contracted to Eq. 63 and 64, and thus the overall dM dx x term reduces to Eq. 65. dQ dS ΛQT = βA BAT dx dx 6
(63)
While the symbol B is overloaded because it is also used for the input weights in the recurrence formula for ADPTNet (Eq. 50), in the scope of this proof it is only used as the shorthand from Eq. 62.
17
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
dQT dS T = βAB T A dx dx
(64)
dS dM dS x = βA( Ba − B T a), a = AT x dx dx dx
(65)
QΛ
Next, going down the derivation chain, the derivative of the skew-symmetric map dS dx requires a tensor outer product ⊗outer (Eq. 66) if computed on its own, since it is the Jacobian matrix of a matrix with respect to a vector. However, the tensors collapse back to matrices when multiplied with the rest of the terms in dM dx x (Eq. 69 and 70), avoiding the need to explicitly compute their individual norms. dk dv dv dk dS = ⊗outer v T + k ⊗outer − ⊗outer k T − v ⊗outer dx dx dx dx dx
(66)
dk 1 = (I − kk T )K dx ||Kx||
(67)
dv 1 = (I − vv T )V dx ||V x||
(68)
dk T dv dv T dk dS Ba = (v Ba) + (kaT )B T − (k Ba) − (vaT )B T dx dx dx dx dx
(69)
dv dk dS dk dv a = B T (v T a) + B T (kaT ) − B T (k T a) − B T (vaT ) dx dx dx dx dx
(70)
BT
Given that ∥B∥2 = B T ∥a∥2 ≤ ∥A∥2 ∥x∥2 then:
2
= λmax , ∥A∥2 =
AT
2
= 1,
dk dx 2
=
dv dx 2
1 , ∥v∥2 = ∥k∥2 = 1 and = ∥x∥ 2
dk dM dv x ≤ β ∥A∥2 ∥B∥2 ∥a∥2 (4 ∥v∥2 + 4 ∥k∥2 ) dx dx 2 dx 2 2 4 4 ≤ βλmax ∥A∥2 ∥x∥2 ( + ) ∥x∥2 ∥x∥2 = 8βλmax
(71)
Since the norm of dM dx x is upper-bounded by 8βλmax , for sufficiently small β, each of the singular values of J is bounded within: max(0, λi − 8βλmax ) ≤ σi (J) ≤ λi + 8βλmax , ∀i ∈ [1, . . . , n]
To simplify analysis going forward, it is assumed that λmin , λmax , and β are constrained such that λmin − 8βλmax > 0. Having established spectral bounds for individual ADPTNet recurrent Jacobian matrices J, it is now possible to formalise theoretical guarantees on long-term behaviour: Corollary 1. The Lipschitz constant of the ADPTNet recurrent step function is λmax + 8βλmax , for x ∈ Rn \ {0}n . Proof. Since Lemma 1 derives an upper bound on the norm of the recurrent Jacobian ∥J∥2 , this also produces, by definition, the Lipschitz constant L = sup ∥J∥2 = M + dM dx x 2 = λmax + 8βλmax . Lemma 2. For sufficiently many trajectory time steps T , the entire Lyapunov spectrum of ADPTNet dynamics is bounded within [ln(min(λmin − 8βλmax )), ln(λmax + 8βλmax )]. 18
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
Proof. In Lemma 1, it is established that all ADPTNet recurrent Jacobian J singular values σi (J) lie within ±8βλmax of the eigenvalues λi . Therefore λmin − 8βλmax ≤ σi (J) ≤ λmax + 8βλmax , ∀i ∈ [1, . . . , n] For notational simplicity, let σ̃min = λmin − 8βλmax and σ̃max = λmax + 8βλmax denote the limits above. As with any arbitrary matrices, for any two Jacobian matrices J1 and J2 from time steps 1 and 2, all the singular values of their product are bounded: σ̃min ∗ σ̃min ≤ σi (J1 J2 ) ≤ σ̃max ∗ σ̃max , ∀i ∈ [1, . . . , n] k If we assume that the long term Jacobian product Gk = J1 J2 . . . Jk (Section 2.3) has singular values bounded by σ̃min k and σ̃max , then
k k σ̃min ∗ σ̃min ≤ σmin (Gk )σmin (Jk+1 ) ≤ σi (Gk Jk+1 ) ≤ σmin (Gk )σmin (Jk+1 ) ≤ σ̃max ∗ σ̃max ⇐⇒ k+1 k+1 σ̃min ≤ σi (Gk Jk+1 ) ≤ σ̃max
Hence, by induction, ADPTNet long-term Jacobians Gk always have singular values bounded by matching powers of σ̃min and σ̃max . Taking the natural logarithm of Gk is then bounded as k k ln(σ̃min ) ≤ ln(σi (Gk )) ≤ ln(σ̃max ) ⇐⇒ ln(σ̃min )k ≤ ln(σi (Gk )) ≤ ln(σ̃max )k
Considering the formula for obtaining the Lyapunov spectrum αi 7 (see Section 2.3), it becomes apparent that ∀ Lyapunov exponents are also bounded ln(σ̃min )k ≤ ln(σi (Gk )) ≤ ln(σ̃max )k ⇐⇒ 1 ln(σ̃min ) ≤ ln(σi (Gk )) ≤ ln(σ̃max ) ⇐⇒ k 1 ln(σ̃min ) ≤ lim ln(σi (Gk )) ≤ ln(σ̃max ) ⇐⇒ k→∞ k ln(σ̃min ) ≤ αi ≤ ln(σ̃max ), ∀i ∈ [1, . . . , n]
Lemma 2 is built on the assumption that the Lyapunov spectrum is derived from the singular values of the long-term Jacobian Gk . However, in practice this is not numerically stable [69, 50, 47], and, as such, the Lyapunov spectrum is typically computed as the sum of individual ln(σi (Jt )) terms, which are always sorted. Therefore, given this implementation detail, it is possible to derive even stronger bounds on the relationship between recurrent eigenvalues and the Lyapunov spectrum in ADPTNet: Lemma 3. When using numerically stable methods to obtain the Lyapunov Spectrum, for sufficiently small β and a given recurrent eigenspectrum Λ of ADPTNet, individual Lyapunov exponents are αi ≈ ln(λi ) + ∆i , where ∆i ∈ [− 8βλλmax , 8βλλmax ], λmax = λ1 ≥ λ2 ≥ · · · ≥ λn = λmin , and αmax = α1 ≥ α≥ · · · ≥ α= αmin . i i Proof. For simplicity, and without loss of generality, it is assumed that mini,j |λi − λj | > 8βλmax , i.e., all eigenvalues are further apart than 8βλmax . Then, let the singular values of each recurrent Jacobian Jt be σ̃1 ≥ σ̃2 · · · ≥ σ̃n , (t) (t) (t) where each σ̃i = λi + ∆i , −8βλmax ≤ ∆i ≤ +8βλmax . As highglighted in Lemma 2, for any given product 2 2 Jt Jt−1 , the only guarantees on its singular value spectrum are that σ̃min ≤ σi (Jt Jt−1 ) ≤ σ̃max . That is because matrix multiplication "mixes" the spectrum and only preserves upper/lower bounds for its extremes. Therefore, when computing the singular values of the final long-term Jacobian product Gt , one can still only bound its extremes. 7
While the notational convention is to denote the Lyapunov spectrum by λ, to avoid confusion with the recurrent eigenspectrum of ADPTNet, α is used instead.
19
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
However, if you compute the spectrum at each time step, for example using QR decomposition [47], it will be sorted (high to low) each time, preventing spectral mixing. (t)
Let αi be the un-normalised partially-computed Lyapunov exponent at time t. Then computing it becomes: (t)
(t−1)
αi = αi
+ ln(σ̃i (t) )
(t)
Assuming β is small enough, then ∆i can be treated as a perturbation, and the logarithm can be expanded with its Taylor series: (t)
(t−1)
+ ln(λi + ∆i )
(t−1)
+ ln(λi ) +
αi = αi
(t)
(t)
= αi
∆i λi
(72)
The final, time-averaged over T steps, Lyapunov exponent is: T
(t)
1X ∆ αi = ln(λi ) + i T t=1 λi 1 = ln(λi ) + λi
T
!
1 X (t) ∆ T t=1 i
!
(73)
PT (t) (t) We denote the time-average as µi = T1 t=1 ∆i . Since, each −8βλmax ≤ ∆i ≤ +8βλmax , then the average also (t) satisfies −8βλmax ≤ µi ≤ +8βλmax . Let ∆i = µλii , then the initial statement in Lemma 3 is proven. (t)
Corollary 2. If ∆i are sampled from a zero-mean distribution (e.g., U(−βλmax , βλmax )), and the sequence length T is sufficiently large, then α ≡ ln(Λ). Lemmas 2 and 3 differ only on an implementation technicality. Regardless of Lyapunov spectrum computational methodology, the LLE and the long-term behaviour of ADPTNet are still prescribed by the recurrent eigenvalues Λ and the choice of β. While the topological conjugate in the ADPTNet provides non-linearity through the matrix exponential and the multiplicative interaction between M (x) and x, one can inject additional non-linear processing by pre-pending activation functions before the KV -projections. For instance, if the state x passes through element-wise ReLU before being projected to k and v and normalised, the spectral bounds from Lemma 1 are unchanged. This is because, unless completely silent, a ReLU diagonal Jacobian has spectral norm ≡ 1, and thus does not influence dM dx x 2 through the derivation chain. This differs from traditional RNNs where non-linearities are typically bottlenecks in the propagation of the recurrent state. It should also be mentioned that "stitching" together linear dynamics to obtain non-linear behaviour is reminiscent of well-established methods in machine learning such as Recurrent Switching Linear Dynamical Systems [120] or the Piecewise-Linear RNN understanding of ReLU networks [23]. However, these methods typically focus on producing certain dynamical behaviours locally around critical points [163]. In terms of global stability properties, the focus in this line of research is to enhance training methodology through regularisation [47] or teacher forcing [89], rather than by the inherent properties of the recurrent architecture, like the ADPTNet model proposed here. 4.2
Eigenvalue Parametrisation and Positional Embeddings
As highlighted in Section 2.2, LTI SSMs have the favourable property of direct control of long-term dynamics through the eigenvalue parametrisation of the recurrent matrix. As established in Lemma 2, ADPTNet also enjoys theoretical guarantees on the direct influence of its recurrent eigenvalues Λ (Eq. 49) over its long-term non-linear dynamics, as 20
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
quantified through its Lyapunov spectrum. Therefore, the parametrisation of Λ plays a crucial role in the model’s performance. The general consensus in the SSM literature is that recurrent weights can be reduced to diagonal matrices containing their complex conjugate eigenspectrum [134]. Their parametrisation and initialisation are designed for stability in continuous time. Hence, eigenvalues are typically constrained through parametrisation to negative real parts (e.g., negative exponential). For simulation, the recurrent dynamics are discretised using techniques such as Zero-Order Hold (ZOH), and typically the discretisation step size ∆t is also introduced as a trainable parameter. ∆t also serves the secondary purpose of normalising the recurrent state and preventing magnitude explosion over long horizons [140]. Altogether, using ZOH, discrete eigenvalue λ̄ parametrisation takes the form of Eq. 74, where ∆t, λreal and λimag are i i initialised to the logarithms of the desired discretisation real and imaginary parts of the eigenvalues. ∆t
λ̄i = e−e
real
(eλi
imag
±ieλi
)
(74)
To ensure long-term slow decay, λi are initialised close to the unit circle [140, 76]. Therefore, ∆t is typically initialised in the range [10−4 , 10−1 ] or [10−3 , 10−1 ] depending on sequence length8 , and αireal = ln( 12 ). For the purposes of this work, the distribution of the imaginary parts at initialisation follows the S4D-Inv scheme proposed by Gu et al. [76] (Eq. 75) e
λimag k
n = π
n − 1 , n = dim(Λ)/2 2k + 1
(75)
At this stage, it is worth emphasising the logic behind SSM complex eigenvalues. Considering, without loss of generality, an unrolled recurrent state xt , i.e., viewed as the sum of inputs uk for k < t (Section 2.2), that uses a generic simplified view of the eigenvalue parametrisation in Eq. 74, then xt can be rewritten as: real
+iΛimag )
(tΛreal
tiΛimag
xt = et(Λ =e
tΛreal
=e
real
= etΛ
(e
real
u0 + e(t−1)(Λ u0 ) + e
imag
R(u0 , tΛ
(t−1)Λreal
) + ··· + e real
ũ0 + · · · + eΛ
+iΛimag )
real
u1 + · · · + e(Λ
(t−1)iΛimag
(e
Λreal
+iΛimag )
u1 ) + · · · + e
imag
R(ut−1 , Λ
Λreal
ut−1 + ut
iΛimag
(e
imag
ut−1 ) + (e0∗iΛ
ut )
) + R(ut , 0)
ũt−1 + ũt , ũk = R(uk , (t − k)Λimag ) imag
Where, because in SSMs eigenvalues come in conjugate pairs by convention, the products eiΛ real-valued block-diagonal rotation matrix products: cos(θ1 ) − sin(θ1 ) 0 iθ e u = R(u, θ) = .. . 0 0
sin(θ1 ) cos(θ1 ) 0 .. . 0 0
0 0 .. .
··· ··· ..
··· ···
. 0 0
0 0 .. . 0 cos(θn ) − sin(θn )
u can be rewritten as
u1 u2 .. . 0 sin(θn ) cos(θn ) u2n 0 0 .. .
(76)
In this rewritten form, xt becomes the weighted sum of rotated inputs ũ. One can change the indexing of each rotation so that instead of having frequencies relative to the current time step t, ũk = R(uk , (t − k)Λimag ), they are relative to the initial state, i.e., ũk = R(uk , kΛimag ) without any meaningful loss of expressivity. In this new form, ũ becomes the well-known Rotary Position Embedding (RoPE), widely employed in modern LLMs [166]. Therefore, complex-diagonal SSMs’ recurrent states are a weighted sum of RoPE-like positionally embedded inputs, with an exponentially decaying memory. This is an important distinction to make in the context of parametrising ADPTNet. The Lyapunov spectrum that ADPTNet is theoretically guaranteed to control only describes exponential decay rates, not oscillations. Conversely, that means that while it is possible to use a complex-valued Λ in Eq. 49, its 8 Discretisation step size has the implicit task of normalising the weighted sum of the inputs that is an SSM hidden state. Therefore, for example, if sequence length is ≈ 1000 time steps, one can use a ∆t = 10−1 as the slowest evolving timescale to roughly normalise the sum.
21
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
oscillations may not materialise as intended in the non-linear dynamics. In contexts such as Section5.1.3, to emphasise this explicit link between recurrent eigenvalues and the Lyapunov spectrum, decay rates and oscillation frequencies are functionally separated in ADPTNet, becoming: real
Mt = Qt (±eΛ̄
real
)QTt , Q ∈ O(n), Λ̄real = diag(−e∆tλ1
real
, . . . , −e∆tλn )
imag
xt = Mt xt−1 + ũt , ũk = R(Buk , k Λ̄imag ), Λ̄imag = diag(e∆tλ1 = (Mt Mt−1 . . . M1 )ũ0 + · · · + Mt ũt−1 + ũt
,...,e
∆tλimag n/2
(77)
)
(78)
As opposed to vanilla RoPE, the frequency distribution in Λimag is still governed by S4D-Inv initialisation and parametrisation. A similar separation of recurrent eigenvalue imaginary components into a RoPE-based layer is also present in Mamba-3 [115]. However, there, the positional embeddings are also input-dependent. Grazzi et al. [70] recently suggested that including negative eigenvalues may also improve performance on certain tasks, and therefore, they are also included here for investigation. This resulting real-value eigenspectrum parametrisation derived from S4D-Inv by separation from rotational positional embeddings, is referred to, in this study, as linspace initialisation, since all λi = ln( 21 ) and ∆t ∈ linspace(ln(∆tmin ), ln(∆tmax ), 2n). Following [76], another real-valued eigenvalue initialisation scheme which forgoes rotations completely is S4D-Real, which takes the form: λi = ln(i)
(79)
Notably, while λi is initialised deterministically here, ∆ti ∼ U(ln(∆tmin ), ln(∆tmax )). Unless otherwise specified, experiments in this study default to the linspace initialisation. 4.3
Low-Parameter Linear Layers
As described in Section 4.1, ADPTNet recurrence makes use of three linear projection matrices for u at every time step: input (B), Key (K), and Value (V ). This means that the width h, and, intrinsically, the number of timescales for the dynamics, come with a taxing parameter budget of 3h2 just for linear projections alone. This imposes a sharp trade-off between temporal feature expressivity/context compression and overparametrisation. By contrast, for instance, SSMs such as S4 [74] or S5 [164] avoid this pitfall since they lack K and V projections. Therefore, to maintain comparability with such benchmark baselines in Section 5, ADPTNet layers need a more parameter-efficient parametrisation. This work explores two such methods. Firstly, Low-Rank Matrices (Eq. 80) are ubiquitous in both machine learning and deep learning research through methods such as Singular Value Decomposition (SVD)-based data compression or Low-Rank Adaptation (LoRA) [97] for memory-efficient LLM fine-tuning, respectively. This factorisation reduces the number of parameters from n2 to 2nr, where typically r ≪ n. Furthermore, low-rank vector-matrix multiplication A(B T v) avoids materialising a full n × n matrix and thus also reduces memory and compute overheads. M = AB T , M ∈ Rn×n , A, B ∈ Rn×r , r ≤ n
(80)
Secondly, Kronecker products (Eq. 81) can also drastically reduce parameter counts, from n2 to Na2 + Nb2 , where Na , Nb ≪ n. Similar to low-rank matrices, Kronecker factorisation also admits faster matrix-vector multiplication without full n × n materialisation, using the vec trick (Eq. 82). In addition, they have also been extensively studied in neural networks, for instance, for Fisher information matrix approximation [125] or even alternatives to LoRA [44]. a11 B a21 B M = A ⊗Kron B =
aNa 1 B
a12 B a22 B .. . aNa 2 B
... ...
a1Na B a2Na B , M ∈ RNa Nb ×Na Nb , A ∈ RNa ×Na , B ∈ RNb ×Nb (81)
. . . aNa Na B
M x = (A ⊗Kron B)x = ((B T x.reshape(Na × Nb ))A).reshape(Na Nb ), x ∈ RNa Nb 22
(82)
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
Both factorisations inject priors into the networks they parametrise, as visualised in Figure 6. Figure 6 consists of a compression task, where the goal is to reconstruct the underlying image using low-parameter matrix factorisations. Evidently, low-rank factorisation restricts the weight space to low-dimensional subspaces. This is useful in applications such as Low-Rank RNNs [127], where the goal is specifically to create interpretable dynamics in a low-dimensional state space. However, as a general-purpose substitute for dense linear layers, this is potentially severely limiting. As shown in Figure 6, this is intuitively equivalent to over-smoothing an image and losing the fine detail. Likewise, while Kronecker products are full-rank, since rank(A ⊗Kron B) = rank(A) ∗ rank(B), they suffer from different strong inductive biases. Namely, the asymmetry in A ⊗Kron B creates a global-local resolution trade-off. As can be observed in Figure 6, a larger Na places emphasis on variance between image patches, and enables a globally complex reconstruction. Conversely, a larger Nb improves the local resolution of individual patches. While it fails for a high-variance image such as the one in Figure 6, it would lead to a higher-fidelity reconstruction of repeating patterns, for instance. In an abstract setting such as a generic linear layer, one cannot know a priori which bias is more appropriate. To hedge between these various inductive biases, this study employs the heuristic of combining the two: low-rank plus Kronecker product. The resulting LowParam linear layer (Eq. 83) is full-rank and admits faster than dense matrix multiplication. As shown in Figure 6, the low-rank adjustments can also help mitigate the harsh trade-off between A and B in the Kronecker product. While not as direct a sum as here, combining Kronecker products and Low-Rank matrices has been seen before in deep learning research, once again for LoRa-like LLM fine-tuning [159]. It should also be noted that while sparsification methods such as pruning are also heavily used, especially in neuromorphic literature [14], they typically do not provide any speed-ups compared to dense matrices [60].
LowParam(x) = A(B T x) + (C ⊗Kron D)x 23
(83)
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
ORIGINAL REFERENCE Rank: 4 Loss: 0.0210 Params: 32768
Rank: 16 Loss: 0.0156 Params: 131072
Na: 64, Nb: 64 Loss: 0.0169 Params: 8192
Na: 128, Nb: 32 Loss: 0.0146 Params: 17408
Na: 32, Nb: 128 Loss: 0.0192 Params: 17408
Na: 256, Nb: 16 Loss: 0.0123 Params: 65792
Na: 16, Nb: 256 Loss: 0.0217 Params: 65792
Rank: 2 Na: 64, Nb: 64 Loss: 0.0164 Params: 24576
Rank: 2 Na: 128, Nb: 32 Loss: 0.0143 Params: 33792
Rank: 2 Na: 32, Nb: 128 Loss: 0.0181 Params: 33792
Rank: 4 Na: 64, Nb: 64 Loss: 0.0159 Params: 40960
Rank: 4 Na: 128, Nb: 32 Loss: 0.0141 Params: 50176
Kronecker
Low Rank
Rank: 2 Loss: 0.0263 Params: 16384
Rank: 2 Na: 256, Nb: 16 Loss: 0.0121 Params: 82176
Low Param
Rank: 4 Na: 32, Nb: 128 Loss: 0.0173 Params: 50176
Figure 6: Patterns of Image Reconstruction Error Each row represents a matrix factorisation method, each producing different patterns of distortion when trained with gradient descent to reconstruct a target image on a low parameter budget. The reconstruction loss reported is Mean Squared Error (MSE). On one hand, Low-Rank matrices are shown to oversmooth the image on very low ranks ∈ {2, 4}. On the other hand, Kronecker products have much more accurate reconstructions on lower parameter budgets, but the performance varies massively depending on the ratio between Na and Nb . The proposed LowParam heuristic is shown to partially mitigate the extremes of Kronecker products with low-rank corrections.
Since the goal is for the LowParam layer to behave similarly to a typical dense linear layer, its initialisation is constructed to match the output pre-activation distribution from weight matrix [83]. For a square matrix, this qa Kaiming-initialised q entails sampling individual weights from ∼ U −
3 n, +
3 n
. The heuristic initialisation solution used here employs
1 the LoRA initialisation scheme for the low-rank component, where B ≡ 0 and A ∼ N (0, rank ). For the Kronecker q q 1 1 product C ⊗Kron D, C ≡ 1 and D ∼ U − Na Nb , Na Nb . As shown in Figure 7, the output distribution closely
matches that of the baseline Kaiming Uniform-initialised dense, with a Kullback-Leibler (KL) divergence ≈ 10−2 . 24
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
0.8
Kaiming Init Low Param Kaiming Init KDE LowParam KDE
0.7
Kaiming Init: Mean=-0.0011, Std=0.5774 Low Param: Mean=-0.0060, Std=0.5804 KL Divergence: 0.009805
0.6
Density
0.5 0.4 0.3 0.2 0.1 0.0
3
2
1
0 1 Pre-Activation Output
2
3
Figure 7: Pre-Activation Distribution Comparison between LowParam and Kaiming Initialisation The plot shows the output distribution of the proposed LowParam heuristic compared to a Kaiming-Uniform initialised dense matrix. The inputs are 1024 random vectors ∈ R1024 with entries sampled from N (0, 1). The LowParam layer uses rank = 16 for the low-rank component and Na , Nb = 128, 8 for the Kronecker product. While less smooth than the baseline, the initialisation scheme closely reproduces the mean and variance of the original, with a KL-divergence of ≈ 10−2 between the two distributions. 4.4
Memory and Compute-Efficient Matrix Exponential
Section 4.1 highlighted that the purpose of using orthogonal matrices in the parametrisation of the topological conjugates is to avoid the computational inefficiency and numerical instability of generic matrix inverses. One may notice, however, that the Cayley Map (Eq. 31) used to map the key-value pairs to the orthogonal manifold also employs a matrix inverse. Because the skew-symmetric matrix S = kv T − vk T is rank-2, it could be argued that the inverse does not incur a computational cost since it can be derived efficiently using the Woodbury Identity [189]. However, in practice, in this particular case, it suffers from even worse conditioning than baseline inversion, and even degrades with wider networks [81], thus making it effectively unserviceable. A more efficient and effective method for computing the orthogonal retraction is to take full advantage of the idiosyncratic rank-2 structure of S in the matrix exponential function, rather than the Cayley Map approximation. To make the rank-2 structure explicit, S can be rewritten as S = EF T , E = (k
−v) ∈ Rn×2 , F = (v
k) ∈ Rn×2
It is assumed that the β step size parameter from Eq. 51 is absorbed into k =
√
Kx β ||Kx|| and v =
√
β ||VV xx|| .
To start, Theorem 1.35 from Higham [91] (Eq. 84) provides a means to compute f (M ) for any function f , including exponentials, such that if M ∈ Rn×n is rank r < n (M = AB T , A, B ∈ Rn×r ), then the problem can be reduced to computing f for a smaller r × r matrix instead of the full n × n. In this case, the α scalar is set to zero, and f is the matrix exponential function, resulting in Eq. 85 for eS . f αIn + AB T = f (α)In + A(B T A)−1 (f (αIr + B T A) − f (α)Ir )B T , α ∈ R, Ik = eye(k) T
eS = In + E(F T E)−1 (eF E − I2 )F T C = FTE =
T k v kT k
−v T v −v T k
=
k·v β
−β −k · v
(84) (85)
(86)
Computing the exponential retraction is thus reduced to finding the inverse and exponential of a 2 × 2 matrix C = F T E (Eq. 86). From Eq. 86, it should be noted that the determinant det(C) = β 2 − (k · v)2 is always ≥ 0 since |k · v| ≤ ∥k∥ ∥v∥. In practice, the determinant is strictly > 0 since K and V are trained separately, producing distinct k and v, and a small perturbation can also be introduced to enforce strict positivity. Because T r(C) = 0 and the determinant is positive, computing eC simplifies to 25
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
p p sin( det(C)) eC = cos( det(C))I2 + p C det(C)
(87)
−1 And C −1 = det(C) C. Because of the explicit formulation of eS in terms of low-rank components, the final orthogonal Q never has to be materialised in full n × n, allowing for memory and compute-efficient matrix-vector multiplication:
Qx = eS x = (In + EC −1 (eC − I2 )F T )x = x + E((C −1 (eC − I2 ))(F T x)) 4.5 4.5.1
ADPTNet Parallelisation ADPTNet DEER Convergence Properties
In Lemma 2, it is established that the Lyapunov spectrum of ADPTNet is determined by the recurrent eigenvalues Λ and the β manifold step size parameter. This level of parametric control has secondary benefits in determining the convergence properties of the DEER algorithm [118, 68] (see Section 2.4) when applied to ADPTNet. DEER displays different convergence properties depending on the region of the trajectory guess s state space (Eq. 18). As proven in Gonzalez et al. [69], DEER is globally guaranteed to converge to the correct trajectory s∗ at least linearly. In other words, the distance of the current guess s(i) to s∗ decays by at least a global constant factor ∈ (0, 1) between Newton iterations. Furthermore, DEER can even converge quadratically, i.e., practically in a constant number of iterations, if the residual r(s(i) ) is sufficiently small. The quadratic convergence condition, as stated in Theorem 5 from Gonzalez et al. [69], is defined by the Lipschitz constant L of the recurrent Jacobian and the LLE αmax of the system being parallelised:
r(s(i) )
< 2
2 2 a L
eαmax − 1 eαmax T − 1
2 (88)
In Eq. 88, T denotes the total number of time steps of the trajectory, and a ≥ 1 is a constant designed to capture the distortion effect of transient dynamics that may push the observed LLE to be larger than its actual long-term value (i.e., overshoot): ∥Jt+k−1 . . . Jt ∥ ≤ aeαmax , t > 1, k ≥ 0
(89)
It should be observed that since the LLE α is present in the definition of the quadratic convergence region, the Λ parametrisation in ADPTNet exerts direct influence over it. In practical terms, a lower λmax and a sufficiently low β can increase the domain where DEER converges quadratically. Moreover, it should be re-emphasised that the Lipschitz constant L for DEER is inherited from the Lipschitz-ness of the recurrent Jacobian J of the dynamics being parallelised (Theorem 2 from Gonzalez et al. [69]). Concretely, L is defined as either Eq. 90 or Eq. 91 [109]:
∥J(x) − J(y)∥2 ≤ L ∥x − y∥
L = sup x∈Rn
dJ(x) dx 2
(90)
(91)
Based on Eq. 91, for ADPTNet dynamics, using Eq. 54 - 65 and the A and B shorthand notation from Section 4.1, Lipschitz-ness can be derived as: 26
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
L = sup x∈Rn
dJ(x) dx 2
d(M + dM dx x) n dx x∈R 2 dM dM d2 M = sup x + + dx dx dx2 x∈Rn 2 dS d dS dS T T dS = sup 2 βA B−B A +β A B−B A x dx dx dx dx dx x∈Rn 2 dS d dS dS dS = β sup 2 A B − BT A + A B − BT A x dx dx dx dx dx x∈Rn 2 = sup
(92)
Without having to actually evaluate the upper bound on the norm in Eq. 92, it can be seen that L ∝ β (L = βc, for a constant c ≥ 0). Therefore, the size of the quadratic convergence region for DEER is itself ∝ β1 . In addition, as highlighted in Lemma 2, the LLE αmax is always bounded ∈ [ln(σ̃min ), ln(σ̃max )], regardless of the trajectory length or region in state space. Therefore, one can define the constant 0 ≤ d ≤ 1 such that αmax = d ln(σ̃min ) + (1 − d) ln(σ̃max ) Altogether, the formula for the DEER quadratic convergence basin as a function of ADPTNet parametrisation can be summarised as: (i)
r(s
2 < 2 a βc 2
ed ln(σ̃min )+(1−d) ln(σ̃max ) − 1 e(d ln(σ̃min )+(1−d) ln(σ̃max ))T − 1
2 (93)
Eq. 93 has two important implications. Firstly, the lower the extremal eigenvalue parameters λmin and λmax , the lower the bounds on the LLE αmax and thus the larger the basin of quadratic convergence. This is intuitive since a lower LLE effectively imposes a more quickly decaying memory for the recurrent dynamics. In turn, any error in the current state (i) (i+1) guess at time t, st , will influence fewer subsequent k time steps st:t+k in the next Newton iteration (i.e., αmax ∝ k, where k is the number of subsequent time steps affect by the current one). Conversely, if the LLE is positive, the recurrent dynamics become chaotic, and the quadratic convergence region vanishes. In practice, this chaotic regime forces a number of Newton iterations close to or equal to the total number of time steps T being simulated. If β is selected sufficiently low, ADPTNet prevents this by parametrising the eigenspectrum Λ ∈ [0, 1]n (see Section 4.2). Secondly, a lower β implies a larger basin of quadratic convergence. Consider the extreme case where β = 0, then L = 0. L = 0 is equivalent to a linear system which converges in a single DEER iteration, with an infinitely-sized quadratic convergence region. This can be easily verified since β = 0 implies that Q = In , and the similarity transform in Eq. 49 becomes In ΛInT = Λ, equivalent to an LTI SSM that can be easily parallelised with a single convolution/parallel scan. To conclude this section, it should also be mentioned that ADPTNet can also be reduced to linear dynamics by enforcing a constant eigenspectrum λmin = λmax = λ. In that case, the recurrent dynamics become QλIn QT x = λQQT x = λx, regardless of the value of β. 4.5.2
Conv-DEER and Forward-DEER
Computing the full n × n recurrent Jacobians Jt required for vanilla DEER is both computationally and memory intensive, effectively limiting its scalability. Gonzalez et al. [68] introduced Quasi-DEER to account for this, replacing the full Jacobian with its diagonal diag(Jt ). In ADPTNet dynamics, for sufficiently small β, the Jacobian Jt = M + dM dx x is equivalent to the forward dynamics M plus a perturbation βP (see Section 4.1). Taking advantage of this, ADPTNet admits two efficient, Jacobian-free, Quasi-DEER heuristic approximations. In this context, this is understood to mean an approximation scheme for diag(Jt ) which does not require computing any derivatives. 27
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
The most inexpensive approximation approach proposed here is to replace diag(Jt ) with the recurrent eigenvalue spectrum Λ. This comes at a minimal computational cost, since Λ is known a priori for all time steps and only requires to be materialised through ZOH discretisation once per batch 9 . Furthermore, sharing Λ across batched execution provides immediate memory savings proportional to the batch size. Using constant Λ at each time step also converts each DEER iteration into an LTI system that can be computed using global and efficient O(T log(T )) convolutions based on FFTs (see Section 2.2), in addition to parallel associative scans. Hence, this method is referred to as Conv-DEER. This opens the door to taking advantage of highly optimised CUDA libraries such as FlashFFTConv [58], or even potentially using dedicated FFT hardware accelerators [62]. To understand why Conv-DEER could be a viable heuristic, it is important to revisit how DEER is implemented in practice [68]. First, the term b is computed for each time step: (i)
(i)
bt = f (st−1 ) − At st−1 Where f is the recurrent function and At is the recurrent Jacobian Jt for DEER, diag(Jt ) for Quasi-DEER, and Λ for Conv-DEER. The next trajectory guess s(i+1) is then computed as: (i+1)
st
= At−1 . . . A1 b0 + At−1 . . . A2 b1 + · · · + At−1 bt−2 + bt−1
(94)
In vanilla DEER, the products of the form At At−1 . . . At−k+1 are exactly long-term Jacobians Jt . . . Jt−k+1 whose effect on bt−k , e.g., exponential decay, is described in some capacity by the Lyapunov spectrum and especially the LLE (see Eq. 89). As stated in Lemma 2, in the case of ADPTNet the bounds on the entire Lyapunov spectrum are constrained by the choice of Λ and β. Therefore, intuitively, using Λk−1 instead of At At−1 . . . At−k+1 should produce a similar long-range propagation effect, subject to the choice of β. The second heuristic proposed here is to take the diagonal of the forward recurrent dynamics diag(Mt ) as an approximation to the full Jacobian diagonal diag(Jt ) from Quasi-DEER. Accordingly, this method is referred to as Forward-DEER. For ADPTNet using the low-rank-aware matrix exponential from Section 4.4, this does not require materialising the full recurrent matrix Mt . The notation in Eq.85 can be updated to emphasise how the exponential is an adjustment to the identity with left L and right R low-rank decomposition terms: T
eS = In + E(F T E)−1 (eF E − I2 )F T T
= In + LRT , L = E(F T E)−1 (eF E − I2 ) ∈ Rn×2 , R = F The full matrix Mt and its diagonal become: M = (In + LRT )Λ(In + RLT ) = Λ + ΛRLT + LRT Λ + LRT ΛRLT diag(M ) = Λ + 2Λ(l1 ⊙ r1 + l2 ⊙ r2 ) + (ˆl1 l1 + ˆl2 l2 ), L̂ = L((RT Λ)R)
(95)
Where lower-case l, r, and ˆl denote individual columns of L, R, and L̂ respectively. It should be noted that no full n × n matrix-matrix or matrix-vector multiplications are required to compute diag(Mt ). Both Jacobian-free heuristics proposed in this study evidently introduce approximation errors compared to vanilla or even Quasi-DEER. In this context, the step-wise approximation errors ϵapprox are quantified as the norm of the t difference between approximated b̃t and the actual bt obtained with vanilla DEER At = Jt : ϵapprox = bt − b̃t t
2
= ∥f (st−1 ) − Jt st−1 − f (st−1 ) + At st−1 ∥2 = ∥(At − Jt )st−1 ∥2 = (At − Mt − 9
dMt st−1 )st−1 dst−1 2
For more details about exponential parametrisation see Section 4.2
28
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
dMt xt−1 from Section 4.1. Starting with On the last line, Jt is instantiated with the ADPTNet Jacobian Jt = M + dx t−1 Conv-DEER where At = Λ the approximation error takes the form:
dMt ϵapprox = (Λ − Mt − st−1 )st−1 t dst−1 2 dMt T Λ − Qt ΛQt 2 + ≤ ∥st−1 ∥2 st−1 dst−1 2
(96)
= (ϵsim + ϵder t t ) ∥st−1 ∥2 Here, ϵsim and ϵder notations are introduced for the approximation errors resulting from the unaccounted effect of the t t similarity transform (i.e., "misalignment") and missing derivative term, respectively. ϵder is upper-bounded by 8βλmax , t as derived in the proof of Lemma 1. For ϵsim , it is important to reiterate that Λ is positive real-valued and diagonal, and t so it is also, by extension, Hermitian. Q from the similarity transform is orthogonal and thus unitary as well. Therefore, Λ and QΛQT exist in the same unitary orbit, and the following bound on ϵsim applies [90, 38]: t ϵsim = Λ − Qt ΛQTt t
2
≤ λmax − λmin
(97)
Forward-DEER trivially presents the same derivative approximation error ϵder as Conv-DEER, however, instead of a t diag misalignment error, it introduces an off-diagonal entry error ϵt = ∥diag(Mt ) − Mt ∥. As showcased in Figure 5, the magnitude of the off-diagonal entries in Mt is directly controlled by β. To highlight further distinctions between Forward and Conv-DEER, one can take Quasi-DEER as a baseline and consider the low-rank-aware computation of the matrix exponential:
diag(Jt ) = diag(Mt ) + diag(
dMt st−1 ) dst−1
dMt = Λ + 2Λ(l1 ⊙ r1 + l2 ⊙ r2 ) + (ˆl1 l1 + ˆl2 l2 ) + diag( st−1 ) dst−1 It can be observed that Conv and Forward-DEER are effectively truncated approximations of the diagonal of the full Jacobian from Quasi-DEER. Furthermore, the difference between the two approximations ϵft wd = 2Λ(l1 ⊙ r1 + l2 ⊙ r2 ) + (ˆl1 l1 + ˆl2 l2 ) ∝ λmax since L̂ also contains Λ in its composition. It should also be 2 noted that when using the efficient matrix exponential from Section 4.4, the k and v vectors comprising E, F , and thus √ L and R as well, each have norm β. Hence β also determines the bounds of the ϵft wd approximation error. So far, the focus has been placed on local approximation errors, stemming from computing individual bt terms. However, (i+1) (i) each state estimate st is a weighted sum of all bk for k < t (Eq. 94). Therefore, all local errors spread to subsequent time steps: (i+1)
s̃t
= At−1 . . . A1 ϵapprox + At−1 . . . A2 ϵapprox + · · · + At−1 ϵapprox + ϵapprox 0 1 t−2 t−1
(98)
Here, s̃ denotes the accumulation of approximation errors in the trajectory guess. The ϵapprox terms are used with a k wider scope than originally in Eq. 96, since it now differs based on the approximation and baseline methods being (i+1) compared (see Table 1 for summary) The influence of local error terms ϵapprox over future time steps s̃k+j depends k on the long-term products Ak+j−1 . . . Ak . This "transmission" term can itself introduce errors in the DEER iteration. For instance, if one compares Conv-DEER with vanilla DEER, the A-products will themselves differ by at most (λmax − λmin + 8βλmax )j (Eq. 96 and 97). However, given Lemma 2, the overall effect of Jk+j−1 . . . Jk compared to Λj is tightly related to the LLE of the recurrent dynamics. In other words, the magnitude of a local ϵapprox , after being k propagated long term over j time steps, will decay at a similar exponential rate ≈ eαmax j . Therefore, the scope here will be limited to measuring the overall magnitude of accumulating local ϵapprox terms, omitting "transmission" errors. k j Considering the approximation of the long-term effect of Ak+j−1 . . . Ak ≈ eαmax j ≤ σ̃max , the accumulation of local approximation errors at a given time step t becomes Eq. 99. The overall upper bound on the accumulation over all time (i+1) steps T of the trajectory of local errors S̃ = ΣTj=1 s̃j is Eq. 100.
29
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
Approx. Method
Conv-DEER
Forward-DEER
ϵsim ≤ λmax − λmin k ϵder ≤ 8βλmax k
N/A
Baseline Vanilla DEER Quasi-DEER
ϵfk wd = 2Λ(l1 ⊙ r1 + l2 ⊙ r2 ) + (ˆl1 l1 + ˆl2 l2 )
Forward-DEER
ϵfk wd = 2Λ(l1 ⊙ r1 + l2 ⊙ r2 ) + (ˆl1 l1 + ˆl2 l2 )
ϵder ≤ 8βλmax k
2
2
∝ λmax , β
∝ λmax , β
ϵder ≤ 8βλmax k N/A
Table 1: Summary of the different local approximation error terms that make up ϵapprox depending on the baseline k established method (DEER and Quasi-DEER) and the approximation methods proposed here (Conv and ForwardDEER). The error terms listed are also multiplied by the norm of the state sk−1 , to obtain ϵapprox = (ϵ + . . . ) ∥sk−1 ∥2 . k
(i+1)
s̃t
t−1 approx t−2 approx = σ̃max ϵ0 + σ̃max ϵ1 + · · · + σ̃max ϵapprox + ϵapprox t−2 t−1
(99)
T −1 1 T −2 1 S̃ (i+1) = (σ̃max + · · · + σ̃max + 1)ϵapprox + (σ̃max + · · · + σ̃max + 1)ϵapprox 0 1
+ · · · + (σ̃max + 1)ϵapprox + ϵapprox t−2 t−1 T T −1 − 1 approx σ̃max − 1 approx σ̃max ϵ0 ϵ + σ̃max − 1 σ̃max − 1 1 1 − 1 approx σ̃ 2 − 1 approx σ̃max ϵt−2 + ϵ + · · · + max σ̃max − 1 σ̃max − 1 t−1 1 T −1 = ((σ̃ T − 1)ϵapprox + (σ̃max − 1)ϵapprox 0 1 σ̃max − 1 max approx approx 2 1 + · · · + (σ̃max − 1)ϵt−2 + (σ̃max − 1)ϵt−1 )
=
(100)
The main takeaway from laying out the sketch of a loose upper bound on accumulating local approximation errors in Eq. 100 is that, perhaps unsurprisingly, it increases with sequence length T and σ̃max , which, in turn, is λmax + 8βλmax . In addition, if the overview in Table 1 is also considered, it is possible to make a number of predictions regarding the performance of the proposed approximation schemes. Conv-DEER should generally require more Newton iterations than vanilla, Quasi, and Forward DEER. In particular, considering ϵsim and ϵder error terms, its convergence should be impacted by the spread of the recurrent eigenvalues (λmax − λmin ), β, and the LLE (also determined by λmax and β). This expectation also rests on the intuition that a higher spread of the eigenvalues allows a smaller β to have higher variance in decay rates "chosen" per element of the recurrent state xt . Forward-DEER should be expected to be less sensitive to eigenvalue spread and generally track closer to Quasi-DEER than Conv-DEER. However, a larger LLE also entails, by definition, longer-term memory, which generally "smears" state-guess errors further into the future. Thus, both approximations, much like DEER in general, incur a penalty in convergence speed when the network has long-range memory. 4.5.3
Damping
Gauss-Newton iterative methods such as DEER are known to suffer from instability. Accordingly, as first introduced in Gonzalez et al. [68], a notable DEER stabilisation strategy is to leverage trust regions, i.e. the Levenberg-Marquardt method. The resulting solution, Evaluating Levenberg-Marquardt with Kalman (ELK), effectively dampens the eigenvalues of the recurrent Jacobians Jt to avoid "exploding" behaviour when computing the J −1 r product in DEER (see Section 2.4). Furthermore, to minimise computational overhead, as also proposed in Gonzalez et al. [68], eigenvalue damping can also be achieved by simply multiplying recurrent Jacobians by a scalar constant ∈ [0, 1] (Scale-ELK). With the added approximation errors introduced by the Conv and Forward-DEER methods proposed here, numerical instability may also be a challenge. As established in Section 74, ADPTNet is constructed with stable recurrent Jacobians by design. However, throughout DEER convergence, intermediary trajectory guesses s(i) are not subject 30
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
to the same normalisation constraints as in valid ADPTNet trajectories s∗ . Moreover, Section 4.5.2 shows that the magnitude of local approximation errors depends on the norm of these state guesses st as well as parameters such as β and λmax . Propagating unnormalised local errors ϵapprox in Eq. 94 may in itself cause instabilities, and thus damping may be necessary. conv
fwd
Let Ãt and Ãt denote Jacobian-free approximations of Quasi-DEER diag(Jt ) from Conv or Forward-DEER respectively. Using Quasi DEER in conjunction with Scale-ELK damping Damping(At ) = kAt , k ∈ [0, 1] on ADPTNet recurrent Jacobians results in:
k(diag(Jt )) = k(diag(Mt ) + diag(
dMt xt−1 )) dxt−1
dMt xt−1 )) dxt−1 dMt conv = k(Ãt + 2Λ(l1 ⊙ r1 + l2 ⊙ r2 ) + (ˆl1 l1 + ˆl2 l2 ) + diag( xt−1 )) dxt−1 dMt fwd xt−1 )) = k(Ãt + diag( dxt−1
= k(Λ + 2Λ(l1 ⊙ r1 + l2 ⊙ r2 ) + (ˆl1 l1 + ˆl2 l2 ) + diag(
(101)
dMt xt−1 and, implicitly, its diagonal, are As established in Lemma 1, the bounds on the norm of the derivative term dx t−1
determined by β. If β is chosen sufficiently small, then
dMt dxt−1 xt−1
≪ Ãfwd t . Therefore, for a certain damping k, conv
k(diag(Jt )) ≈ k Ãfwd . In t . While including additional approximation error, a similar logic can also be applied to Ãt other words, damping should align Jacobian-free approximations with Quasi-DEER (and similarly vanilla DEER) and lead to more and more similar convergence behaviour as k increases.
4.5.4
Block-Wise DEER
As proven in Gonzalez et al. [68], even in the worst-case, DEER and its more efficient approximations are guaranteed to converge correctly on the first i states of the trajectory within the first i Newton iterations. If j more Newton iterations are required to converge to the fully correct trajectory s∗ , compute and memory will still be wasted recomputing the already converged first i steps, at least. When training, for example, all DEER intermediary trajectory guesses have to be stored for back-propagation, imposing a particularly heavy memory cost (Subfigure 8a). In contrast, fully sequential simulation, by definition, only has to compute each time step once, i.e., no wasted resources. However, this evidently comes at the cost of no GPU parallelism and slower wall-clock simulation10 . Therefore, one solution towards a Pareto-optimal balancing between memory utilisation and GPU parallelism during training is to execute DEER in a block-wise sequential manner (Subfigure 8b). T Intuitively, the original T -length sequence is split into B B-sized blocks. DEER or DEER-approximate parallelisation is then applied to each individual block. The blocks are sequentially computed and chained together by passing the final state from one as the initial state to the next. In sum, the algorithm pseudocode is summarised in Algorithm 1.
10
In the event that the dynamics are effectively parallelisable by DEER [69].
31
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
Algorithm 1 Block-Wise DEER Require: T > 0, B, inputs∈ RBatch Size×D×T , initial state ∈ RBatch Size×D , states∈ RBatch Size×D×T , niters Ensure: T %B = 0 Nblocks ← T /B input blocks ← inputs.chunk(Nblocks , dim = −1) states blocks ← states.chunk(Nblocks , dim = −1) states out ← empty list i←0 while i < nblocks do j←0 states guess ← state blocks[i] while j < niters do states guess ← DEER(states guess, initial guess, input blocks[i]) j ←j+1 end while i←i+1 initial state = states guess[..., -1] states out.append(states guess) end while return cat(states out, dim=-1)
(a) DEER
(b) Block-Wise DEER
Figure 8: Saving memory and compute with Block-Wise DEER execution. Subfigure 8a shows how DEER uses compute and, especially, memory to store intermediary copies of already converged states (the red shaded area). Subfigure 8b shows how splitting the sequence into 4 equally split sequential blocks, as an example, results in roughly 4× reduction in memory required for storing intermediary Newton iterates.
Block-wise execution is a prevalent strategy in maximising GPU utilisation in parallel architectures (see Section 3.2) [35, 58, 194, 121]. However, the form of Block-Wise DEER investigated here is most similar to the ParaRNN [33]. There, block-wise sequential processing for DEER-like algorithms is discussed in more detail with respect to hardware I/O optimisation, rather than DEER idiosyncrasies. 4.6
State Expansion
In this study, unless otherwise specified, inputs ut in the ADPTNet recurrent step are dense linear projections of the overall input to the layer ot ∈ Rh (ut = W ot , W ∈ Rh×h ). However, SSMs such as S4, the LRU, or Mamba [74, 140, 71] typically rely on parameter-efficient projections that expand the size of the recurrent state to n by a factor nstate compared to the hidden dimension h of the network (n = nstate h), to allow for higher memory capacity. The projections are implemented with vector-scalar multiplications per input dimension ui:i+nstate −1 [t] = wi ·oi [t], w ∈ Rnstate 11 . 11
It should be noted, for notational convenience, that ut refers to the entire input vector at time step t, whereas ui [t] denotes the i dimension of the input at time t. th
32
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
4.7
Input Normalisation
A crucial consideration in systems with slowly fading memory, such as SSMs or ADPTNet, is preventing recurrent state magnitude explosion. In fact, Orvieto et al. [140] showed how input and state normalisation at initialisation can determine whether SSMs converge beyond random accuracy or not on long-range sequence modelling tasks. One common solution is to adopt leaky integration (see Section 2.1) or exponential moving average (EMA) parametrisation [122]: xt = Λ̄xt−1 + (1 − Λ̄)u
(102)
Another normalisation technique, used in Orvieto et al. [140] to construct the Linear Recurrent Unit (LRU), is: xt = Λ̄xt−1 + (
p In − Λ̄2 )ut
(103)
Finally, one can also tie input normalisation to the step size ∆t, by extending ZOH discretisation to inputs as proposed in [76]: (e−e
∆t Λ
e
− 1) ut (104) eΛ For simplicity, from now on let γ ∈ Rn denote the element-wise input normalisation term in the ADPTNet recurrent p −e∆ti eλi −1) step, either γi = 1 − λi for EMA, γi = 1 − λ2i for LRU normalisation, or γi = (e for ZOH. In addition, eλ i following Orvieto et al. [140], γ can either be set as a function of Λ throughout training, or simply initialised in relation to Λ and then allowed to be trained independently. xt = Λ̄xt−1 +
In the context of ADPTNet, a further consideration is the similarity transform applied to the recurrent eigenspectrum QΛQT . As a consequence of the rotations, there is no element-wise correspondence between xt−1 and ut to perfectly match individual γ terms. However, several strategies can be investigated for mitigation in the eventuality it is necessary. Firstly, one could apply a rotation to the input as well to preserve element-wise matching between Λ and 1 − Λ. This rotation-aware normalisation could take the form of: xt = Qt (ΛQTt xt−1 + γ ⊙ ut )
(105)
It should be observed that any form of rotation-aware normalisation introduces an additional term in the recurrent Jacobian that has a norm proportional to the input ut (Equation 106). Hence, this breaks the Jacobian norm bound guarantees derived in Lemma 1 and thus does not trivially inherit the theoretical properties of ADPTNet. dM dQ x+ (γ ⊙ u) (106) dx dx A second solution would be not to consider the effect of the EMA mismatch. It could be argued that the 2-norm bounds of the recurrence step are invariant to any rotation applied to the input, and thus the recurrence can remain unchanged: J =M+
xt = Qt ΛQTt xt−1 + γ ⊙ ut Finally, it is important to reiterate that Orvieto et al. [140] found significant performance degradation due to a lack of proper normalisation when testing on long sequences, with eigenvalues close to the unit circle. Conversely, if Λ is initialised for quick decay, the problem should be less pressing. Therefore, one could take inspiration from adLIF neurons (Subsection 3.3) [18] to separate long-range and short-range timescales (Λ) into an ADPTNet non-linear module, and a long-term EMA linear memory (Γ): xt = Mt xt−1 + γ ⊙ (ut + gain ⊙ wt−1 ) wt = Γwt−1 + (1 − Γ)xt
(107)
The model resulting from Eq. 107, named here CoupledADPTNet, reverses the tradition of focusing on linear memory units to complement non-linear RNNs [179]. Here, the secondary purpose of the CoupledADPTNet is to study the robustness of the proposed non-linear dynamics to coupling with another system as a function of the feedback gain term, rather than the long-term memory capacity of the linear subnetwork, which is essentially a standard SSM. 33
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
4.8
Input Selectivity
On tasks requiring in-context adaptation or selectivity, to match the formulation of selective SSMs such as Mamba, ADPTNet can be equipped with a data-dependent element-wise input gate (i): xt = Mt xt−1 + i(ut ) ⊙ γ ⊙ ut
(108)
Following the convention from Mamba, the input gate i(ut ) = σ(U DT ut ), where low-rank U, D ∈ Rn×r , and σ := softplus. 4.9
ADPTNet Taylor Series Approximation
For low β, the computational cost of the ADPTNet recurrence (Sec. 4.1) can be reduced by approximating the similarity transform with its truncated Taylor series expansion. For any arbitrary matrix X and orthogonal matrix Q = eβS , where S is skew-symmetric, using the Baker-Campbell-Hausdorff formula [24], a similarity transform can be expanded as: QXQT = eβS Xe−βS = X + β[X, S] + β 2 ∗ [X, [X, S]] + . . .
(109)
Where [X, S] denotes the commutator XS − SX. Accordingly, the ADPTNet recurrence, for sufficiently low β, can be first-order approximated as: xt+1 = MTaylor xt + γut+1 MTaylor = Λ̄ + β(Λ̄(kv T − vk T ) − (kv T − vk T )Λ̄)
(110)
It should be noted that matrix exponentials are no longer required. Furthermore, the diagonal of the commutator term cancels out in the subtraction, and thus MTaylor ’s diagonal is exactly the parametrised eigenspectrum, while off-diagonal entries are directly scaled by β. Hence, for TaylorADPTNet, Conv and Forward DEER are identical. 4.10
Linear ADPTNet
One can notice that the topological conjugation backbone of ADPTNet’s recurrence does not inherently have to be non-linear. Removing the state dependence, the derivative term dM dx x vanishes, and the Lyapunov spectrum of the now LTV system is identical to its eigenspectrum. Therefore, to isolate the expressivity of the similarity transform itself, separate from any non-linear recurrence effects, LinearADPTNet can be used: xt+1 = MLinear xt + γut+1 MLinear = Q(ut+1 )Λ̄Q(ut+1 )T , k =
Kut+1 V ut+1 ,v = ∥Kut+1 ∥2 ∥V ut+1 ∥2
(111)
Importantly, as an LTV system, LinearADPTNet can be parallelised using a single parallel associative scan [71]. However, this entails materialising all recurrent n × n MLinear matrices, at a memory cost of O(T n2 ), where T is the sequence length. Furthermore, all materialised recurrent weights have to be multiplied with O(T n3 ) computational cost. Therefore, one could still use Conv or Forward DEER in this context, to propagate off-diagonal entries. The k and v projection can be computed once at the first iteration and cached for the entire simulation, since they do not depend on state trajectory guesses. Then, the memory cost becomes O(T n) for all DEER iterations12 . In terms of computation, the new scaling O(niter T n) is advantageous as long as the number of DEER iterations required for convergence niter < n. To the author’s knowledge, Quasi-DEER or any diagonal form of DEER has not previously been used a means for more efficient full-matrix LTV system parallelisation. As with non-linear ADPTNet, for small enough β, LinearADPTNet can also be approximated to first order with a Taylor Series expansion. Interestingly, in this case, since dM dx x = 0, Quasi, Conv, and Forward DEER are all equivalent. 12
This cost assumes gradient checkpointing is applied during training, and intermediary DEER iterates are not stored for the backward pass. The result applies universally during inference.
34
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
4.11
Spiking ADPTNet
To materialise the neuromorphic inspiration from the auditory cortex (Section 1), ADPTNet layers can be equipped with position-wise spiking activations for information propagation across across model depth [165]. Using Eq. 5, the output of each layer becomes: outi [t] = s(xi [t]) =
0, xi [t] < θ 1, xi [t] ≥ θ
(112)
In Section 5.8, where ADPTNet-based SNNs are employed, following Fabre et al. [54], the Boxcar surrogate gradient is used in the backward pass: ds = dx 4.12
0, x ∈ / [−0.5, 0.5) 1, x ∈ (−0.5, 0.5)
(113)
Fine-Grained Control over ADPTNet Non-Linearity
The efficient matrix exponential function from Section 4.4 reveals how the non-linearity, or the level of off-diagonal interaction, is heavily influenced by the choice of β, but not entirely. More specifically, p if one revisits Eq. 85 and 87, it can be observed that the off-diagonal E and F -derived entries are scaled by sin( β 2 − β 2 (k · v)2 ), where the √ β scaling applied to the normalised k and v is explicitly taken out. As β approaches 0, the sin also follows, and thus off-diagonal non-linear integrations vanish. However, similarly, if the dot product k · v approaches 1, i.e., k and v become identical, the non-linear interaction also disappears. Therefore, one could completely control the level of non-linearity in ADPTNet layers by parametrically setting k · v to a value p:
k̂ =
(I − vv T )k ∥(I − vv T )k∥2
(114)
2
k ← (1 − p )k̂ + pv
5
Results
5.1
Lyapunov Spectrum
Lemmas 2 and 3 provide claims regarding the relationship between the parametrised "strength" of the systems’ non-linearity as expressed through β, the recurrent eigenvalue spectrum Λ and the effective timescales of the non-linear ADPTNet dynamics measured through its Lyapunov spectrum. This section provides empirical evidence for these specific claims while also testing their robustness to less strict parametrisation assumptions. 5.1.1
Untrained Orthogonal RNN Baseline
As briefly discussed in Section 2.5, one can parametrise the recurrent weights of RNNs with orthogonal matrices to ensure a degree of stability and long-term memory [6]. Before showcasing empirical evidence supporting the theoretical properties of ADPTNet from Section 4.1, it is helpful to examine Orthogonal RNNs as a baseline. Consider a vanilla recurrence rule for Orthogonal RNNs: xt = ReLU(Qxt−1 + But )
(115)
Where Q ∈ O(n), B ∈ Rn×n , x, u ∈ Rn . As long as there are no silent time steps, where the ReLU activation is ≡ 0, Orthogonal RNNs of this form are guaranteed to have at least the LLE αmax = 0. Crucially, there are no theoretical bounds for the rest of the Lyapunov spectrum in terms of how low they can be. Figure 9 visualises how this limitation materialises in practice. The Figure shows the entire Lyapunov spectrum13 of a single-layer untrained Orthogonal RNN with 128 hidden units. The network is driven by random inputs ∼ N (0, 1) of lengths ∈ {64, 4096, 16384}. It should be mentioned that the longer the input sequence is, the better the long-term dynamics of the system can be observed and Lyapunov spectrum estimated [50]. 13
Computed using a numerically stable algorithm as described in Lemma 3
35
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
Value (ln(Eigenvalue) / LE)
Here, the most relevant takeaway is that while the entire singular value spectrum of σ(Q) = 1, the Lyapunov spectrum only partially aligns with ln(σ(Q)) = 0. In fact, as can be seen in Figure 9, only approximately half of the spectrum is close to 0, tied to the sparsity of ReLU activations at initialisation. Moreover, while the network can be regularised during training, there are no parametrised fine-grained controls over the sparsity of the activations. Therefore, as is also observed in Figure 9, short-term time scales are effectively vanishing, with no controls over their decay rate (e−40 ). Even if a smoother activation function (e.g., sigmoid, GELU) were considered, one would still not have any more control over the lower bounds of the spectrum, and would also lose the LLE properties of ReLU. ln(All Eigenvalues) Lyapunov Spectrum
T=64
0
LE Bounds (Min/Max)
T=4096
0 5
LLE: -0.3064
10
LLE
5
LLE: -0.3209
10
30
20
15
25
20
30
25
35
40 0
20
40
60
80
100
120
140
40
LLE: -0.3221
10
15
20
T=16384
0
30 0
20
40
60
80
100
120
140
0
20
40
60
80
100
120
140
Index (Rank) Figure 9: Lyapunov Spectrum of Orthogonal RNN at Initialisation Each subfigure shows the estimated Lyapunov spectrum of a single-layer untrained Orthogonal RNN (Eq. 115) at initialisation with hidden size = 128 for a given random input of length ∈ {64, 4096, 16384}. The blue line shows the logarithm of the eigenspectrum of the orthogonal recurrent matrix, which = 1, and is only partially aligned with the Lyapunov spectrum of the system. 5.1.2
ADPTNet at Initialisation
This section examines the robustness of the claims regarding the parametric control over ADPTNet’s Lyapunov spectrum at initialisation. Here, untrained networks are driven by inputs that are randomly drawn ∼ N (0, 1). Unless otherwise specified, all networks tested used the configuration in Table 2. Algorithm 2 is used to compute the Lyapunov Spectra, extracting Jacobian singular values at each time step. Rot.-Aware ✗
Neg. Eigvals. ✗
p
γ
RoPE
D
Exp.
KV
(λmin , λmax )
(∆tmin , ∆tmax )
1 − λ̄2
✓
16
Low Rank
Ortho.
(2, 100)
(10−4 , 10−4 )
Table 2: Baseline Untrained ADPTNet Configuration Rotation-Awareness refers to the correction for the element-wise mismatch from the similarity transform (Eq. 105). Negative Eigenvalues can be included as mentioned in Section 4.2. The RoPE column indicates whether the custom RoPE embeddings with S4D-Inv parametrisation are included. Otherwise, no positional embeddings are used. D refers to the hidden state dimension. Matrix Exponential (Exp.) refers to whether the low-rank structure-aware formula (Section 4.4) or the Cayley Map is used (Section 2.5). γ refers to the input normalisation scheme used, EMA or LRU (Section 4.7). KV refers to the matrix parametrisation used for the K and V linear projections: baseline unconstrained (Kaiming initialisation), orthogonal, or LowParam (Section 4.3). Using the ZOH discretisation scheme (Section 4.2) eigenvalues λ̄i = e−λi ∆ti , where λi is ∼ U(λmin , λmax ) and ∆ti is obtained with linspace(∆tmin , ∆tmax , D). Sequence Length and Hidden State Size Figure 10 shows the effect of the input sequence length on the measured Lyapunov spectrum. Firstly, and most crucially, it should be noted that all Lyapunov exponents are well contained within the spread of the logarithms of the recurrent eigenspectrum, implicitly meeting the bounds set out in Lemma 2. Secondly, the hidden state size D of an ADPTNet layer does not affect the alignment between the observed Lyapunov spectrum and the parametrised eigenvalues. Finally, as expected, the input sequence length distorts the Lyapunov exponents measured from the network, with short sequences effectively "squeezing" the spectrum. Intuitively, there 36
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
is not sufficient long-term information to extract the most slowly evolving timescales, as by definition Lyapunov exponents become more accurate with T → ∞ (Section 2.3). Furthermore, in the longest sequence length configuration, αi ≈ ln(λ̄i ), supporting the claims in Lemma 3 as well. Robustness of K and V Linear Projection Parametrisation to Choice of β In Lemma 1, it is assumed that K and V are constrained to orthogonal matrices which results in ∥K∥2 = ∥V ∥2 = 1. This is helpful because it allows the norms of several derivative terms to cancel and produce the final 8β λ̄max factor, which is independent of network width, state x norm, etc. However, in practice, constraining dense and full-rank linear layer weights to the orthogonal manifold is computationally expensive and slow, while also effectively halving the degrees of freedom ( D(D+1) ) compared to 2 dk dv unconstrained GL(D) matrices. Assuming an arbitrary invertible parametrisation , the norms of the dx and dx terms dM inside dx x 2 in Eq. 71 become:
∥Kx∥2 ≥ σmin (K) ∥x∥2 ⇒
∥K∥2 ∥x∥2 σmax (K) ∥x∥2 1 1 ≤ ⇒ ≤ = κ(K) ⇒ ∥Kx∥2 σmin (K) ∥x∥2 ∥Kx∥2 σmin (K) ∥x∥2
dk dv dM x ≤ βλmax + ∥x∥2 dx dx dx 2 2 ∥K∥2 ∥x∥2 ∥V ∥2 ∥x∥2 ≤ 4βλmax + ≤ 4βλmax (κ(K) + κ(V )) ∥Kx∥2 ∥V x∥2
(116)
(K) Where κ(K) = σσmax denotes the condition number of K. In other words, if K or V are ill-conditioned at min (K) initialisation or become ill-conditioned during training, the bounds on the Lyapunov spectrum of the network also become looser. As with other random matrices, Kaiming-initialised dense linear weights have condition numbers that, on average, scale linearly with network width, but may have large outliers [45]. Furthermore, if the weights are low-rank or non-invertible, the bounds are undetermined < ∞. Using the initialisation scheme from Section 4.3, a LowParam linear layer would in theory have an undefined condition number14 . Either unpredictable or undefined, one cannot directly predict the conditioning of the two alternative orthogonal parametrisations considered here.
Regardless, Figure 11 shows that ADPTNet dynamics are generally robust to this choice. More specifically, for a lower β = 0.125, all parametrisation schemes display Lyapunov spectra ln(λ̄min ) ≤ αi ≤ ln(λ̄max ). With β = 1, all parametrisations are still generally close to the original ln(Λ̄). A more significant gap between orthogonal and alternatives becomes apparent with a large, and perhaps unrealistic, β = 1000. In terms of the most quickly decaying timescales, while orthogonal constraints show a ≈ −0.002 "undershoot", LowParam and dense have an order of magnitude higher discrepancies of < −0.04 and < −0.01, respectively. Moreover, both alternatives technically display chaotic dynamics since their LLEs are > 0, while the orthogonal-constrained configuration is still stable. It is important to note, however, that all configurations are still well within the bounds prescribed by Lemma 2. In the β = 1 case, the orthogonal parametrisation for which they are guaranteed, the theoretical bounds would be roughly −∞ < αi < ln(λ̄max + 8 ∗ λ̄max ), which are evidently far outside the observed behaviour and thus respected, and likewise for the more extreme case β = 1000. One could argue that this provides evidence that mitigation methods such as spectral normalisation are not necessary for ADPTNet, especially in the more realistic range of β ≪ 1.
14
The low-rank component would be all-0 at initialisation. The remaining Kronecker component is a product between a dense random matrix and an all-1 matrix, which results in an undefined division by zero when computing the condition number.
37
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
Lyapunov Spectra Bounds vs. Eigenspectrum Bounds for No. Timesteps (T) × Hidden Size (D) Combinations T=64, D=16
0.000
T=4096, D=16
0.000
0.001
T=16384, D=16
0.000
0.001
0.001
0.002
0.002
0.003
0.003
0.004
0.004
0.005
0.005
0.002 0.003
Value (ln(Eigenvalue) / Lyapunov Exponent)
0.004 0.005 0.006 0.007 0.008 0
2
4
0.000
6
8
10
T=64, D=512
12
14
16
0.006
0
2
4
0.000
6
8
10
T=4096, D=512
12
14
16
0.006
0.002
0.002
0.004
0.004
0.004
0.006
0.006
0.006
0.008
0.008
0.008
0.010 0
100
200
300
400
500
ln(All Eigenvalues) Eigenvalue Bounds (Min/Max)
2
4
0.000
0.002
0.010
0
6
8
10
12
14
400
500
T=16384, D=512
16
0.010 0
100
200
300
Index (Rank)
LE Bounds (Min/Max) Lyapunov Spectrum
400
500
0
100
200
300
LE bounds fit within Eigenvalue bounds
Figure 10: Effect of Time Steps and Hidden Size on Lyapunov Spectrum Each subfigure in the grid represents the recurrent eigenvalues Λ and Lyapunov spectrum of untrained ADPTNet layers driven by random input of different sequence lengths T ∈ {64, 4096, 16384}. Each subfigure is obtained using a different randomly initialised network. Both D configurations have β = 0.125. It should be observed how all Lyapunov exponents αi ≈ ln(λi ) over long sequences, and are universally contained within the bounds of the recurrent eigenspectrum as shown by the yellowshaded areas.
38
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
Lyapunov Spectra Bounds vs. Eigenspectrum Bounds for Weight Parametrisations × Beta Combinations Low-Parameter, = 0.125
0.000
Low-Parameter, = 1
0.000
Low-Parameter, = 1000 0.00
0.002
0.002
0.004
0.02
0.004
0.03
0.006
0.006
0.008
Value (ln(Eigenvalue) / Lyapunov Exponent)
0.01
0.04 0.05
0.008 0
2
4
0.000
6
8
10
12
Orthogonal, = 0.125
14
16
0
2
4
0.000
0.001
0.001
0.002
0.002
0.003
0.003
0.004
0.004
0.005
0.005
0.006
0.006
0.007
0.007
0.008
6
8
10
Orthogonal, = 1
12
14
16
2
0.000
4
6
8
10
12
Baseline Linear, = 0.125
14
16
2
0
2
0
2
4
6
8
10
12
14
16
4
6
8
10
12
14
16
4
6
8
10
12
14
16
0.000
Orthogonal, = 1000
0.002 0.004 0.006 0.008
0.008 0
0
0.010 0
2
4
0.000
6
8
10
12
Baseline Linear, = 1
14
16
Baseline Linear, = 1000
0.0000
0.001
0.0025
0.002
0.002
0.0050
0.003 0.004
0.0075
0.004
0.0100
0.005 0.006
0.0125
0.006
0.007
0.0150
0.008
0.008 0
2
4
6
8
10
12
14
16
ln(All Eigenvalues) Lyapunov Spectrum
0.0175 0
2
4
6
8
Index (Rank)
Eigenspectrum Bounds (Min/Max) LE Bounds (Min/Max)
10
12
14
16
LE bounds fit within Eigenspectrum bounds LE bounds do not fit within Eigenspectrum bounds
Figure 11: Interaction of KV Weight Parametrisation and Choice of β Each subfigure shows the Lyapunov spectrum and ln(Λ̄) for random ADPTNet networks at initialisation with different parametrisation schemes. The LowParam setting uses rank = 2 and Na = Nb = 4. Inputs are driven by T = 4096 random inputs. Notably, only for very large β = 1000 do the dynamics of orthogonal constraints start to diverge significantly from the baseline dense linear with Kaiming initialisation and LowParam.
Effect of Matrix Exponential Function The proofs for the Lemmas 2 and 3 rely on the assumption that the matrix exponential required for mapping skew-symmetric matrices to the orthogonal manifold is computed using the Cayley Map. However, in practice, the Cayley map is prohibitively expensive in terms of compute and memory and, thus, an alternative that takes into account the low-rank structure of the skew-symmetric mapping for the k ⊗ v outer product is preferable (see Section 4.4). However, since the Cayley map is only a truncated approximation of the matrix exponential, one could argue that switching exponential functions may affect the theoretical claims in Section 4.1. Figure 12 provides evidence to suggest that the ADPTNet recurrence is resilient to this switch. It can be observed that the low-rank-aware exponential function produces very close, if not nearly identical, Lyapunov spectra relative to the parametrised ln(Λ̄). Even in the large β = 1000 setting, the two functions produce "undershoots" of similar proportions ≈ −0.002. 39
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
Lyapunov Spectra Bounds vs. Eigenspectrum Bounds for Matrix Exp Function × Beta Combinations Cayley Map, = 0.125
0.000
0.001
0.002
0.002
0.003
0.003
Value (ln(Eigenvalue) / Lyapunov Exponent)
0.001
0.004
0.004
0.005
0.005
0.006
0.006
0.007
0.007
0.008
0.008 0
2
0.000
4
6
8
10
12
Low Rank Matrix Exp, = 0.125
14
16
0.001 0.002
0.003
0.003
0.004
0.004
0.005
0.005
0.006
0.006
0.007
0.007
0.008
0.008 2
4
6
8
10
12
0.004 0.006 0.008 0.010 2
0.000
0.002
14
ln(All Eigenvalues) Lyapunov Spectrum
16
Cayley Map, = 1000
0.000 0.002
0
0.001
0
Cayley Map, = 1
0.000
4
6
8
10
12
Low Rank Matrix Exp, = 1
14
16
0
2
0
2
0.000
4
6
8
10
12
14
16
4
6
8
10
12
14
16
Low Rank Matrix Exp, = 1000
0.002 0.004 0.006 0.008 0.010 0
2
4
6
8
Index (Rank)
Eigenspectrum Bounds (Min/Max) LE Bounds (Min/Max)
10
12
14
16
LE bounds fit within Eigenspectrum bounds LE bounds do not fit within Eigenspectrum bounds
Figure 12: Effect of Matrix Exponential Function Choice The two rows of the grid show the two different matrix exponential functions tested: the Cayley map and the low-rank structure-aware exponential, as described in Section 4.4. All randomly initialised networks in the grid are driven by T = 4096-long random inputs. It is apparent that for any scale of β, both exponentiation methods produce very similar behaviour in terms of Lyapunov spectra.
Effect of Recurrent Condition Number Since the ADPTNet recurrent matrices Mt are symmetric positive definite by construction, their singular values are ≡ Λ̄. Therefore, their 2-norm condition number is = λ̄λ̄max . Figure 13 shows min how this ratio does not influence the eigenspectrum’s relationship to the observed Lyapunov exponents. While the ≈ 1.01 setting shows slightly higher discordance between the parametrised and measured spectra, compared to the other condition numbers, alignment appears to be driven mostly by β. As in Figure 11 and 12, a large β = 1000 causes more "undershoot" compared to the β = 1. This evidence suggests that regardless of the spread of parametric timescales in the network (e.g., if the spread increases during training), that should not affect the resulting Lyapunov spectrum. 40
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
Lyapunov Spectra Bounds vs. Eigenspectrum Bounds with Condition Number × Combinations = 1, Cond Num = 9.87
= 1, Cond Num = 4.95
0.0
= 1, Cond Num = 1.01
0.8
0.002
1.0
0.5
1.2 1.0
0.004
Value (ln(Eigenvalue) / Lyapunov Exponent)
1.4 0.006
1.6
1.5
1.8 0.008
2.0
2.0
2.2 0
20
40
60
80
100
120
140
0.010 0
20
= 1000, Cond Num = 9.87
40
60
80
100
120
140
0
20
40
= 1000, Cond Num = 4.95
0.0
0.75
0.5
1.00
80
100
120
140
120
140
0.002
1.25
1.0
60
= 1000, Cond Num = 1.01
0.004
1.50 0.006
1.75
1.5
2.00
2.0
0.008
2.25 0.010
2.5
0
20
40
60
80
100
120
140
ln(All Eigenvalues) Lyapunov Spectrum
0
20
40
60
Index (Rank)
Eigenspectrum Bounds (Min/Max) LE Bounds (Min/Max)
80
100
120
140
0
20
40
60
80
100
LE bounds fit within Eigenspectrum bounds LE bounds do not fit within Eigenspectrum bounds
Figure 13: Effect of Recurrent Weights Condition Number Each subfigure shows the effect of different ratios λ̄λ̄max min on the aligment between the eigenvalue and Lyapunov spectra. Each random network with D = 128 takes in T = 4096 long random inputs. There are three condition number configurations, obtained by setting λmin = λmax = 1 and (∆tmin , ∆tmax ) ∈ {(10−2 , 2.3), (0.7, 2.3), (10−3 , 10−2 )}. It should be noted that alignment is driven more by the choice of β than by the condition number. 5.1.3
Trained ADPTNet on Copy Memory Task
Neg. Eigvals. ✗
γ 1 − |λ̄|
RoPE ✓
D 128
Exp.
KV
Low Rank
Dense
(λmin , λmax ) 0.5
(∆tmin , ∆tmax ) −3
(10
−1
, 10
)
β 0.0125
Table 3: Baseline Trained ADPTNet Configuration Column headers follow the same conventions as in Table 2. Notably, γ is parametrised as a function of λ̄ throughout training, not just initialised to the value. While Rotationalawareness is omitted from the table, it is assumed to be absent from all models tested. KV Dense refers to unconstrained dense linear layers with Kaiming initialisation. Copy Memory Task This section explores how robust the alignment between parametrised recurrent eigenvalues and observed Lyapunov exponents is throughout training epochs. Table 3 shows the baseline configuration used for all models tested, unless otherwise specified. All networks are trained on the Copy Memory Task [92]. The task consists of a sequence of target tokens, followed by a series of distractors, and finally output prompting tokens to elicit recollection of the target tokens (Figure 14). Following Romero et al. [150], all models are trained to store 10 tokens in memory from a vocabulary of integers ∈ {1, 2, . . . , 8}. The task difficulty is modulated by the number of distractors (e.g., 100, 1000), since the loss computed from the final 10 time steps of the recall must be backpropagated through time over increasing numbers of time steps. This section emphasises the behaviour of ADPTNet Lyapunov spectra over training iterations, rather than raw accuracy. Typically, delayed copying is not intended as a challenge in itself, but rather, it is a pass/fail test of whether information can travel across time in sequence models. Here, passing is quantified as achieving > 90% accuracy within at most 100 training epochs with a learning rate of η = 0.004, cosine scheduling, and Adam optimiser [110]. All reported models pass this standard. 41
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
Figure 14: Copy Memory Task All models in this section are trained on the Copy Memory Task, which measures long-term recall ability. Target tokens (integers ∈ {1, 2, . . . , 8}) need to be retrieved in the correct order. The loss is then computed on the final recall window and backpropagated through the distraction steps.
Effect of β on Trained ADPTNet Figure 15 shows how the choice of β affects ln(Λ̄) and Lyapunov exponents over training iterations. Firstly, it can be observed that a lower β induces a sharper slope and increased spread in the Λ̄ distribution. Concretely, at epoch 5, ln(λ̄min ) is an order of magnitude lower for β = 0.0125 compared to β ∈ {0.125, 1}, going from −2.5 compared to −0.35 and −0.14 respectively. The trend continues up to the final epoch 50, with higher β progressively increasing ln(λ̄min ) orders of magnitude from ≈ −7 to ≈ −2.5 and ≈ −0.35, for β = 0.0125, 0.125, and 1, respectively. Intuitively, a lower β implies a lower rotation angle in the similarity transform QΛ̄QT , where the plane of rotation is dictated by k ⊗ v. Therefore, to produce more non-linear behaviour, the network may learn to spread the eigenvalues further apart, achieving higher output variance with a limited rotation angle β. Regarding the Lyapunov spectrum, once again, increasing β decreases alignment with the ln(λ̄) distribution. Both β = 0.0125 and β = 0.125 have Lyapunov exponents within the log-spread ln(Λ̄), evidently meeting the bounds in Lemma 2. For β = 1, the LLE is positive, so the dynamics are slightly chaotic. However, the LLE is still < 0.5, and the theoretical bound ln(λ̄max + 8β λ̄max ) ≈ ln(9) = 2.19, which is evidently larger than the observed value. Therefore, the theoretical bounds hold up over training and are robust to the choice of β.
KV-Param
Dense
Orthogonal
LowParam
5
LLE = 0.00344 Bound ≈ 0.0952
LLE = 0.00048 Bound ≈ 0.095
LLE = 0.00779 Bound ≈ 0.095
25
LLE = 0.00835 Bound ≈ 0.0953
LLE = 0.01043 Bound ≈ 0.0953
LLE = 0.00723 Bound ≈ 0.0953
50
LLE = 0.0108 Bound ≈ 0.0953
LLE = 0.0098 Bound ≈ 0.0953
LLE = 0.01026 Bound ≈ 0.0953
Epoch
Table 4: LLEs and their Maximum Theoretical Bounds Each cell in the table contains the LLE associated with each KV-parametrisation scheme along with the theoretical bound on it obtained with Lemma 2. While the bounds ln(λ̄max + 8β λ̄max ) are technically specific to orthogonal weights, it can be observed that all parametrisations are well within them.
Effect of KV Parametrisation on ADPTNet Trained Dynamics As shown for ADPTNet initialisation (Fig. 11), Figure 16 shows the same effect persisting over training epochs. The Lyapunov spectra closely track the recurrent eigenspectra, with the lower limits of all Lyapunov exponents remaining above their parametrised counterparts across all configurations and epochs. It can, however, also be observed that the LLEs do "overshoot" ln(λ̄max ) by ≪ 0.05, especially towards the end of training. The mismatch remains within the bounds prescribed by Lemma 2 as detailed in Table 4. 42
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
Evolution of Lyapunov Spectra vs. Eigenspectrum Bounds across Epochs: Choice of Beta: 0.0125 | Epoch: 5
0.0
Beta: 0.0125 | Epoch: 25
0
0.5 1.0
1
1
2
2 3
3
4
4
1.5
5
5
2.0
6
6 2.5 0
20
0.00
40
60
80
100
Beta: 0.125 | Epoch: 5
7
120
7 0
20
40
60
80
100
Beta: 0.125 | Epoch: 25
120
ln(| |) / LE)
0
20
0
20
0
20
0.0
0.0
0.05
40
60
80
100
120
40
60
80
100
120
40
60
80
100
120
Beta: 0.125 | Epoch: 50
0.5
0.5
0.10
Beta: 0.0125 | Epoch: 50
0
1.0
0.15
1.0 1.5
0.20 1.5
0.25 0.30
2.0 2.5
2.0
0.35 0
20
0.00
40
60
80
Beta: 1.0 | Epoch: 5
100
120
0
20
40
60
80
Beta: 1.0 | Epoch: 25
100
120
0.00
0.02 0.04 0.06
0.00
0.05
0.05
0.10
0.10 0.15
0.15
0.08 0.10
0.20
0.12
0.25
0.14
0.30 0
20
40
60
80
100
120
Beta: 1.0 | Epoch: 50
0.20 0.25 0.30 0.35 0
20
40
60
80
100
120
Index (Rank) ln(| |) Lyapunov Spectrum
Eigenspectrum Bounds LE Bounds
LE within Eig Bounds LE outside Eig Bounds
Figure 15: Effect of β over ADPTNet Trained Dynamics Each row represents a trained ADPTNet model with a given β, and the columns are snapshots from different epochs (epoch 50 is the final one). The networks are trained on the Copy Memory task with 44 distraction tokens for a total sequence length of T = 64. The initial learning rate is η = 0.005. It should be observed how the bounds from Lemma 2 hold up over training and how lower β induces a higher eigenvalue spread.
43
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
Evolution of Lyapunov Spectra vs. Eigenspectrum Bounds: KV Parametrisation Baseline | Epoch: 5
Baseline | Epoch: 25
0.00
0.00
0.025
0.02
0.05
0.050 0.075
0.04
0.10
0.100
0.06
0.125
0.15
0.150
0.08 LLE: 0.00344 ln( max): -0.00016
0.10 0
20
0.00
Value (ln(|Eigenvalue|) / LE)
Baseline | Epoch: 50
0.000
40
60
80
100
Orthogonal | Epoch: 5
120
0.02 0.04 0.06 0.08 LLE: 0.00048 ln( max): -0.00027
0.10 0
20
40
60
80
LowParam | Epoch: 5
100
120
0.175 0.200
LLE: 0.00835 ln( max): -0.00002 0
20
40
60
80
100
Orthogonal | Epoch: 25
0.20
120
0
0.000
0.00
0.025
0.02
0.050
0.04
0.075
0.06
0.100
0.08
0.125
0.10
0.150 0.175
LLE: 0.01043 ln( max): -0.00001 0
20
40
60
80
100
LowParam | Epoch: 25
120
0.14
0.00
0.00
0.025
0.05
0.05
0.050
0.10
0.10
0.15
0.100
0.175 0.200
LLE: 0.00779 ln( max): -0.00035 0
20
40
60
80
100
120
80
100
120
LLE: 0.00980 ln( max): -0.00000 0
20
40
60
80
100
LowParam | Epoch: 50
120
0.25
0.30 0.35
60
0.20
0.25
0.150
40
Orthogonal | Epoch: 50
0.15
0.20
0.125
20
0.12
0.000
0.075
LLE: 0.01080 ln( max): -0.00000
LLE: 0.00723 ln( max): -0.00001 0
20
40
60
80
100
LLE: 0.01026 ln( max): -0.00000
0.30
120
0
20
40
60
80
100
120
Index (Rank) ln(|Eigenvalues|) Lyapunov Spectrum
Eigenspectrum Bounds LE Bounds
LE within Eig Bounds LE outside Eig Bounds
Figure 16: Effect of KV-parametrisation over Lyapunov Spectrum During Training Each subfigure in the grid shows the Lyapunov spectrum and the ln(Λ̄) for each KV parametrisation option. The networks are trained on the Copy Memory task with 108 distraction tokens and a total of T = 128 time steps. Baseline denotes dense and unconstrained linear layers, and for the LowParam configuration, rank = 32, Na = 32 and Nb = 4. As also reinforced in Table 4, all configurations have Lyapunov spectra within the bounds from Lemma 2.
Effect of Including Negative Eigenvalues As mentioned in Section 4.2, for the purposes of this study, ADPTNet recurrent eigenvalues are constrained parametrically to |λ̄i | ≤ 1 by double exponentials, following Gu et al. [76] (Eq. 77). Because the M = QΛ̄QT structure creates symmetric matrices, even with negative λ̄i , singular values are still σi (M ) = |λi |. Consequently, the singular values σ̃i of the recurrent Jacobian Jt are invariant to this choice, and, thus, alignment with the Lyapunov spectrum should not be affected. Figure 17 provides empirical evidence supporting this. In particular, Subfigure 17a shows that trained networks with either completely positive eigenspectra or negative eigenvalues show no visible differences in alignment with their respective Lyapunov spectra. However, building on the evidence from Figure 15, the networks do display different spectral distributions. 44
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
In Figure 17a, the fully-positive eigenspectrum configuration registers a |λ̄min | ≈ e−7 , orders of magnitude lower than the "negative" configuration’s |λ̄min | ≈ e−0.04 . The LLE of the negative eigenvalue condition is slightly positive, which indicates weakly chaotic dynamics, compared to the stability of the fully positive condition. Furthermore, Subfigure 17b shows how both configurations maximise eigenvalue spread over training epochs. For the "positive" setting, maximising spread means pushing eigenvalues to the boundaries of [0, 1]. The "negative" setting can instead cluster eigenvalues around the ±1 extremities. As highlighted in Section 4.5.2, the topological conjugates QΛ̄QT live within the same unitary orbit, which means that the distance between M (x) and M (y) for any given x and y is bounded by the eigenvalue spread: ∥M (x) − M (y)∥2 = Q(x)Λ̄Q(x)T − Q(y)Λ̄Q(y)T
2
≤ |λ̄max − λ̄min |
(117)
Therefore, as argued before, ADPTNet layers may be learning to maximise variance, or the "sharpness of the turns" (Fig. 4), possible at each time step.
Effect of Negative Eigenvalues on Lyapunov Spectra vs. Eigenspectrum Bounds Positive Eigenvalues
Negative Eigenvalues
0
0.00
1 0.01
ln(| |) / LE
2 3
0.02
4 5
0.03 6 7
0.04 0
20
40
60
Index
80
100
ln(| |) Lyapunov Spectrum
120
0
20
Eigenspectrum Bounds LE Bounds
40
60
Index
80
100
120
LE within Eig Bounds LE outside Eig Bounds
(a) Fully Trained Singular Value and Lyapunov Spectra Evolution of Eigenvalue Spread (| max
min|) During Training
2.00 1.75
Spread Magnitude
1.50 1.25 1.00 0.75 0.50 0.25 0.00
Standard Model Negative Eig Model 20
40
Epoch
60
80
100
(b) Eigenvalue Spread
Figure 17: Effect of Including Negative Eigenvalues of Trained ADPTNet Spectra. Subfigure 17a shows the different distributions of log-singular values and Lyapunov exponents for a network with a fully positive eigenspectrum and a second network with 32 eigenvalues constrained to be negative throughout training. Both networks are trained on the Copy Memory task with 1004 distraction tokens (total sequence length T = 1024). The "positive" model reaches 95% test accuracy at Epoch 100, while the "negative" model reaches 90%. The distribution snapshots are taken at the end of training (Epoch 100). Subfigure 17b shows the evolution of the eigenvalue spread for the two model configurations throughout training. It should be noted that the Lyapunov spectra are close to the log-singular values for both settings. However, the distribution has a lower minimum for the "positive" case. It can also be observed that both spreads tend to their respective maximums over training. 45
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
(t)
∆i Distribution Changes over Training Epochs So far, all observed Lyapunov spectra strongly align with their parametrised ln(|Λ̄|) counterparts. Since only the stable Lyapunov spectrum computation algorithm (Algorithm 2) has (t) been used so far, all evidence thus far points toward confirming Lemma 3. In Lemma 3, the notation of ∆i = σ̃i − σi was introduced, describing the difference (i.e., perturbation) between the singular values of the recurrent Jacobian Jt and the parametrised |Λ̄|. As noted in Corollary 2, if the distribution of such perturbations over sequence length has mean = 0, then they effectively cancel out in the long term, and the Lyapunov spectrum matches ln(Λ̄). (t)
Figure 18 provides evidence supporting this claim. As shown, throughout training, the distribution of ∆i is centred around 0, leading to the strong match between log-singular values and Lyapunov exponents. However, the variance of the perturbations does increase between Epoch 5 and the end of training (Epoch 100). Besides β, the size of the perturbations also depends on the various k · v dot products inside the Jt formulation. Therefore, as training progresses, ADPTNet layers may also learn to control the distance between k and v and the plane of rotation in Q more precisely. This further supports the argument that by increasing eigenvalue spread, the network dynamics tend towards increased non-linearity throughout training.
Distribution throughout Training 8000
Epoch 5
7000
Epoch 100
=1.02e-06 =2.48e-04 max| |=3.91e-03 =-1.75e-07 =5.40e-04 max| |=1.31e-02
Probability Density
6000 5000 4000 3000 2000 1000 0 0.0020
0.0015
0.0010
0.0005
0.0000
(t) Pertubation ( (t) i = i
(t)
0.0005 i)
0.0010
0.0015
0.0020
(t)
Figure 18: ∆i Distribution Throughout Training ∆i perturbations are computed at each time step using as the (t) difference σi (Jt ) − λ̄i = σ̃i − λ̄i , and averaged over sequence length for a Copy Memory task sample. µ denotes the mean over time steps, while σ represents the standard deviation. The perturbations are collected from the "positive" model in Figure 17. The key trends to observe are that for both Epoch 5 and Epoch 100 (final epoch) the mean of the (t) perturbation distribution is 0. In addition, ∆i variance increases between the two epochs visualised. 5.2
Parallelisation with Conv and Forward DEER
As in Section 5.1, this section is split between analysing single-layer ADPTNet models at initialisation and throughout training. Using the insights from Section 4.5, the goal is to assess the degree to which the approximation errors of Conv and Forward-DEER compared to Quasi-DEER impact parallelisation performance in practice. 5.3
Conv and Forward DEER Computational Overhead
Before analysing convergence properties for the Jacobian-free DEER variations proposed in this study, it is important to underscore the limited computational overhead they incur. Here, computational overhead is defined as additional Multiply-Accumulate (MAC) operations required to obtain At (Section 4.5), after computing Mt xt−1 for the forward 46
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
dynamics of the network. This does not include the actual cost of computing DEER iteration parallel scans until convergence, as those vary by variable factors such as ADPTNet layer dynamics or Scale-ELK damping (see Section 4.5.3). In Table 5, it can be seen that since Conv DEER consists of At = Λ̄, no additional floating point operations are required, as the discretised eigenspectrum is already necessary to compute Mt xt−1 in the ADPTNet recurrence and thus can be reused. Forward DEER introduces only a linearly scaling MAC overhead. Evidently, Quasi-DEER would also incur a linearly scaling cost, with a higher constant owing to the terms in dM dx x in Eq. 65. DEER Approx.
MACs
Conv
0
Forward
20 ∗ D
Table 5: Computational Overhead for Jacobian-Free DEER Variations Conv-DEER does not require any additional MACs to obtain its At terms, while Forward DEER scales linearly. The MAC count estimate for Forward-DEER is obtained using the torchprofile library and the implementation of Eq. 95 in Appendix B. Here, D refers to the model hidden state size.
5.3.1
ADPTNet Parallelisation at Initialisation Neg. Eigvals.
γ
RoPE
D
Exp.
KV
(λmin , λmax )
T
✗
1 − |λ̄|
✓
64
Low Rank
Dense
0.5
1024
Table 6: Baseline Untrained ADPTNet Configuration for DEER Parallelisation Ablations Overall, the baseline configuration here is close to the trained model configuration in Section 5.1.3. T refers to the total number of time steps in the random sequences fed into the networks as inputs. To collect convergence metrics, the networks are then simulated with these random drivers following Gonzalez et al. [68]. Similar to Section 5.1.2, this section explores the DEER parallelisation characteristics of single untrained ADPTNet layers at initialisation. Inputs are once again random signals drawn ∼ N (0, 1). The default configuration for all experiments is in Table 6. Eigenvalue Spread and β Ablation In Section 4.5.2 it is hypothesised that increasing β decreases the size of the basin of quadratic convergence for DEER. In practice, this would mean that the increased non-linearity of a larger β should also increase the number of Newton iterations required for DEER convergence. Furthermore, considering the various local approximation error terms (Table 1), one could also predict that increasing λ̄max and the spread |λ̄max − λ̄min | should also increase the required iterations for convergence, and also increase the disparities between Quasi-DEER and the Jacobian-free alternatives proposed here (Conv and Forward-DEER). Figure 19 provides empirical evidence supporting these hypotheses. Firstly, Subfigure 19a shows how decreasing ∆tmin (i.e., increasing λ̄max = e−λmin ∆tmin ) also increases the number of DEER iterations required for all variants tested. Furthermore, the increase in iterations required for convergence is larger for larger β, and a larger β in itself raises the total number of iterations required for the same ∆tmin . In other words, a larger β increases the non-linearity of the system, and combined with a larger λ̄max , the bounds for the LLE also increase (Lemma 2). These two factors (increased non-linearity and larger LLE) are already known to slow down DEER convergence [69]. However, unlike other non-linear RNNs, the evidence in Figure 19 suggests that ADPTNet can parametrically control them and, in turn, DEER convergence speed. It is worth noting that the differences between Quasi, Conv, and Forward DEER become more pronounced with lower ∆tmin and larger β. Namely, for ∆tmin = 10−3 and β = 100, Forward and Conv-DEER have progressively higher penalties compared to Quasi-DEER, reflecting their increasingly aggressive truncations of the terms making up Quasi-DEER’s diag(Jt ) (Table 96). Additionally, since ∆tmax is held constant between configurations, the eigenvalue spread increases, which may contribute to worse convergence in the Jacobian-free methods. Secondly, Subfigure 19b shows a complementary story to Subfigure 19a. Holding ∆tmin constant and decreasing ∆tmax effectively decreases eigenvalue spread, while also stretching all timescales of the system for longer-term memory. Subfigure 19b shows how this decrease in spread leads to faster average convergence across all DEER variations. In fact, it should be noted that, for the final setting where ∆tmax = ∆tmin = 10−3 , the layers are actually linear, and thus only require a single DEER iteration by definition (see Section 4.5). This is reflected in the average convergence for 47
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
Forward and Quasi DEER. However, the mean iterations required for Conv-DEER on this configuration are slightly > 1, which suggests that numerical precision issues can emerge even in trivial cases. Finally, Subfigure 19c adds nuance to the differences in convergence behaviour between the DEER versions. Its two histograms show the distribution of DEER iterations required for convergence as different random input sequences are fed to the network for simulation. The vertical lines show the median number of iterations required for convergence across all seeds for each algorithm. For β = 1, all algorithms have comparable median iterations to convergence, but Conv-DEER shows higher maximum outliers than Forward and Quasi DEER. Differences are further amplified by β = 100, with significantly larger outliers for all algorithms, but also a clearer ordering for both mean and median convergence between DEER versions: Quasi < Forward < Conv. This is unsurprising, as it again confirms the accumulation of approximation errors caused by truncating Jacobian diagonal constituent terms. Newton Iterations for And tmin Ablations 256 128
8
Mean Newton Iterations
Mean Newton Iterations (Log Scale)
Newton Iterations for and Eigenspectrum Spread Ablations
10 0.125 1.0 100.0 Algorithm Quasi DEER Conv DEER Forward DEER
64 32 16 8
6
4
2
4
1 0
2
10 2
10 3
tmin
0.01
0.005
(a) Increasing Long-Term Memory
0.001
tmax
(b) Decreasing Eigenvalue Spread
Distribution of Newton Iterations Required for Convergence = 1.0 Algorithm Stats
14
Algorithm Stats
Quasi DEER ( =30.5, =9.7) Quasi DEER Median: 27 Conv DEER ( =33.5, =13.5) Conv DEER Median: 29 Forward DEER ( =31.5, =10.0) Forward DEER Median: 28
12 10
Frequency
= 100.0 Quasi DEER ( =125.5, =119.9) Quasi DEER Median: 55 Conv DEER ( =236.0, =223.1) Conv DEER Median: 228 Forward DEER ( =187.7, =161.4) Forward DEER Median: 165
8 6 4 2 0
30
40
50
60
Newton Iterations
70
80
0
100
200
300
400
500
Newton Iterations
600
700
(c) DEER Iteration Distribution
Figure 19: Eingspectrum Effects on DEER Convergence. All Subfigures are produced using 15 random samples for each configuration. For Subfigure 19a, ∆tmin controls the longest time scale of the layer through λ̄max and takes values ∈ {10−2 , 10−3 }. ∆tmax controls the shortest timescale through λ̄min and is constant at 10−1 . The subfigure highlights how increasing the eigenvalue spread and β leads to more DEER iterations across all algorithms tested. Subfigure 19b holds ∆tmin fixed at 10−3 , while varying ∆tmax ∈ {0.01, 0.005, 0.001}. This shows that DEER iterations for long-range, slowly decaying dynamics can be reduced by reducing the eigenvalue spread along with β. Subfigure 19c uses the same samples and models as Subfigure 19a and shows that Forward, and to a great extent Conv DEER, suffer from more high-iteration-count outliers. 48
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
DEER Iteration Scaling with Sequence Length The total number of time steps in the trajectory T appears explicitly in the expression for the size of the quadratic convergence zone (Eq. 88 adapted from Gonzalez et al. [69]). This formalises the intuition that a longer sequence would require more DEER iterations to converge. Figure 20 clearly reflects this notion as well. Furthermore, it shows that the mean convergence iterations for Conv and Forward DEER largely match those of baseline Quasi-DEER, with a slight sign of upward divergence on the longest T = 16K setting for Conv DEER. Focusing on the maximum outliers for each algorithm, however, one can notice that while they increase with a sharper slope than the mean for all DEER variants, they are universally higher for the Jacobian-free approximations. Still, they are in the same order of magnitude.
Newton Iterations Required for Sequence Length T Quasi DEER Conv DEER Forward DEER Max Outlier
18
Mean Newton Iterations
16 14 12 10 8 6 4
512
1024
4096
Sequence Length T
8192
16384
Figure 20: Scaling with Sequence Length This figure shows the averaged number of iterations required for convergence over 15 random samples of varying length T ∈ {512, 1204, 4096, 8192, 16384}. β = 0.125 , ∆tmin = 10−4 , and ∆tmax = 10−1 for all sequence length configurations. The main takeaway is that, unsurprisingly, the number of Newton iterations needed for convergence grows with sequence length. However, importantly, the average number of convergence iterations is the same between Conv, Forward, and Quasi DEER.
Block-Wise DEER Simulation As highlighted in Section 4.5.4, by construction, all DEER algorithms suffer from a degree of redundant computation. This resource waste can be mitigated by introducing sequential computation between blocks of parallel DEER execution. Figure 21 visualises the trade-off that occurs. Subfigures 21a and 21b show how increasing the number of blocks decreases the memory consumption at the cost of increasing simulation wall-clock time. Moreover, for both sequence lengths shown, memory savings appear to plateau, with seemingly exponentially diminishing returns, while simulation times increase steadily and linearly. Subfigure 21c displays how, by increasing the number of blocks and implicitly reducing each block’s length, the number of DEER iterations required to converge, per block, also decreases. Together with Subplots 21a and 21b, this suggests that, in this case, the savings in iterations per block do not outweigh the slowness of sequential processing. Finally, Subfigure 21d offers more detail into the relative performance improvements of Conv over Forward DEER. Most importantly, Conv DEER reduces execution time by > 25% compared to Forward DEER. Still, mirroring the general trend seen in Subfigures 21a and 21b, peak memory allocation advantages vanish with increased block counts. Since both use the same implementation of parallel scans15 , the efficiency gains are largely driven by the lower computational overhead in computing At terms for Conv DEER (Table 5). It is important to emphasise that these findings are strongly tied to PyTorch idiosyncrasies. Danieli et al. [33] provides a more in-depth analysis of hardware I/O-aware optimisations for DEER-like algorithms.
15
Parallel associative scans are computed in PyTorch using https://github.com/proger/accelerated-scan
49
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
Conv-DEER Forward-DEER
800
Conv-DEER Forward-DEER
1200
B=32
B=32
600
B=32 B=16
Execution Time (ms)
Execution Time (ms)
1000
B=16
400
B=8 B=8 B=4
200
B=32
800 600 400
B=4 B=2
24
26
28
B=8 B=8 B=4
200
B=2
B=1
B=1
30
32
Peak GPU Memory (MB)
34
36
38
30
(a) Memory vs Speed (T = 4096)
B=2
B=1
B=2
B=1
35
40
45
Peak GPU Memory (MB)
50
55
(b) Memory vs Speed (T = 8192) Algorithm & T
8
B=4
0
0 22
B=16 B=16
1.0
Algorithm Conv DEER Forward DEER T 4096 8192
6
1.00
0.96
0.92
1.00
0.82
0.8
Ratio (Conv / Forward)
Newton Iterations
7
Time Ratio (Conv/Fwd) Memory Ratio (Conv/Fwd) Equal Performance (1.0)0.88
0.74
0.73
0.72
0.70
0.73
0.73
0.6
5
0.4
4 3
0.2
2 1 2
4
8
16
Block Size
0.0
32
(c) Iterations Required for Convergence
1
2
4
Block Size
8
16
32
(d) Conv vs Forward DEER Performance
Figure 21: Block-Wise DEER Trade-Offs Subfigures 21a and 21b show the wall clock time and peak GPU memory utilisation depending on the different number of sequential blocks (B ∈ {1, 2, 4, 8, 16, 32}) used for Conv/Forward DEER simulation. A random input is sampled for each sequence length T ∈ {4096, 8192}, along with a randomly initialised ADPTNet layer. For all configurations, β = 0.125, ∆tmin = 10−4 , and ∆tmax = 10−1 . Subfigure 21c shows the number of DEER iterations used for each block/ sequence length configuration (B, T ). For B > 1, the number of DEER iterations per sequential block may vary. Therefore, the iteration counts in Subfigure 21c are the median number of iterations to convergence across blocks in a (B, T ) simulation. This median is then used in all blocks in each (B, T ) setting to obtain the time and memory metrics. Subfigure 21d shows the efficiency gains in simulation time and peak memory usage of Conv DEER compared to Forward DEER. Note that the Conv DEER implementation here is suboptimal because it recomputes Λ̄ for At instead of caching it from the forward pass. In contrast, Forward DEER reuses the same intermediary results from the forward pass where possible, minimising overhead in accordance with Table 5. Peak memory utilisation is obtained with torch.cuda_max_memory_allocated(), while simulation time is obtained as an average of 7 trials after 2 warm-up runs. The recordings are taken without saving gradients for the backwards pass.
Scale-ELK Damping The hypothesis presented in Section 4.5.3 is that applying a scalar damping factor to Jacobian-free Quasi DEER approximations would "squeeze" the local approximation errors introduced by the diag(Jt ) truncation (Table 1). Figure 22 shows this phenomenon in effect. Subfigure 22a displays how increasing the damping effect first removes high-iteration-count outliers for both Conv and Forward DEER. Then, all algorithms progressively enter an over-damping regime where iterations to convergence rise because actual information, rather than local approximation errors, vanishes before it propagates to target future time steps (Eq. 101). Subfigure 22b brings additional detail to the gradually growing alignment between Conv/Forward DEER and Quasi DEER as the damping effect increases. Interestingly, on the particular settings tested here, only mild damping of 0.999 or 0.99 is sufficient to remove most outliers for both Conv and Forward DEER. Still, Conv-DEER retains a slightly higher mismatch in iterations to convergence than Forward DEER, relative to Quasi DEER, across all damping-factor configurations tested. This perhaps reflects the non-trivial influence of the ϵf wd error term (Table 1). 50
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling Iterations Required for Convergence Quasi DEER Conv DEER Forward DEER Damping: 0.999 Damping: 0.99
Damping: 1.0
250
Damping: 0.95
Damping: 0.8
Iterations
200 150 100 50
32
30
0
30
26
30
26
26
30
75
30
75
238
75
238
238
Algorithm Comparison (Quasi, Conv, Forward)
(a) Total Iterations for Convergence
Differences in Convergence Iterations between Approx. and Quasi DEER Baseline (Quasi)
Damping = 1.0
Frequency
10 8 6 4 2 0
8
0
8
16
24
32
40
48
56
18 16 14 12 10 8 6 4 2 0
Damping = 0.999
Conv - Quasi
Forward - Quasi
Damping = 0.99
Damping = 0.95 40
56
24
35
48
30
20
2
4
6
8
10
12
14
16
24
15
8
10
4
5
0
32
20
12
0
40
25
16
2
Damping = 0.8
28
2
1
0
1
0
Difference in Iterations (Algorithm - Quasi DEER)
16 8 1
0
1
0
1
0
1
(b) Reducing Approximation Errors
Figure 22: Effect of Damping on DEER Convergence. Subfigure 22a shows the effect of Scale-ELK damping on the absolute number of iterations required for convergence. The 30 dots in each column represent the convergence iterations for different random input samples, while the column shows their median. Each sample is also obtained with its own randomly initialised ADPTNet layers, with 16 eigenvalues Λ̄ set to negative, leading to a higher median count for T = 1024 than reported in Figure 20. All models tested have β = 0.125, ∆tmin = 10−4 , and ∆tmax = 10−1 . Subfigure 22b shows the relative difference in iterations to convergence between Conv/Forward DEER and Quasi DEER. The smooth curve shows the probability density. It can be noted that increasing the damping effect quickly reduces the number of outliers.
5.3.2
ADPTNet Parallelisation during Training
In Section 5.1.3, it is shown that the spread of Λ̄ values tends to increase over training epochs. At the same time, Figure 19c underlines how a higher eigenvalue spread leads to slower DEER convergence. Therefore, one could expect that DEER convergence would slow down throughout training. In Figure 23 it can be seen that this is indeed likely to happen in practice. Taking all subfigures into consideration, it can be observed how (in the absence of damping), the sharp increase in DEER iterations during the early training epochs coincides with the fastest rate of growth in the Λ̄ spread (i.e., during the first ≈ 20 − 30 epochs). In addition, the fully parallel simulations in Subfigure 23a appear prone to occasionally collapsing into sequential-step parity (DEER iterations = T ) for the Jacobian-free DEER versions. Scale-ELK damping seems to eliminate extreme outliers from Forward-DEER simulations, at the cost of considerably more baseline iterations to convergence. In the case of Conv-DEER, however, even severe damping, which leads to orders of magnitude higher baseline convergence iterations, does not fully eliminate T -iteration outliers. In contrast, block-wise DEER simulations appear stable and even equivalent for all DEER variants, without requiring any damping (Subfigure 23b). Given that block-wise execution may be the only viable option for training on memory-constrained hardware, this raises the prospect of Conv and Forward DEER as viable, faster and more efficient alternatives to Quasi-DEER. 51
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
DEER Convergence over Training Epochs Using Damping
No Damping
1024
Damping: 0.95
Damping: 0.8
512 256
Iterations
128
Algorithm Quasi DEER Forward DEER Conv DEER
64 32 16
20
40
Epoch
60
80
100
20
40
Epoch
60
80
100
20
40
Epoch
60
80
100
(a) Damping Fully Parallel DEER
Block-Wise DEER Convergence over Training Epochs
Linear Eigenvalue Spread | max
Newton Iterations
7 6 5 4 3 2
Evolution of Eigenvalue Spread over Epochs 1.0 min|
Algorithm Quasi DEER Forward DEER Conv DEER
8
20
40
60
Training Epoch
80
100
0.8
0.6
0.4
0.2
20
(b) Block-Wise DEER
40
60
Training Epoch
80
100
(c) Eigenvalue Spread Evolution
Figure 23: DEER Convergence over Training Epochs. The convergence data in this figure are obtained from the same fully positive Λ̄ model snapshots used to produce the spectral analysis in Figure 17 (Table 3). That also means that the training iterations correspond to learning the T = 1024 Copy Memory Task. Consequently, Subfigure 23c is a reproduction of the data in Subfigure 17b included here to aid convergence data analysis. For the block-wise simulations in Subfigure 23b, the number of sequential blocks is B = 8. The shown convergence iteration counts are obtained from the same randomly selected Copy Memory test sample. The main observations are that convergence iterations for all DEER variants increase over epochs in line with increasing Λ̄ spread. Additionally, in this particular case, Conv-DEER appears unstable for full DEER parallelism in Subfigure 23a, while block-wise execution in Subfigure 23b appears to mitigate the convergence pathology. 5.4
Measuring Non-Linearity with State Tracking
Recall from Section 2.2 that SSMs place non-linear activation functions position-wise, between layers. Therefore, a single SSM is, in theory, restricted to modelling fading memory models that eventually decay back to a singular resting state / fixed point. However, despite their linear recurrence, there is substantial evidence to suggest that fixed-depth stacked SSMs are capable of modelling non-linear dynamical systems [140, 98], such as Mackey-Glass [27, 2]. In other words, for practical purposes, SSM systems can model non-linear dynamics. Conversely, dynamical systems are not fully adequate to answer the question of the extent to which the ADPTNet topological conjugate recurrence induces non-linearity. Instead, in recent years, the de facto standard methodology for probing non-linear sequential modelling has become state tracking [128]. As the name suggests, state tracking refers to the capability of a single-layer sequence model to track changes to a state over time. For instance, Merrill et al. [129] provides the analogy of tracking a game of chess. Sequential piece moves affect the state of the board at each time step, and the network needs to correctly encode these state changes all throughout the game to correctly output its end-state. The same principle can be more abstractly encapsulated by the word problem over the symmetric group Sn 16 [129]. The symmetric group encodes permutations applied to a set of items. For example, given the set of symbols, [a, b, c, d, e], the S5 group encodes all possible permutations(e.g., swap(1, 3) → [c, b, a, d, e] ∈ S5 ). The world problem over S5 consists of tracking a sequence of permutations applied to the base set (e.g., what is the end result of sequentially applying swap(1, 3) → swap(2, 4) → . . . ). 16
S is once again being overloaded. In SNN contexts, it represents the spiking activation; in the ADPTNet recurrence, it represents a skew-symmetric matrix; and here it is the standard notation for the symmetric group
52
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
Crucially, in a seminal finding, Merrill and Sabharwal [128] theoretically proved that state tracking tasks such as S5 cannot be solved by fully parallel architectures such as Transformers without scaling network depth with sequence length. In contrast, traditional non-linear RNNs can perfectly simulate S5 using a single layer. Merrill et al. [129] extended the result to SSMs, which despite the recurrent formulation, lack the expressivity of non-linear RNNs and cannot solve state-tracking at fixed depth either. Table 7 summarises concrete challenges that Mamba, for instance, faces when applied to the S5 task, as reported in Schöne et al. [155]. The same contrast between linear and non-linear recurrences can be observed in Table 8, where a vanilla RNN generalise to 2x and 4x the training sequence length. At the same time, a single-layer,linear, and diagonal ADPTNet (β = 0) cannot surpass ≈ 31.2% accuracy on the training sequence length, and collapses to random accuracy on longer sequences. Furthermore, not even increase the width of the system can improve performance. Grazzi et al. [70] showed how including negative eigenvalues in linear SSMs can help the networks learn parity state tracking tasks (e.g., determine the final sign after mulitplication with sequence of {±1}: 1, −1, 1, 1, −1, . . . ). Interestingly, even for β = 0 on a non-parity state tracking task such as S5 , there appears to be a slight improvement of performance when setting half the eigenspectrum negative. Model
Seq.Len.
N. Layers Required
Extrapolation
Mamba
8 32
4 16
× ×
RNN
8 32
1 1
✓ ✓
Table 7: Mamba / RNN State Tracking Performance Summarised from Schöne et al. [155] This table summarises the number of layers required for Mamba and non-linear RNNs to converge on the S5 task, as reported in Schöne et al. [155]. Of note, Mamba requires more layers as sequence length grows, while not being able to extrapolate to longer sequences at test time (i.e., 2x or 4x the training length). In contrast, a single layer RNN can not only solve the task, but also extrapolate to longer input sequences.
Eig. Init.
Model
Width
Sparsity
Neg. Eig.
Accuracy (%) L=8 1×
L=16 2×
L=32 4×
–
ReLU RNN ReLU RNN
512 864
– –
– –
100.0 100.0
100.0 100.0
99.6 100.0
Linspace
ADPTNet (β = 0) ADPTNet (β = 0) ADPTNet (β = 0)
512 624 512
0% 0% 75%
✓ ✓ ✓
31.2 31.2 26.5
1.1 1.0 0.8
0.9 0.9 0.9
S4D-Real
ADPTNet (β = 0) ADPTNet (β = 0) ADPTNet (β = 0) ADPTNet (β = 0)
512 512 512 512
0% 0% 75% 75%
✓ × ✓ ×
28.6 26.7 26.9 25.3
1.2 0.9 1.1 0.9
0.8 0.8 0.8 0.8
Table 8: Baseline S5 Accuracy Results are obtained by averaging over three seeds for each configuration and following the training setup from Appendix C. One can observe the strong alignment with the results in Table 7 with respect to the stark gap in accuracy between linear and non-linear recurrence.
Table 9 summarises the impact on state tracking abilities that the choice of ADPTNet initialisation and parametrisation can have. Overall, it can be observed that the best predictor of S5 accuracy is β, with performance monotonically increasing with a higher β, regardless or eigenspectrum initialisation/parametrisation. Moreover, for the largest setting β = 4, almost all eigenspectrum choices enable signs of extrapolation up to 4× the training sequence length. This evidence supports the intuition that as β increases, the "sharpness of turns" of the topological conjugates with β-bound Lispchitz-ness (Corollary 1), also increases to the point where the system is sufficiently non-linear to model state tracking. Stated differently, the choice of β is effectively a gradual parametric "dial" for increasing the non-linearity of the recurrent dynamics. 53
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
While β dominates S5 performance, it does not singularly control it. A second visible factor is the choice of eigenspectrum parametrisation, which comprises whether 75% of the eigenvalues are masked to 0 permanently, and whether half of the non-zero eigenvalues are constrained to be negative. For β < 1, negative eigenvalues appear to improve accuracy regardless of sparsity, while for β = 1, sparsity appears to becomes the dominant driver of performance. Interestingly, for β = 4, sparsity and negative eigenvalues become mutually exclusive, with both reaching comparable peak accuracies only when separated from the other. The effect does appear conditional on the initialisation used, with Linspace reaching its highest performance with a sparse and positive spectrum, while S4D-Real does so on the reverse. In terms of maximum extrapolation performance, Linspace initialisation outperforms S4D-Real. Hence, it can be argued that sparsity and negative eigenvalues engender non-linearity through different inductive biases. On one hand, sparsity can be seen as forcing the network to learn ReLU-like forgetting, with the similarity transform effectively rotating information into the 0-masked dimensions to be discarded. The "surviving" information is then projected onto a low-dimensional state space akin to the action of low-rank RNNs [127]. On the other hand, negative eigenvalues maximise the possible spread of the eigenspectrum. This, in turn, maximises the output variance from the similarity transforms, increasing model expressivity. At the lowest tested β = 0.00125, both mechanisms are necessary to boost accuracy slightly over the linear β = 0 baseline (32.4% > 31.2%). However, with higher β, as previously highlighted, they become incompatible.
Eig. Init.
Sparsity
β
Accuracy (%)
Neg. Eig. L=8 1×
L=16 2×
L=32 4×
0%
4 4
✓ ×
43.4 51.5
1.7 1.3
0.8 0.9
75%
0.00125 0.0125 0.125 1 4 4
✓ ✓ ✓ ✓ ✓ ×
28.8 36.2 43.6 76.9 79.9 98.1
0.8 1.2 0.9 30.5 44.2 62.4
0.8 1.1 0.8 1.5 8.7 2.9
0%
0.00125 0.00125 0.0125 0.0125 0.125 0.125 0.5 0.5 1 1 4 4
✓ × ✓ × ✓ × ✓ × ✓ × ✓ ×
30.3 29.4 50.3 42.5 50.7 45.1 49.6 52.0 64.3 56.2 85.3 66.1
1.3 1.2 1.4 1.3 1.7 1.3 1.6 1.5 18.1 1.5 60.3 27.4
0.8 0.8 0.8 0.8 0.8 0.8 0.8 0.9 5.5 0.9 23.0 2.3
75%
0.00125 0.00125 0.0125 0.0125 0.125 0.125 0.5 0.5 1 1 4 4
✓ × ✓ × ✓ × ✓ × ✓ × ✓ ×
32.4 29.4 48.4 43.8 57.3 44.4 53.6 51.2 55.9 53.7 64.5 83.3
1.4 1.4 1.3 1.2 1.6 1.5 2.5 1.4 1.7 1.3 29.5 45.5
0.9 0.8 0.8 0.8 0.8 0.8 0.8 0.9 0.9 0.9 10.2 4.2
Linspace
S4D-Real
Table 9: Effect of β and Eigenspectrum on ADPTNet State Tracking Performance Each result in this table is obtained by averaging over three seeds, as in Table 8. In addition, the training configuration is in Appendix C. It can be seen that the most reliable indicator of S5 word problem accuracy is the choice of β. To maximise performance, however, eigenvalue initialisation and parametrisation must also be considered.
54
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
Table 10 shows the effect of including additional k and v processing steps on S5 -word problem accuracy, this time on training sequence length L = 32. As mentioned in Section 4.1, applying a ReLU non-linearity before K and V projections does not affect long-term behaviour guarantees related to the Lyapunov spectrum (Lemma 2). However, as shown in Table 10, it does make a substantial difference in state-tracking ability. More specifically, including ReLU doubles accuracy on the training length (perfectly learning the task with 100% accuracy), and almost completely bridges the gap to vanilla RNNs for 2× length extrapolation (97%). Even in the 4× extrapolation setting, maximum accuracy is visibly higher with ReLU than all results on the shorter sequences in Table 9. Here, it is important to recall that the orthogonal matrices used in the ADPTNet recurrence are obtained via Riemannian gradient descent steps from the origin I, of step size β and in a high-dimensional direction defined by k and v outer and inner products. Therefore, there are two hyperparameters for ADPTNet non-linearity and expressivity that do not pertain to β or the inclusion of ReLU activations. Firstly, as argued in Section 4.12, the strength of off-diagonal interactions is controlled, in part, by the dot-product similarity between k and v. The < k, v > setting in Table 10 uses the a priori dot product control method from Section 4.12 to either enforce or not enforce k · v = 0.5. As a result, it can be seen that strict constraints on the variance of the dot product do reduce model expressivity and hurt S5 performance to some extent. For example, compared to 49.2% accuracy on L = 32 for the baseline configuration (no LayerNorm, ReLU or k · v control), including k · v reduces accuracy by almost 6% (43.4%). Secondly, LayerNorm can be applied to Kx and V x before unit-normalisation to constrain their variance to 1. In other words, the variance in the directions to travel across the manifold is constrained. Table 10 shows how including this effectively erases state-tracking abilities. Conversely, this suggests that ADPTNet actively learns to control the interaction between k and v to achieve state tracking. However, including ReLU largely mitigates constraints on both the inner-product and outer-product. Mean Accuracy (%) LN
ReLU
⟨k, v⟩
× ✓ × × ✓ ×
× × ✓ × ✓ ✓
× × × ✓ × ✓
Best (%)
L=32 1×
L=64 2×
L=128 4×
L=128 4×
49.2 5.3 100.0 43.4 100.0 100.0
26.3 3.6 97.0 21.8 93.6 89.8
9.9 2.5 5.1 8.7 0.8 15.8
10.3 2.9 36.8 8.9 8.8 29.7
Table 10: Ablation on L=32 S5 As with Tables 8 and 9, results are obtained by averaging over 3 seeds and following the training configuration from Appendix C. LN refers to whether LayerNorm is applied to Kx and V x before unit normalisation. ReLU refers to including a ReLU activation in KReLU(x) and V ReLU(x). < k, v > denotes whether the dot product between k and v is constrained to 0.5 using the methods from Section 4.12. Overall, ReLU can be observed to significantly benefit state tracking performance, mitigating even the effects of adding LN and < k, v > constraints, which otherwise can severely impact accuracy.
5.5
Measuring Adaptability with Selective Copy
Sections 5.1.3 and 5.3.2 employ the Copy Memory task (Fig. 14), which measures a network’s ability to learn fixed delay lines. While it effectively probes long-range dependencies, it requires no selectivity or adaptability. Instead, Gu and Dao [71] showed that Selective Copying task performance can distinguish data-dependent, adaptive networks, such as Mamba, from fixed-parameter LTI SSMs such as S4. In Selective Copying, target tokens are randomly spread throughout the input sequences, interleaved with distractors (Fig. 24). As in the fixed Copy Memory task, the network is prompted to recall all the target tokens at the end of the input sequence. Gu and Dao [71] presented evidence that, while S4 struggles to generalise on Selective Copy, Mamba can fully solve it given the same parameter budget. The most widely accepted explanation is that a non-data-dependent network effectively dilutes its memory with distractors and cannot discard unhelpful information. Conversely, adaptive/selective models dynamically discard distractors. Tables 11a and 11b show how ADPTNet performs on the Selective Copying task. Here, the focus is on ablating the effect of β on selectivity rather than establishing absolute performance metrics compared to existing methods. In this regard, both tables mirror the state-tracking results from Table 9. Namely, there is a robust and monotonic improvement in Selective Copying accuracy stemming from higher β. The effect is present regardless of whether the topological conjugacy operation is linear (i.e., input-driven, see Section 4.10) or non-linear (i.e. recurrent state-driven). Remarkably, Table 11a shows how even relatively small β = 0.00125, which effectively forces off-diagonal interactions 55
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
to be weak couplings/perturbations, shows visible improvements in selectivity over the diagonal β = 0 baseline. The non-linear ADPTNet recurrence in Table 11b similarly shows dramatic improvements in selectivity with a relatively low β = 0.0125, more than doubling Selective Copy accuracy. It is important to note that the results between the two tables discussed here are not directly comparable, since Table 11a is the result of a larger model with L = 1024 sequence length, compared to the L = 128 sequence length in Table 11b.
Figure 24: Selective Copy Task The Selective Copying task inherits the same general formulation as the Copy Memory task (Fig. 14). The difference is that here the target tokens are spread out randomly among the distractors.
Model
Acc (%)
Baseline (β = 0)
46.3
Model
Acc (%)
β = 0.00125
64
Baseline (β = 0)
36
β = 0.0125
69
β = 0.0125
80
β = 0.125
75.6
(b) Non-Linear ADPTNet
(a) LinearADPTNet
Table 11: Effect of β on Selective Copy Accuracy Details of the experimental setup are available in Appendix D. It should be mentioned that each baseline in Tables 11a and 11b is obtained by taking the maximum accuracy over three seeds across six learning rate × weight decay configurations. The main observation from both tables is that using a larger β improves Selective Copying accuracy. Table 11a is obtained using sequence length L = 1024, while Table 11b uses L = 128. In addition, Table 11a uses a model hidden size of h = 128, and an state expansion of nstate=2 (see Section 4.6). For Table 11b, hidden size is h = 64 and state expansion is nstate = 4. It is important to highlight that the specific hyperparameter selection here is not optimised for producing the absolute highest Selective Copy accuracy. Instead, it purposefully underparametrises models to the point where differences in expressivity can be isolated. Evidently, drastically scaling the number of parameters could allow non-selective architectures to simply memorise all possible target token positions.
5.5.1
ADPTNet Task-Dependent Adaptation
Given how closely related the Copy Memory and the Selective Copying tasks are, one can use them to probe how ADPTNet learns to use the topological conjugate for different purposes. Namely, by setting the same vocabulary size, number of target tokens and distractors, model size, and training hyperparameters, it is possible to study how the exact same randomly initialised network evolves depending on the learning objective: selectivity or fixed delay-line learning. The specific shared hyperparameters between the twin experimental conditions are included in Appendix D. Lyapunov Spectrum A key question worth answering is whether training for selectivity affects the Lyapunov spectrum’s alignment with the parametrised recurrent eigenspectrum. In this regard, Figure 25 shows how, irrespective of the task objective or network layer, the Lyapunov spectrum rests within the bounds of ln(|Λ̄|), and thus implicitly meets the prediction of Lemma 2. Interestingly, as one would intuitively suspect, the eigenspectra resulting from the Selective Copying tasks have lower minima than those from Copy Memory. Lower eigenvalues yield shorter timescales, allowing 56
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
quicker decay and disposal of distractors. It should also be noted that while Figure 25 appears to show a stronger "squeezing effect" of the Lyapunov spectrum compared to the eigenvalues, it is only so relative to the spread of the spectrum. Specifically, on all layers for both tasks, the absolute difference between the minimum Lyapunov exponent αmin and ln(|Λ̄|)min is on the order of ≈ 1−2 .
Fixed , Copy Memory Layer 0 Lyapunov Exponents against the Transition Eigenvalues
Fixed , Copy Memory Layer 1 Lyapunov Exponents against the Transition Eigenvalues
0.00
0.00 0.01
Log Growth Rate per Step
Log Growth Rate per Step
0.01
0.02
0.03 ln| | Lyapunov Exponents Eigenspectrum Bounds Lyapunov Bounds Within Eigenvalue Bounds Outside Eigenvalue Bounds
0.04 0
100
200
300
Mode, Sorted Ascending
0.02
ln| | Lyapunov Exponents Eigenspectrum Bounds Lyapunov Bounds Within Eigenvalue Bounds Outside Eigenvalue Bounds
0.03 0.04 0.05
400
500
0
(a) Copy Memory Task - First Layer
200
300
Mode, Sorted Ascending
400
500
(b) Copy Memory Task - Second Layer
Fixed , Selective Copying Layer 0 Lyapunov Exponents against the Transition Eigenvalues
Fixed , Selective Copying Layer 1 Lyapunov Exponents against the Transition Eigenvalues
0.00
0.00
0.02
0.05
Log Growth Rate per Step
Log Growth Rate per Step
100
0.04 ln| | Lyapunov Exponents Eigenspectrum Bounds Lyapunov Bounds Within Eigenvalue Bounds Outside Eigenvalue Bounds
0.06 0.08 0.10
ln| | Lyapunov Exponents Eigenspectrum Bounds Lyapunov Bounds Within Eigenvalue Bounds Outside Eigenvalue Bounds
0.10
0.15
0.20
0.12 0.25 0
100
200
300
Mode, Sorted Ascending
400
500
0
(c) Selective Copy - First Layer
100
200
300
Mode, Sorted Ascending
400
500
(d) Selective Copy - Second Layer
Figure 25: Lyapunov Spectrum learning for Matching Copying Tasks. Each subfigure here shows the Lyapunov spectrum compared to the eigenspectrum for each layer in both networks trained on Selective Copying and Copy Memory tasks with matching training configurations and model hyperparameters (Appendix D). Note that the empirical Lyapunov exponents are bounded and closely match ln(|Λ̄|) across all layers examined. To obtain the spectra, a random sample is chosen from both datasets for each corresponding model.
Trainable β Another question is whether, if ADPTNet layers are permitted to include β in the training, the learning objective might push it towards making the network more/less diagonal-dominated, non-linear, and selective. Using the same training and model setup from Appendix D, networks are now initialised with β = 0.0125 and allowed to train, with clipping applied to keep it within [0.0001, 0.025]. Table 12 shows that, regardless of task, learning pushes β towards its maximum, saturating the upper bound clipping value. This suggests that the similarity transforms are being actively used to learn the task, whether it is Selective Copying or Copy Memory. In other words, the models are not simply learning to circumvent the dynamic off-diagonal couplings. Still, Table 12 does also show how including trainable β only improves accuracy over the fixed baseline Selective Copying, following the trend of higher accuracies for higher β from Table 14. 57
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
Task
Layer
Learned β
Acc. Diff.
Copy Memory
0 1
0.026 0.029
-1.5%
Selective Copy
0 1
0.027 0.027
+2.4%
Table 12: Trained β across Tasks and Layers Both networks are initialised with β = 0.0125, and it can be observed that training pushes it to saturation of the upper clipping bound 0.025. The Acc. Diff. column shows the change in final accuracy compared to the fixed β models used for Figure 25. In absolute terms, fixed β resulted in 89.69% accuracy for Copy Memory and 63.70% for Selective Copy. For trainable β, the accuracies are 88.20% for Copy Memory and 66.07% for Selective Copying.
Recurrent Weights Entropy In this context, normalised Shannon entropy [157] H measures the degree of uniformity across the recurrent matrix entries, with H = 1 being achieved for all constant entries, and H = 0 for a single large entry that dominates. It should be mentioned that since the diagonal of M (Eq. 49) is dominant regardless of task, it is excluded from this computation. Hence, for the off-diagonal entries |mi | from M ∈ Cn×n , the discrete probabilities pi , i = 1, 2, . . . , t, t = n(n − 1) and the normalised entropy H are shown in Equation 118.
|mi |2 pi = Pt 2 j |mj | Pt pi log(pi ) H=− i log(t)
(118)
Figures 26 and 27 show how the Copy Memory and Selective Copy tasks shape the data-dependent ADPTNet recurrent matrices across depth and training epochs. In this regard, the most striking effect observable is the higher degree of uniformity emerging from training on Selective Copy. On both layers, Selective Copying engenders higher off-diagonal entropies compared to Copy Memory at the end of training, (0.77 > 0.66 for the first layer and 0.71 > 0.61 for the second layer). Visually, Figure 26 in particular shows the largest discrepancy emerging from training on the two tasks, with the Selective Copying-induced first layer dynamics evidently showing a denser "all-to-all" coupling pattern between recurrent neurons, compared to Copy Memory. Since the weights are data-dependent, slight differences are also visible at initialisation. Namely, while the overall patterns largely match layer-wise between the two tasks, Selective Copy displays persistently higher amplitudes in both Figures 26 and 27. These observations, combined with the more quickly decaying time-scales present for Selective Copy (Figure 25), point towards the networks effectively learning to not only create "forget" channels, but also learning to move distractor information to those channels for disposal. In other words, some neurons learn to have quickly fading memory, and the network dynamically assigns information to those neurons to induce selective forgetting via changing connectivity. As a concrete example, on Copy Memory, the minimum log-eigenvalue is ln(λ̄min ) ≈ −0.04 and the sample recurrent matrix entropy is ≈ 0.61 for the second layer of the network. By contrast, for Selective Copying, the second layer has ln(λ̄min ) ≈ −0.25 and a sample recurrent matrix entropy of ≈ 0.71. Therefore, Selective Copying induces higher entropy (more uniform connectivity) and a lower minimum timescale (faster decay), consistent with the hypothesis of the network routing information to "forget" neurons. 58
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
Fixed , Copy Memory At Initialisation Off-Diagonal Entropy 0.7323
0
10 1 10 2 10 3 10 4 10 5 10 6 0 10 6
150
10 5 10 4 10 3 10 2 10 1
200
50
100 150 State Index (Column)
200
10 1 10 2 10 3 10 4 10 5
50
State Index (Row)
State Index (Row)
100
0
After 100 Epochs of Training Off-Diagonal Entropy 0.6613
0
50
250
Layer 0: Transition Matrix
100
10 6 0 10 6
150
10 5 10 4 10 3 10 2 10 1
200
250
250
0
50
100 150 State Index (Column)
200
250
(a) Copy Memory - First Layer
Fixed , Selective Copying At Initialisation Off-Diagonal Entropy 0.7005
0
10 1 10 2 10 3 10 4 10 5 10 6 10 7
0 10 7 150
10 6 10 5 10 4 10 3 10 2 10 1
200
0
50
100 150 State Index (Column)
200
10 1 10 2 10 3 10 4 10 5 10 6
50
State Index (Row)
100
250
After 100 Epochs of Training Off-Diagonal Entropy 0.7734
0
50
State Index (Row)
Layer 0: Transition Matrix
100
010107 7 150
10 6 10 5 10 4 10 3 10 2 10 1
200
250
250
0
50
100 150 State Index (Column)
200
250
(b) Selective Copy - First Layer
Figure 26: Recurrent Matrix M Materialisation during Training - First Layers. Figures 26 and 27 show how the data-dependent recurrent matrices M adapt depending on the task, Copy Memory or Selective Copy. Here, the same models and training experiments are used as in Fig. 25. On the left, the matrices are extracted at initialisation (before training) from their respective networks, while on the right the matrices are materialised at the end of training. Each entry in the matrices is the coupling strength between two recurrent neurons, so each matrix also implicitly represents a network topology. The main differences pertain to how evenly-distributed the off-diagonal entries are, as quantified by the normalised Shannon entropy (Eq. 118). It can be observed that Selective Copying leads to higher overall entropy, and thus more evenly distributed all-to-all recurrent connectivity.
59
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
Fixed , Copy Memory At Initialisation Off-Diagonal Entropy 0.7279
0
10 1 10 2 10 3 State Index (Row)
10 6 0 6 10 150
10 5 10 4 10 3 10 2
200
50
100 150 State Index (Column)
200
10 4 10 5
100
10 6 0 6 10 150
10 5 10 4 10 3 10 2
200
10 1 0
10 1 10 2 10 3
50
10 4 10 5
100
250
After 100 Epochs of Training Off-Diagonal Entropy 0.6137
0
50
State Index (Row)
Layer 1: Transition Matrix
250
250
10 1 0
50
100 150 State Index (Column)
200
250
(a) Copy Memory - Second Layer
Fixed , Selective Copying At Initialisation Off-Diagonal Entropy 0.7725
0
10 1 10 2 10 3 10 4 10 5 10 6 0 10 6
150
10 5 10 4 10 3 10 2 10 1
200
50
100 150 State Index (Column)
200
10 1 10 2 10 3 10 4 10 5 10 6 0 10 6
50
State Index (Row)
State Index (Row)
100
0
After 100 Epochs of Training Off-Diagonal Entropy 0.7108
0
50
250
Layer 1: Transition Matrix
100
150
10 5 10 4 10 3 10 2 10 1
200
250
250
0
50
100 150 State Index (Column)
200
250
(b) Selective Copy - Second Layer
Figure 27: Recurrent Matrix M Materialisation during Training - Second Layers. This figure accompanies Figure 26, showing the effect of training on different tasks on the materialised recurrent matrix M . While the effect on the second layers is less dramatic than on the first, there is still a visiable difference in entropy, with more evenly spread out connectivity for Selective Copying than Copy Memory.
DEER Iterations Figure 28 shows how Conv-DEER iterations required for convergence change after training and between tasks. Firstly, the same trend from Figure 23 is also visible for both Copy Memory and Selective Copying. Namely, convergence slows down as training progresses. Furthermore, as seen in Figure 25, eigenvalue spread is higher in the deeper layers, and, accordingly, DEER convergence is slower on the second layers in all experimental conditions. Corroborating Figure 19, on all layers for both tasks, it can be seen that pushing β to the 0.025 ceiling after training invariably increases Newton iterations. Interestingly, however, undermining the trends set out in Section 5.2, Figure 28 shows how the convergence for Copy Memory is slightly slower than that of Selective Copying. This is despite the fact that, as highlighted in Figure 25, Selective Copy induces a higher eigenvalue spread. One could hypothesise that higher off-diagonal entropy partially mitigates the effect of higher spread by effectively smoothing information across recurrent neurons. While β remains a 60
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
10 2 10 3 10 4 10 5 2
4
6
8 10 Solver Iteration
12
Fixed , Copy Memory
Layer 0
14
16
Before Training After Training DEER Tol. (0.0005)
10 1 10 2 10 3 10 4 10 5 2
4
6
8 10 Solver Iteration
Trainable , Selective Copying
12
14
16
Layer 0 Before Training After Training DEER Tol. (0.0005)
10 1 10 2 10 3 10 4 10 5 2
4
6
8 10 Solver Iteration
Trainable , Copy Memory
100
12
14
16
Layer 0 Before Training After Training DEER Tol. (0.0005)
10 1 10 2 10 3 10 4 10 5 2
4
6
8 10 Solver Iteration
12
14
16
Largest Deviation from the Exact Solution
Before Training After Training DEER Tol. (0.0005)
Largest Deviation from the Exact Solution
Layer 0
Largest Deviation from the Exact Solution
Fixed , Selective Copying 10 1
Largest Deviation from the Exact Solution
Largest Deviation from the Exact Solution
Largest Deviation from the Exact Solution
Largest Deviation from the Exact Solution
Largest Deviation from the Exact Solution
reliable predictor of DEER convergence, perhaps eigenvalue spread also requires consideration of effective recurrent connectivity to predict DEER convergence. Fixed , Selective Copying
Layer 1 Before Training After Training DEER Tol. (0.0005)
10 1 10 2 10 3 10 4 10 5
2
4
100
8 10 Solver Iteration
12
Fixed , Copy Memory
6
Layer 1
14
16
Before Training After Training DEER Tol. (0.0005)
10 1 10 2 10 3 10 4 10 5
2
4
6
8 10 Solver Iteration
Trainable , Selective Copying
12
14
16
Layer 1 Before Training After Training DEER Tol. (0.0005)
10 1 10 2 10 3 10 4 10 5 2
4
6
8 10 Solver Iteration
Trainable , Copy Memory
100
12
14
16
Layer 1 Before Training After Training DEER Tol. (0.0005)
10 1 10 2 10 3 10 4 10 5 2
4
6
8 10 Solver Iteration
12
14
16
Figure 28: Training Effect on DEER Convergence. This figure uses the models from Table 12 and Figure 25, i.e., the trained and fixed β models trained on either Selective Copying or Copy Memory. Using Conv-DEER, it can be seen that deeper layers are slower to converge across all settings. Notably, Selective Copy appear to converge more quickly than Copy Memory, despite the higher spread eigenspectrum (see Fig. 25).
5.6
Measuring Loose Coupling Expressivity on Mechanistic Architecture Design
The key question being answered in this section is whether a small β = 0.00125, which induces loose perturbation-like coupling, is expressive enough to compete with existing Transformer efficient alternatives. For this purpose, this study uses the Mechanistic Architecture Design (MAD) benchmark [146]. The MAD benchmark is a synthetic test suite designed to predict large-scale language modelling capabilities from small token manipulation tasks. This enables identification of how specific model characteristics compare with established sequence architectures, including self-attention and efficient alternatives such as SSMs. Furthermore, all tasks within MAD use a fixed model dimension, standardised training hyperparameter sweeps, and the same task difficulty settings, maximising comparability across architectures. Thus, in this work, MAD serves as a proxy for validating ADPTNet adaptability, as required for language modelling, against prior methods. 61
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
One of the tasks in MAD is the now familiar Selective Copying. The overall MAD score for Selective Copying is computed by varying its difficulty through increased numbers of distractors, larger vocabulary sizes, more target tokens to store and smaller training dataset sizes (see Appendix E for details). Another task is Compression, which tests a network’s ability to compress the concatenated input token sequence into a fixed-sized context (e.g., input : [a, b, c, d, e] → output : cat(abcde)). The basic premise is reminiscent of the Long-Range Arena (LRA) benchmark [169], which tests a network’s ability to model long-range dependencies and implicitly requires a degree of context compression. However, while some of the LRA’s tasks comprised tokens simply consisting of greyscale pixels, Compression synthetically and systematically constructs larger and more diverse vocabularies, requiring a higher degree of expressivity from the network. As with Selective Copying, the details of the difficulty tunings for Compression are in Appendix E. The final MAD task employed here is Memorisation, which tests key-value associative retrieval of factual memories. This differs from in-context recall tasks because it only requires learning fixed key-value pairs of facts from the training data as ground truth and answering prompts at test time from the same fact collection (e.g., training set (key, value) pairs: [(a, b), (c, d), (e, f), . . . ], test prompt: [a →?, c →?]). Memorisation can be considered a precursor to LLMs storing facts such as the (Key: Capital of Romania, Value: Bucharest). Once again, MAD averages over several Memorisation difficulty settings, including the number of facts to store and training dataset size (Appendix E). It should be noted that, while they do test computational primitives required for language modelling, neither Memorisation nor Compression requires in-context adaptation to the same degree as Selective Copying. Furthermore, key-value associative fact knowledge is typically stored within position-wise Feed-Forward Networks (FFN) / MLPs in LLMs [66]. As such, the Memorisation task is likely to be a stronger reflection of those components, rather than the sequence-mixing layers such as ADPTNet. Table 13 shows how ADPTNet within the perturbation regime (β = 0.00125) compares against existing MAD baselines. Notably, the only architectures with vector states in the table are ADPTNet and Hawk [40]. The rest not only use matrixvalued states, but also expand the recurrent dimension to ∈ R2048×2048 , much larger than C512 for ADPTNet. Despite this, for the Compression task, ADPTNet outperforms all reported matrix-valued architectures and is outperformed only by the Transformer and Hawk. In contrast, Memorisation appears more challenging for the tested ADPTNet configuration, as it only outperforms Mamba2 [34] and DeltaNet. Before examining Selective Copying performance, it is important to establish in more detail how Hawk relates to ADPTNet. As proposed in De et al. [40], Hawk relies on a selective Real-Gated LRU (RG-LRU) sequence-mixing layer:
ht = at ⊙ ht−1 +
q
1 − a2t ⊙ (it ⊙ xt )
(119)
Where at = acrt , and the input-driven gates are rt = σ(Wr xt + br ), and it = σ(Wi xt + bi ), where σ denotes a softplus activation, a is parametrically constrained 0 ≤ a ≤ 1, and c is constant. As argued in De et al. [40], this recurrence rule differs from other selective SSMs such as Mamba because it cannot arbitrarily discard information. More precisely, in Mamba, each new state is obtained by interpolating between new inputs and the previous state, so there is no lower bound on the information that can be discarded at any given step. In other words, the network can decide to completely override its memory with the new input. By contrast, RG-LRU interpolates between its past state and an LRU update. Information is either added or not, but nothing is forgotten faster than the decay rates set by the network’s fixed eigenspectrum through a. In principle, this constitutes a direct alternative linear solution to ADPTNet’s stated goal of controlling the effective long-term dynamics of the recurrence, while enabling adaptability. Therefore, as highlighted in Table 13, for Selective Copy, the task that tests adaptability most directly among the MAD tasks included here, it is important to note that ADPTNet, at a perturbation level β = 0.00125, manages to outperform Hawk (80.9 vs 77.0). This provides evidence that the topological conjugate adaptability implemented in ADPTNet outperforms its closest equivalent in existing literature. Furthermore, it does so using weak perturbative coupling. 62
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
Model
Memorisation
Selective Copy
Compression
Mamba2 [34] GLA [192] xLSTM [12] DeltaNet [193] Gated DeltaNet[194] MesaNet [181] Transformer [181]
42.0 82.5 79.8 40.8 81.7 77.2 84.7
95.4 96.1 95.4 98.8 95.7 99.2 96.0
41.3 42.3 43.4 43.3 45.0 45.4 49.5
Hawk [40] ADPTNet (β = 0.00125)
91.3 46.1
77.0 80.9
47.7 45.9
Table 13: Accuracies on MAD Tasks The ADPTNet network tested here is situated in a perturbation regime, where off-diagonal couplings are scaled by β = 0.00125. The baseline results included in this table are reported from von Oswald et al. [181]. GLA stands for Gated Linear Attention [192]. All networks besides Hawk, the Transformer, and ADPTNet, compress information in a matrix-valued state. The Transformer evidently keeps the entire sequence as its state, while ADPTNet and Hawk both have vector-valued states. The most notable results are that ADPTNet performs competitively against models with much larger states on Compression, while outperforming its closest equivalent alternative, Hawk, on Selective Copying.
5.7
Measuring Long-Range Temporal Modelling with Sequential-CIFAR10
Sections 5.4, 5.5, and 5.6 showed how the data-dependent similarity transforms within ADPTNet recurrence can enable state tracking and selectivity in SSMs. The key question, then, is whether ADPTNet preserves the desirable long-range modelling qualities of LTI SSMs. For instance, Mamba achieves a markedly lower accuracy than S4 on the LRA [197]. Could the topological conjugation operation ruin the coherence of long-range features? To check this, this study uses Sequential CIFAR-10 (sCIFAR), a well-established benchmark for modelling long-range temporal dependencies [113, 79, 74, 169]. It consists of 32 × 32 images from the CIFAR-10 dataset being converted into L = 1024-long 1-dimensional temporal signals by processing each pixel at a time. The sequences are produced by concatenating each image’s rows together. Tables 14a and14b show how ADPTNet, with a β = 0.0125 which improves performance on both Selective Copying and state tracking (see Sections 5.4 and 5.5), compares to the baseline β = 0 as well as existing non-linear recurrent or selective/adaptive architectures. Firstly, Table 14b shows how state-driven non-linear ADPTNet does not degrade performance compared to the same baseline model with β = 0, and actually marginally improves accuracy. It should also be noted that, as mentioned in Section 4.2, ADPTNet is only theoretically guaranteed to preserve, within given bounds, the decay rates encoded in the eigenvalue spectrum. Because the baseline β = 0 and the non-linear β = 0.0125 models both use S4D-Inv complex initialisation and parametrisation, the slight improvement in accuracy suggests that ADPTNet preserves the imaginary/oscillatory components of the recurrent eigenspectrum to some extent in its long-term dynamics. Compared to existing non-linear RNNs, Table 14b also shows how small 2-layer non-linear ADPTNet is capable of outperforming them. Secondly, Table 14a shows how a larger 4-layer input-driven LinearADPTNet outperforms current selective/adaptive architectures. To the best of the authors’ knowledge, this is the highest recorded accuracy on this particular task (RGB colour sCIFAR) for a selective model. Moreover, the baselines in Table 14a use deeper networks (6 layers) than both ADPTNet networks tested, yet even the smaller 2-layer non-linear ADPTNet configuration outperforms them. Nevertheless, the main takeaway here is that ADPTNet retains baseline LTI SSM performance for long-range sequence modelling while also improving state tracking and Selective Copy. The topological conjugation data-dependent adjustments improve adaptability while not damaging long-range features. 63
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
Model
Acc (%)
Transformer [174]
62.2
Model
Acc (%)
CKConv [150]
63.74
Llama [173]
0.629
TrellisNet [9]
73.42
Mamba [71]
0.765
r-LSTM [174]
72.2
RWKV-4 [144]
0.757
UR-LSTM [73]
71.00
LSTM
0.756
UR-GRU [73]
74.4
xLSTM
0.761
HiPPO-RNN [72]
61.1
ADPTNet (β = 0.0125
81.34
LipschitzRNN [51]
64.2
Ablation S4D-Inv (β = 0)
76.39
ADPTNet (β = 0.0125)
76.62
(a) LinearADPTNet - Four Layers
(b) Non-Linear ADPTNet - Two Layers
Table 14: sCifar Accuracy Details of the experimental setup are available in Appendix F. The baseline accuracies in Table 14a are reported from Beck et al. [12], while Table 14b collects baseline results from Gu et al. [74]. It should be observed that both the 2-layer non-linear ADPTNet and the 4-layer LinearADPTNet outperform all of the larger selective architectures in Table 14a. Most importantly, however, the non-linear ADPTNet configuration matches and even slightly outperforms the baseline S4D-Inv model (β = 0) of the same size.
5.8
Neuromorphic Speech Processing with Spiking Speech Commands
As mentioned in Section 1, the neuroscientific inspiration for ADPTNet stems from the auditory cortex’s input-invariant time scales. As argued before, this implies that non-linear neuronal activity, synaptic modulation, and plasticity do not affect the temporal scales of population-level dynamics as captured by localised electrodes. Moreover, the same section detailed that the ultimate goal of ADPTNet is to provide a viable energy-efficient alternative to the Transformer, which could be achieved through an ADPTNet-based SNN (Section 4.11). Therefore, this section addresses whether SpikingADPTNet can play a role in neuromorphic auditory processing. The de facto standard tasks in this regard are the Spiking Heidelberg Digits (SHD) and Spiking Speech Commands (SSC) [31]. Both datasets are constructed by passing spoken language recordings through a biologically plausible artificial cochlea model. For SSC, the original recordings consist of the Google Speech Commands dataset, which is also prevalent in mainstream sequence processing [184, 74]. It should be noted that the SHD dataset is smaller, and state-of-the-art models have recently approached saturation near 100% accuracy [167]. Therefore, this work focuses only on the SSC task. The SpikingADPTNet setup used in this section is effectively a drop-in replacement for the LIF neurons within SiLIF [54]. The training, hyperparameters, network topology, and eigenvalue initialisation are held constant with the experiments in Fabre et al. [54]. This is achieved by adding spiking activations to CoupledADPTNet (see Section 4.7, which mirrors the AdLIF backbone of SiLIF (see Section 3.3). Both architectures are non-linear spiking units coupled with a slower-evolving "adaptation" linear state. To keep the parameter count comparable, LowParam linear layers are deployed within the input, K, and V projections. Full details of the training setup are included in Appendix G. Two SpikingADPTNet configurations are tested: one setting has its parameter count and recurrent state dimension matched to SiLIF, while the other expands the recurrent state ×4 to match the size of the state-of-the-art spiking Transformer baseline, the SpikCommander [182]. Table 15 shows the performance of SpikingADPTNet compared to the current state-of-the-art on SSC17 . The first important takeaway is that the parameter-matched SpikingADPTNet model outperforms the baseline SiLIF result (82.5 > 82.03 ± 0.25) and bridges the gap to the current state-of-the-art, DelRec [147]. Although only one seed is tested for the small SpikingADPTNet configuration, it lands slightly higher than DelRec’s mean accuracy (82.42). Scaling up, the larger SpikingADPTNet configuration sets a new state-of-the-art accuracy on the SSC dataset, averaging 83.56 over 4 random seeds, more than 1 point higher than the previous state-of-the-art. Improvements in accuracy with higher parameter counts are not trivial on SSC, as Table 15 shows several architectures within the ≈ 1M parameter range still outperformed by the methods proposed here. Moreover, as reported in [147], DelRec is based on delay 17
Official ranking is available at https://zenkelab.org/resources/spiking-heidelberg-datasets-shd/.
64
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
learning, which requires buffering past outputs, even as far back as 59 time steps. Therefore, while the recurrent state size is set to 256 (Table 15), the effective memory footprint is closer to the large ADPTNet configuration, strengthening the comparability of the two models. It should be mentioned that the accuracies reported in Table 15 contain only models trained on SSC without data augmentation. The highest accuracy recorded on SSC in Schöne et al. [154] relies on several spike-data augmentation techniques and does not use spiking activations within the model itself, leading to > 88% accuracy. Similarly, SpikCommander, proposed in Wang et al. [182], a Transformer-based spiking architecture, reaches at most 85.98%, using spike-dropping data augmentation. Interestingly, the 1.1M SpikCommander configuration reaches 83.26%, which, even with data augmentation, is below the comparably-sized SpikingADPTNet proposed here. Model
Acc(%)
State Size
N. Layers
N. Param.
RSNN [31] Adaptive RSNN [196] EventProp [130] RSNN with Adaptation [18] d-cAdLIF [41] SE-adLIF [11] DCLS [80] Adapt. Skip Rec. Connection SNN [191] SiLIF [54] RSNN DelRec [147]
51.1 ± 1.1 74.2 76.1 ± 1.0 77.4 80.23 ± 0.07 80.4 ± 0.3 80.7 ± 0.2 81.93 82.03 ± 0.25 82.42 ± 0.23
– – – – – – – – 512(×2) 256(*)
– – – – – – – – 2 3
– – – – 0.35M 1.6M 1.2M – 0.35M 0.37M
SpikingADPTNet
82.5 83.56 ± 0.15 (max : 83.76)
512(×2) 2048(×2)
2 2
355,689 1,069,481
Table 15: Accuracies on SSC The ranking of prior state-of-the-art work reported in this table is reported from the official SSC leaderboard (Available at https://zenkelab.org/resources/spiking-heidelberg-datasets-shd/). As shown, the smaller SpikingADPTNet configuration performs on par with the current state-of-the-art, while the larger one sets a new state-of-the-art by over 1 point. (*) denotes that the state size in practice is larger due to axonal/synaptic delays requiring explicit buffers of past outputs.
6
Discussion
Motivation Revisited As stated in Section 1, the goal of this work is to lay the foundations for large-scale neuromorphic systems as viable alternatives to the energy-intensive Transformer-GPU paradigm that dominates AI today. This requires an efficient neuromorphic alternative to the Transformer. Therefore, ADPTNet is proposed here, combining: (i) non-linear recurrent dynamics, (ii) adaptability, (iii) fine-grained parametric control over the time-scales of the system (to achieve long-term temporal modelling), and (iv) parallelisability. These four qualities are either already present and crucial to the Transformer’s performance (i.e., adaptability (ii), long-range sequence modelling (iii), and parallelisability (iv)), or a key shortcoming that sets it back (i.e., vanilla fixed-depth Transformers lack scalable non-linear recurrence (i), which is now considered important for tasks such as reasoning [103, 155, 63]). The challenge in combining all four properties within a single recurrent architecture is that they are, to some extent, antithetical (Fig. 1). Non-linear recurrence (i) impedes GPU-parallelism (iv). In addition, non-linear recurrence (i) and adaptability (ii) typically hinder fine-grained control over timescales, and thus long-range sequence modelling as well (iii). As shown through both theoretical proofs and empirical evidence throughout this study, ADPTNet either mitigates or outright removes these trade-offs. Non-Linearity and Parallelism Trade-Off Until recently, non-linear recurrence traditionally enforced sequential step-by-step simulation, which effectively prevented GPU-parallel training. With the popularisation of parallel-in-time simulation methods for RNNs, such as DEER, this limitation has gradually softened. Nevertheless, as highlighted in Gonzalez et al. [69], non-linear RNN parallelisation still depends on the properties of the dynamics, with penalties on larger Lipschitz constants and LLEs, and thus, implicitly on expressivity and long-term memory [51]. Moreover, training traditional non-linear RNNs such as LSTMs or GRUs is susceptible to bifurcations [46], whereby recurrent dynamics can unpredictably become chaotic and, thus, difficult to parallelise. While researchers have parallelised novel recurrent architectures by design in the past [56], ADPTNet is the first to use dynamical systems to provide theoretically principled control over its parallelisation. Namely, as prescribed in Sections 4.5 and 4.1, and empirically validated in Section 5.2, the step size β and recurrent eigenvalues Λ̄ reliably predict and control the network’s non-linearity and 65
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
long-term behaviour, and, thus, its DEER convergence behaviour. One can parametrically and gradually "dial" the level of parallelism vs non-linearity within ADPTNet, unlike existing non-linear RNNs. Furthermore, the ADPTNet recurrence rule requires only efficient low-rank/element-wise multiplications and two dense linear projections, for a total computational cost comparable to GRUs and LSTMs. However, knowing the recurrent eigenvalues a priori, and thus also long-term behaviour, enables more efficient adaptations to Quasi-DEER compared to those established architectures. Namely, this work introduces Conv-DEER, which requires zero computational overhead and no derivatives or Jacobians. In addition, Conv-DEER reduces memory overhead by using a fixed kernel for the entire batch, and allows a choice between FFT-based signal convolutions and parallel scans for computing DEER iterations. Even with the efficiency gains, Conv-DEER retains average convergence rates similar to Quasi-DEER. With little additional cost and still no Jacobians, Forward-DEER matches baseline convergence. To the authors’ knowledge, this is the first example of co-designing non-linear RNN dynamics and efficient parallel simulation algorithms, with direct user control over both. Adaptability and Long-Range Sequence Modelling Trade-Off As showcased by Mamba’s degradation in performance compared to S4 on long-range sequence modelling tasks [197], in recurrent architectures, adaptability can come at the cost of stable long-term memory. Existing methods such as Hawk [40] aim to mitigate this by disallowing arbitrarily strong forget gates and enforcing a lower bound on memory decay rates. In principle, ADPTNet follows a similar strategy. Information can only decay as fast as the fastest recurrent eigenvalue, perturbed by a term controlled by β. However, in practice, ADPTNet outperforms Hawk in adaptability, as quantified by the superior accuracy on the Selective Copying task (Section 5.6). Furthermore, Section 5.7 shows how this does not come at the cost of longrange sequence modelling performance. Both linear and non-linear ADPTNet outperform existing selective/adaptive architectures, while using significantly fewer parameters on sCIFAR, matching baseline LTI SSM performance. For both the MAD Selective Copying benchmark and sCIFAR, ADPTNet employs complex S4D-Inv and a perturbation-level β. Hence, interestingly, this study shows for the first time how linear oscillators with loose but dynamic and non-linear coupling can enable competitive selectivity compared to traditional forget gates, while retaining long-range memory. This finding contributes to the recently increasing interest in oscillators as a computational primitive in sequence models [36, 176]. Non-Linearity and Fine-Grained Timescale Control Trade-Off Currently, controlling the non-linear dynamics of RNNs for long-range sequence modelling has been a matter of shaping the LLE. As showcased in Section 5.1, if one uses RNNs with orthogonally-constrained recurrent weights to set the LLE close to zero, this still does not prevent an arbitrarily quickly decaying rest of the Lyapunov spectrum. That differs from LTI SSMs where the entire spectrum can be fine-tuned to initialise and parametrise timescales according to task needs. The closest existing solution to attaining a form of full-spectrum control is Gradient Flossing [47]. Still, as detailed in Section 2.3, Gradient Flossing is a regularisation method relying on a loss term "nudging" dynamics towards the desired spectrum. It is sensitive to, and to some extent reliant on, the trajectories sampled for training. Furthermore, without continually optimising the flossing objective, the spectrum drifts unpredictably from its target. In terms of computational cost, it requires not only materialising full recurrent Jacobians but also computing their QR decomposition up to the number of exponents targeted for flossing, a cubically scaling operation. By contrast, ADPTNet has theoretical guarantees on parametric control of its Lyapunov spectrum without the need for loss-based regularisation, without spectral drift over training, and regardless of inputs. Lemmas 2 and 3 provide the proofs for these guarantees, and Section 5.1 reinforces them with empirical evidence. Furthermore, Section 5.7 shows that non-linear ADPTNet networks can still retain SSM state-of-the-art long-range sequence modelling performance, enabled by adopting their powerful recurrent eigenspectrum parametrisation and initialisation. At the same time, Section 5.4 shows that, while ADPTNet can offer guarantees on its Lyapunov spectrum, it is still sufficiently non-linear to perform state-tracking. Complementary to how β can smoothly "dial" the selectivity and parallelisability of the network (Sections 5.2 and 5.5), improvements in state tracking show by proxy how β also incrementally increases non-linearity. It is worth emphasising how this interpolation behaviour differs from existing Almost Linear RNNs [23]. Existing methods tune non-linearity by applying non-linear activations only to a subset of recurrent neurons (e.g., 10% of neurons). That requires discrete increments in non-linear units, while β is a continuous control. This enables, by contrast with Almost Linear RNNs, direct tuning of β via gradient descent, allowing the network to set the level of non-linearity appropriate for the task (Table 12). State-of-the-Art Neuromorphic Speech Processing ADPTNet takes inspiration from the computational and dynamical systems properties of the auditory cortex. This work argues that the input-invariant time scales in the auditory cortex can be abstracted as predictable Lyapunov exponents in a non-linear, adaptive RNN. In Section 5.8, it is shown how these brain-inspired principles lead to state-of-the-art accuracy on the SSC dataset, a de facto standard for benchmarking neuromorphic speech processing. Furthermore, it does so without the need to include large explicit buffers for delays as the previous state-of-the-art, which become more and more difficult to scale with sequence length (see Section 3.3).
66
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
SpikingADPTNet also outperforms state-of-the-art spiking Transformers on SSC (the SpikCommander), given a comparable parameter budget, without the need for data augmentation. Within the broader deep learning research landscape, ADPTNet is the first architecture to employ a theoretically principled adaptation of linear Transformers/DeltaNet online optimisation rules to Riemannian manifolds. Furthermore, this is the first study to connect modern linear Transformers/DeltaNet to dynamical systems theory using Chaos Theory (i.e., Lyapunov spectra). This latter connection extends the principles behind Neuromorphic Intermediate Representations (NIR) [143]. NIR connects heterogeneous neuromorphic hardware and software platforms through dynamical systems formalisms. This study likewise connects brain computing principles, GPU-parallelisation, longrange sequence modelling, and state-of-the-art notions of adaptability/selectivity through the shared framework of dynamical systems theory. Significance ADPTNet pushes forward the state-of-the-art in multiple distinct fields. It advances sequence modelling as the first combination of linear DeltaNet concepts with principles from dynamical systems theory, particularly Lyapunov exponents, and Riemannian manifold optimisation. It proposes a novel, more efficient paradigm for parallelising nonlinear dynamics that exploits highly controllable long-term dynamics to reduce computational and memory overhead and offer fine-grained control over parallelisation properties. It also contributes to the RNN vanishing/exploding gradients mitigation literature, namely Gradient Flossing, as a more efficient alternative that achieves stable gradients by construction, not regularisation, to significantly improve efficacy and reliability. In addition, SpikingADPTNet sets a direct state-of-the-art result on real-world neuromorphic speech processing (83.56% ± 0.15), outperforming the current state-of-the-art spiking Transformer in the process. Ultimately, all the steps taken within this work bring neuromorphic architectures closer to competing with the Transformer, even outperforming it on long-range dependency modelling and state tracking. As a result, ADPTNet is also a step towards a viable neuromorphic and energy-efficient alternative to current LLM systems. This is a vital research direction, as an uncontrolled rise in AI energy consumption can have severe negative effects on both the environment and society at large.
7
Future Work
Limitations Although ADPTNet makes strides towards capturing the qualities that support the Transformer’s dominance, several challenges remain. First, while ADPTNet outperforms its closest existing equivalent, Hawk, it remains unclear whether it can outperform other adaptive networks on Selective Copying, since only a small perturbative β = 0.00125 is tested here. Evidence from Table 14 suggests testing larger β should further close that performance gap. Second, while Selective Copy (and MAD) and state tracking are synthetic predictors of language modelling and reasoning performance, respectively, they are not sufficient to conclude whether ADPTNet performs on par with the Transformer at scale. A more informative test would be to scale ADPTNet to billions of parameters and directly test it on full language modelling and reasoning benchmarks. Third, the ADPTNet and Conv/Forward DEER implementations are in PyTorch, which may be suboptimal compared to fused CUDA kernels and may also be limiting scalability at the moment. Generally, to fully evaluate the potential of ADPTNet, more compute than was available for the scope of this work is needed. Koopman-DEER Briefly, Koopman Operator theory [112] posits that any non-linear dynamics can be approximated by a linear operator K if lifted to an appropriate infinite-dimensional state space. In practice, a high-dimensional embedding, such as a delay embedding or the recurrent state of an SSM or ADPTNet, can serve as a proxy for computing the linear approximator K ≈ K, which can then effectively describe the system’s non-linear behaviour. The topological conjugates that underpin ADPTNet recurrence have an established connection to the Koopman Operator. Dynamical Similarity Analysis (DSA) [141] uses orthogonally-constrained topological conjugates QKQT to define a metric distance between different non-linear dynamics. In general, finding a Koopman operator K is done using Dynamic Mode Decomposition (DMD) [175], which is essentially a linear least-squares problem of the form ∥(st − Kst−1 )∥2 , where st are all the states in a trajectory 18 . One direction for future work is to build on Conv-DEER toward a Koopman-DEER. Conv-DEER instability largely stems from errors caused by not considering the off-diagonal interactions between recurrent states. A Koopman-DEER implementation would similarly minimise computational overhead compared to traditional Quasi-DEER, but use a single dense matrix approximator K instead of the diagonal eigenvalues Λ̄. Given the QΛ̄QT structure of ADPTNet recurrent dynamics, finding an ADPTNet Koopman operator K = K̂ Λ̄K̂ T may be a question of finding an optimal averaging of rotations K̂ ≈ mean(Q1 , Q2 , . . . , Qt ). Computing DEER iterations is then just a matter of finding K t terms, which are K̂ Λ̄t K̂ T , enjoying similar efficiency and convolution equivalence to Conv-DEER. 18
It is worth mentioning that the DMD optimisation problem bears striking resemblance to the residual minimisation problem within DEER.
67
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
Proximal-DEER More generally, future work could also explore taking advantage of manifold structure within the DEER update itself. In other words, DEER relies on unconstrained Newton steps in the direction minimising the residual. However, in networks such as ADPTNet, each recurrent step has a known and highly structured form. One could potentially devise new architectures and accompanying DEER extensions that can take advantage of Riemannian optimisation techniques to constrain DEER updates to valid manifold-bound state guesses. Depending on the constraints, one can use efficient proximal optimisation algorithms with accelerated convergence [67]. Topological Conjugate / Manifold-Constrained Online and Local Learning Rules ADPTNet relies on similarity transforms and topological conjugations of its recurrent dynamics at each timestep. However, one could space out the updates and accumulate gradient/learning signals over larger sequences. This could yield an online learning rule similar to e-prop [16] or STDP [17], where, however, optimisation updates are constrained to the orthogonal manifold and used to topologically conjugate (i.e. apply similarity transforms to) SNN/RNN recurrent weights. This would mean that the fundamental timescales of the networks remain unchanged throughout learning, avoiding vanishing/exploding history compression pathologies. Improve Energy-Efficiency As presented, SpikingADPTNet could be further optimised for energy efficiency by adopting more neuromorphic computing principles. For instance, while ReLU is currently applied before obtaining k and v, a spiking activation could avoid the dense and expensive K and V vector matrix multiplications. Furthermore, the loose coupling induced by the topological conjugates could also be optimised for sparsity and could eventually be derived stochastically rather than deterministically, to reduce memory overhead. Dynamical Systems Modelling The focus of this work is on stable long-range sequence modelling, and thus the eigespectrum is parametrised accordingly. However, the ADPTNet recurrence imposes no inherent constraints. One could also explore ADPTNet modelling of chaotic dynamical systems. Systems such as the Lorenz 63 attractor have well-known positive Lyapunov spectra which could be used as part of ADPTNet initialisation or as a form of inductive bias towards learning unstable dynamics. Furthermore, one could use the ADPTNet topological conjugate primitive to build interpretable models of different dynamical system topologies, with different similarity transforms for individual critical points/attractor basins in the spirit of recurrent switching SSMs [120]. Tighter Lyapunov Exponent Bounds As shown in Section 5.1, the Lyapunov spectra empirically observed for ADPTNet at initialisation and during training lie within much tighter bounds of the eigenspectrum, relative to the limits prescribed by Lemma 2. Future work should establish whether the theoretical bounds could be further tightened, or whether the observed behaviour is an idiosyncrasy of the data and model setup used in this study. Architectural Extensions and Scaling One of the key limits on ADPTNet expressiveness is the low-rank structure of Key-Value outer products projected onto the orthogonal manifold. For instance, Siems et al. [160] showed how taking multiple gradient descent steps over the DeltaNet learning objective improved state-tracking performance. Similar results could be a useful future direction for ADPTNet. In addition, ADPTNet currently relies on Riemannian gradient descent steps from the origin In at each time step, which may also limit expressiveness. It may be worth exploring alternative starting points, or even dynamic starting points, potentially with matrix-valued memory similar to DeltaNet-like state-of-the-art recurrent architectures. Ultimately, the goal of ADPTNet is to propose computational alternatives to the Transformer. This evidently requires testing ADPTNet at billion-parameter scale
8
Author Contributions Statement
M.I.S. conceived the models under investigation, conducted the experiments, contributed to their design, prepared figures, and wrote the main manuscript text. O.R. supervised the research, contributed to devising the experiments and provided major revisions to the final manuscript text and figures. All authors reviewed the final manuscript.
9
Data Availability Statement
The datasets used in this study are publicly available:
10
Additional Information
Competing interests: The authors declare no competing interests. 68
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
References [1] Pierre Ablin and Gabriel Peyré. Fast and accurate optimization on the orthogonal manifold without retraction. In International Conference on Artificial Intelligence and Statistics, pages 5636–5657. PMLR, 2022. [2] Steven Abreu, Jens E Pedersen, Kade M Heckel, and Alessandro Pierro. Q-s5: Towards quantized state space models. arXiv preprint arXiv:2406.09477, 2024. [3] P-A Absil, Robert Mahony, and Rodolphe Sepulchre. Optimization algorithms on matrix manifolds. Princeton University Press, 2008. [4] AbdalQader AlKilany and Dan FM Goodman. Neuromodulation enhances dynamic sensory processing in spiking neural network models. bioRxiv, pages 2025–07, 2025. [5] Arnon Amir, Brian Taba, David Berg, Timothy Melano, Jeffrey McKinstry, Carmelo Di Nolfo, Tapan Nayak, Alexander Andreopoulos, Guillaume Garreau, Marcela Mendoza, et al. A low power, fully event-based gesture recognition system. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 7243–7252, 2017. [6] Martin Arjovsky, Amar Shah, and Yoshua Bengio. Unitary evolution recurrent neural networks. In International conference on machine learning, pages 1120–1128. PMLR, 2016. [7] Simran Arora, Sabri Eyuboglu, Aman Timalsina, Isys Johnson, Michael Poli, James Y Zou, Atri Rudra, and Christopher Ré. Zoology: Measuring and improving recall in efficient language models. In International conference on learning representations, volume 2024, pages 15664–15730, 2024. [8] Guangji Bai, Zheng Chai, Chen Ling, Shiyu Wang, Jiaying Lu, Nan Zhang, Tingwei Shi, Ziyang Yu, Mengdan Zhu, Yifei Zhang, et al. Beyond efficiency: A systematic survey of resource-efficient large language models. arXiv preprint arXiv:2401.00625, 2024. [9] Shaojie Bai, J Zico Kolter, and Vladlen Koltun. Trellis networks for sequence modeling. arXiv preprint arXiv:1810.06682, 2018. [10] Malyaban Bal and Abhronil Sengupta. P-spikessm: Harnessing probabilistic spiking state space models for long-range dependency tasks. In International Conference on Learning Representations, volume 2025, pages 73912–73927, 2025. [11] Maximilian Baronig, Romain Ferrand, Silvester Sabathiel, and Robert Legenstein. Advancing spatio-temporal processing through adaptation in spiking neural networks. Nature Communications, 16(1):5776, 2025. [12] Maximilian Beck, Korbinian Pöppel, Markus Spanring, Andreas Auer, Oleksandra Prudnikova, Michael Kopp, Günter Klambauer, Johannes Brandstetter, and Sepp Hochreiter. xlstm: Extended long short-term memory. Advances in Neural Information Processing Systems, 37:107547–107603, 2024. [13] Ali Behrouz, Meisam Razaviyayn, Peilin Zhong, and Vahab Mirrokni. It’s all connected: A journey through test-time memorization, attentional bias, retention, and online optimization. arXiv preprint arXiv:2504.13173, 2025. [14] Guillaume Bellec, David Kappel, Wolfgang Maass, and Robert Legenstein. Deep rewiring: Training very sparse deep networks. arXiv preprint arXiv:1711.05136, 2017. [15] Guillaume Bellec, Darjan Salaj, Anand Subramoney, Robert Legenstein, and Wolfgang Maass. Long short-term memory and learning-to-learn in networks of spiking neurons. Advances in neural information processing systems, 31, 2018. [16] Guillaume Bellec, Franz Scherr, Anand Subramoney, Elias Hajek, Darjan Salaj, Robert Legenstein, and Wolfgang Maass. A solution to the learning dilemma for recurrent networks of spiking neurons. Nature communications, 11(1):3625, 2020. [17] Yoshua Bengio, Thomas Mesnard, Asja Fischer, Saizheng Zhang, and Yuhuai Wu. Stdp as presynaptic activity times rate of change of postsynaptic activity. arXiv preprint arXiv:1509.05936, 2015. [18] Alexandre Bittar and Philip N Garner. A surrogate gradient spiking baseline for speech command recognition. Frontiers in Neuroscience, 16:865897, 2022. [19] Guy E Blelloch. Scans as primitive parallel operations. IEEE Transactions on computers, 38(11):1526–1538, 2002. [20] Aleksandar Botev, Soham De, Samuel L Smith, Anushan Fernando, George-Cristian Muraru, Ruba Haroun, Leonard Berrada, Razvan Pascanu, Pier Giuseppe Sessa, Robert Dadashi, et al. Recurrentgemma: Moving past transformers for efficient open language models. arXiv preprint arXiv:2404.07839, 2024. 69
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
[21] Stephen Boyd and Leon Chua. Fading memory and the problem of approximating nonlinear operators with volterra series. IEEE Transactions on circuits and systems, 32(11):1150–1161, 1985. [22] Michael Breakspear and Viktor K Jirsa. Neuronal dynamics and brain connectivity. In Handbook of brain connectivity, pages 3–64. Springer, 2007. [23] Manuel Brenner, Christoph Jürgen Hemmer, Zahra Monfared, and Daniel Durstewitz. Almost-linear rnns yield highly interpretable symbolic codes in dynamical systems reconstruction. Advances in Neural Information Processing Systems, 37:36829–36868, 2024. [24] John E Campbell. On a law of combination of operators bearing on the theory of continuous transformation groups. Proceedings of the London Mathematical Society, 1(1):381–390, 1896. [25] Aakshita Chandiramani, Aaron Blakeman, Abdullahi Olaoye, Abhibha Gupta, Abhilash Somasamudramath, Abhinav Khattar, Adeola Adesoba, Adi Renduchintala, Adil Asif, Aditya Agrawal, et al. Nemotron 3 super: Open, efficient mixture-of-experts hybrid mamba-transformer model for agentic reasoning. arXiv preprint arXiv:2604.12374, 2026. [26] Vinod Kumar Chauhan, Jiandong Zhou, Ping Lu, Soheila Molaei, and David A Clifton. A brief review of hypernetworks in deep learning. Artificial Intelligence Review, 57(9):250, 2024. [27] Narsimha Reddy Chilkuri and Chris Eliasmith. Parallelizing legendre memory unit training. In International conference on machine learning, pages 1898–1907. PMLR, 2021. [28] Kyunghyun Cho, Bart Van Merriënboer, Dzmitry Bahdanau, and Yoshua Bengio. On the properties of neural machine translation: Encoder–decoder approaches. In Proceedings of SSST-8, eighth workshop on syntax, semantics and structure in statistical translation, pages 103–111, 2014. [29] James W Cooley and John W Tukey. An algorithm for the machine calculation of complex fourier series. Mathematics of computation, 19(90):297–301, 1965. [30] Julia C Costacurta, Shaunak Bhandarkar, David Zoltowski, and Scott W Linderman. Structured flexibility in recurrent neural networks via neuromodulation. Advances in Neural Information Processing Systems, 37: 1954–1972, 2024. [31] Benjamin Cramer, Yannik Stradmann, Johannes Schemmel, and Friedemann Zenke. The heidelberg spiking data sets for the systematic evaluation of spiking neural networks. IEEE Transactions on Neural Networks and Learning Systems, 33(7):2744–2757, 2020. [32] Mohammed Dahleh, Munther A Dahleh, and George Verghese. Lectures on dynamic systems and control. A+ A, 4(100):1–100, 2004. [33] Federico Danieli, Pau Rodriguez, Miguel Sarabia, Xavier Suau, and Luca Zappella. Pararnn: Unlocking parallel training of nonlinear rnns for large language models. arXiv preprint arXiv:2510.21450, 2025. [34] Tri Dao and Albert Gu. Transformers are ssms: Generalized models and efficient algorithms through structured state space duality. arXiv preprint arXiv:2405.21060, 2024. [35] Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-efficient exact attention with io-awareness. Advances in neural information processing systems, 35:16344–16359, 2022. [36] Luke Darlow, Ciaran Regan, Sebastian Risi, Jeffrey Seely, and Llion Jones. Continuous thought machines. Advances in Neural Information Processing Systems, 38:21548–21594, 2026. [37] Yann N Dauphin, Angela Fan, Michael Auli, and David Grangier. Language modeling with gated convolutional networks. In International conference on machine learning, pages 933–941. PMLR, 2017. [38] Kenneth R Davidson. Estimating the distance between unitary orbits. Journal of Operator Theory, pages 21–40, 1988. [39] Mike Davies, Andreas Wild, Garrick Orchard, Yulia Sandamirskaya, Gabriel A Fonseca Guerra, Prasad Joshi, Philipp Plank, and Sumedh R Risbud. Advancing neuromorphic computing with loihi: A survey of results and outlook. Proceedings of the IEEE, 109(5):911–934, 2021. [40] Soham De, Samuel L. Smith, Anushan Fernando, Aleksandar Botev, George Cristian-Muraru, Albert Gu, Ruba Haroun, Leonard Berrada, Yutian Chen, Srivatsan Srinivasan, Guillaume Desjardins, Arnaud Doucet, David Budden, Yee Whye Teh, Razvan Pascanu, Nando De Freitas, and Caglar Gulcehre. Griffin: Mixing gated linear recurrences with local attention for efficient language models, 2024. URL https://arxiv.org/abs/2402. 19427. [41] Lucas Deckers, Laurens Van Damme, Werner Van Leekwijck, Ing Jyh Tsang, and Steven Latré. Co-learning synaptic delays, weights and adaptation in spiking neural networks. Frontiers in Neuroscience, 18:1360300, 2024. 70
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
[42] Rodney J Douglas, Christof Koch, Misha Mahowald, Kevan AC Martin, and Humbert H Suarez. Recurrent excitation in neocortical circuits. Science, 269(5226):981–985, 1995. [43] Scott C Douglas and Jiutian Yu. Why relu units sometimes die: Analysis of single-unit error backpropagation in neural networks. In 2018 52nd Asilomar conference on signals, systems, and computers, pages 864–868. IEEE, 2018. [44] Ali Edalati, Marzieh Tahaei, Ivan Kobyzev, Vahid Partovi Nia, James J Clark, and Mehdi Rezagholizadeh. Krona: Parameter-efficient tuning with kronecker adapter. In Enhancing LLM Performance: Efficacy, Fine-Tuning, and Inference Techniques, pages 49–65. Springer, 2025. [45] Alan Edelman. Eigenvalues and condition numbers of random matrices. SIAM journal on matrix analysis and applications, 9(4):543–560, 1988. [46] Lukas Eisenmann, Zahra Monfared, Niclas Göring, and Daniel Durstewitz. Bifurcations and loss jumps in rnn training. Advances in Neural Information Processing Systems, 36:70511–70547, 2023. [47] Rainer Engelken. Gradient flossing: Improving gradient descent through dynamic control of jacobians. Advances in Neural Information Processing Systems, 36:10412–10439, 2023. [48] Rainer Engelken. Sparseprop: Efficient event-based simulation and training of sparse recurrent spiking neural networks. Advances in Neural Information Processing Systems, 36:3638–3657, 2023. [49] Rainer Engelken and Larry Abbott. Analyzing and improving surrogate gradient training in binary neural networks using dynamical systems theory. In ICML 2024 Workshop on Differentiable Almost Everything: Differentiable Relaxations, Algorithms, Operators, and Simulators. [50] Rainer Engelken, Fred Wolf, and Larry F Abbott. Lyapunov spectra of chaotic recurrent neural networks. Physical Review Research, 5(4):043044, 2023. [51] N Benjamin Erichson, Omri Azencot, Alejandro Queiruga, Liam Hodgkinson, and Michael W Mahoney. Lipschitz recurrent neural networks. arXiv preprint arXiv:2006.12070, 2020. [52] Thomas Erneux. Applied delay differential equations. Springer, 2009. [53] Jason K Eshraghian, Max Ward, Emre O Neftci, Xinxin Wang, Gregor Lenz, Girish Dwivedi, Mohammed Bennamoun, Doo Seok Jeong, and Wei D Lu. Training spiking neural networks using lessons from deep learning. Proceedings of the IEEE, 111(9):1016–1054, 2023. [54] Maxime Fabre, Lyubov Dudchenko, and Emre Neftci. Structured state space model dynamics and parametrization for spiking neural networks. arXiv preprint arXiv:2506.06374, 2025. [55] Wei Fang, Zhaofei Yu, Yanqi Chen, Tiejun Huang, Timothée Masquelier, and Yonghong Tian. Deep residual learning in spiking neural networks. Advances in neural information processing systems, 34:21056–21069, 2021. [56] Mónika Farsang and Radu Grosu. Parallelization of non-linear state-space models: Scaling up liquid-resistance liquid-capacitance networks for efficient sequence modeling. Advances in Neural Information Processing Systems, 38:169137–169163, 2026. [57] Wanjin Feng, Xingyu Gao, Wenqian Du, Hailong Shi, Peilin Zhao, Pengcheng Wu, and Chunyan Miao. Efficient parallel training methods for spiking neural networks with constant time complexity. arXiv preprint arXiv:2506.12087, 2025. [58] Dan Fu, Hermann Kumbong, Eric Nguyen, and Christopher Ré. Flashfftconv: Efficient convolutions for long sequences with tensor cores. In International Conference on Learning Representations, volume 2024, pages 9455–9483, 2024. [59] Daniel Y Fu, Tri Dao, Khaled K Saab, Armin W Thomas, Atri Rudra, and Christopher Ré. Hungry hungry hippos: Towards language modeling with state space models. arXiv preprint arXiv:2212.14052, 2022. [60] Trevor Gale, Erich Elsen, and Sara Hooker. The state of sparsity in deep neural networks. arXiv preprint arXiv:1902.09574, 1(2):3, 2019. [61] Martin J Gander. 50 years of time parallel time integration. In Multiple Shooting and Time Domain Decomposition Methods: MuS-TDD, Heidelberg, May 6-8, 2013, pages 69–113. Springer, 2015. [62] Mario Garrido. A survey on pipelined fft hardware architectures. Journal of Signal Processing Systems, 94(11): 1345–1364, 2022. [63] Jonas Geiping, Sean McLeish, Neel Jain, John Kirchenbauer, Siddharth Singh, Brian Bartoldson, Bhavya Kailkhura, Abhinav Bhatele, and Tom Goldstein. Scaling up test-time compute with latent reasoning: A recurrent depth approach. Advances in Neural Information Processing Systems, 38:41340–41391, 2026. 71
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
[64] Karlheinz Geist, Ulrich Parlitz, and Werner Lauterborn. Comparison of different methods for computing lyapunov exponents. Progress of theoretical physics, 83(5):875–893, 1990. [65] Wulfram Gerstner, Werner M Kistler, Richard Naud, and Liam Paninski. Neuronal dynamics: From single neurons to networks and models of cognition. Cambridge University Press, 2014. [66] Mor Geva, Roei Schuster, Jonathan Berant, and Omer Levy. Transformer feed-forward layers are key-value memories. In Proceedings of the 2021 conference on empirical methods in natural language processing, pages 5484–5495, 2021. [67] Anand Gokhale, Alexander Davydov, and Francesco Bullo. Proximal gradient dynamics: Monotonicity, exponential convergence, and applications. IEEE Control Systems Letters, 8:2853–2858, 2024. [68] Xavier Gonzalez, Andrew Warrington, Jimmy T Smith, and Scott W Linderman. Towards scalable and stable parallelization of nonlinear rnns. Advances in Neural Information Processing Systems, 37:5817–5849, 2024. [69] Xavier Gonzalez, Leo Kozachkov, David Zoltowski, Kenneth Clarkson, and Scott Linderman. Predictability enables parallelization of nonlinear state space models. Advances in Neural Information Processing Systems, 38: 19101–19147, 2026. [70] Riccardo Grazzi, Julien Siems, Arber Zela, Jorg KH Franke, Frank Hutter, and Massimiliano Pontil. Unlocking state-tracking in linear rnns through negative eigenvalues. In 13th International Conference on Learning Representations Iclr 2025, pages 1–33. ICLR, 2025. [71] Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00752, 2023. [72] Albert Gu, Tri Dao, Stefano Ermon, Atri Rudra, and Christopher Ré. Hippo: Recurrent memory with optimal polynomial projections. Advances in neural information processing systems, 33:1474–1487, 2020. [73] Albert Gu, Caglar Gulcehre, Thomas Paine, Matt Hoffman, and Razvan Pascanu. Improving the gating mechanism of recurrent neural networks. In International conference on machine learning, pages 3800–3809. PMLR, 2020. [74] Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. arXiv preprint arXiv:2111.00396, 2021. [75] Albert Gu, Isys Johnson, Karan Goel, Khaled Saab, Tri Dao, Atri Rudra, and Christopher Ré. Combining recurrent, convolutional, and continuous-time models with linear state space layers. Advances in neural information processing systems, 34:572–585, 2021. [76] Albert Gu, Karan Goel, Ankit Gupta, and Christopher Ré. On the parameterization and initialization of diagonal state space models. Advances in neural information processing systems, 35:35971–35983, 2022. [77] Albert Gu, Isys Johnson, Aman Timalsina, Atri Rudra, and Christopher Ré. How to train your hippo: State space models with generalized orthogonal basis projections. arXiv preprint arXiv:2206.12037, 2022. [78] David Ha, Andrew Dai, and Quoc V Le. Hypernetworks. arXiv preprint arXiv:1609.09106, 2016. [79] Danijar Hafner, Alexander Irpan, James Davidson, and Nicolas Heess. Learning hierarchical information flow with recurrent neural modules. Advances in Neural Information Processing Systems, 30, 2017. [80] Ilyass Hammouamri, Ismail Khalfaoui Hassani, and Timothée Masquelier. Learning delays in spiking neural networks using dilated convolutions with learnable spacings. In International conference on learning representations, volume 2024, pages 17890–17903, 2024. [81] Behnam Hashemi and Yuji Nakatsukasa. Instability of the sherman-morrison formula and stabilization by iterative refinement. arXiv preprint arXiv:2510.01696, 2025. [82] Jennifer Hasler. Special report: Can we copy the brain?-a road map for the artificial brain. IEEE Spectrum, 54 (6):46–50, 2017. [83] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human-level performance on imagenet classification. In Proceedings of the IEEE international conference on computer vision, pages 1026–1034, 2015. [84] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770–778, 2016. [85] Donald Olding Hebb. The organization of behavior: A neuropsychological theory. Psychology press, 2005. [86] Kyle Helfrich, Devin Willmott, and Qiang Ye. Orthogonal recurrent neural networks with scaled cayley transform. In International Conference on Machine Learning, pages 1969–1978. PMLR, 2018. 72
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
[87] Russell L Herman. A second course in ordinary differential equations of dynamical systems and boundary value problems, 2008. [88] Michiel Hermans and Benjamin Schrauwen. Memory in linear recurrent neural networks in continuous time. Neural Networks, 23(3):341–355, 2010. [89] Florian Hess, Zahra Monfared, Manuel Brenner, and Daniel Durstewitz. Generalized teacher forcing for learning chaotic dynamics. arXiv preprint arXiv:2306.04406, 2023. [90] Fumio Hiai and Yoshihiro Nakamura. Distance between unitary orbits in von neumann algebras. 1989. [91] Nicholas J Higham. Functions of matrices: theory and computation. SIAM, 2008. [92] Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 9(8):1735–1780, 1997. [93] Sepp Hochreiter, Yoshua Bengio, Paolo Frasconi, Jürgen Schmidhuber, et al. Gradient flow in recurrent nets: the difficulty of learning long-term dependencies, 2001. [94] Sara Hooker. The hardware lottery. Communications of the ACM, 64(12):58–65, 2021. [95] John J Hopfield. Neural networks and physical systems with emergent collective computational abilities. Proceedings of the national academy of sciences, 79(8):2554–2558, 1982. [96] Dichao Hu. An introductory survey on attention mechanisms in nlp problems. In Proceedings of SAI intelligent systems conference, pages 432–448. Springer, 2019. [97] Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. Lora: Low-rank adaptation of large language models, 2021. URL https://arxiv. org/abs/2106.09685, 2106, 2021. [98] Zheyuan Hu, Nazanin Ahmadi Daryakenari, Qianli Shen, Kenji Kawaguchi, and George Em Karniadakis. State-space models are accurate and efficient neural operators for dynamical systems. Neural Networks, page 108496, 2025. [99] Eugene Izhikevich. Spiking manifesto. arXiv preprint arXiv:2512.11843, 2025. [100] Eugene M Izhikevich. Polychronization: computation with spikes. Neural computation, 18(2):245–282, 2006. [101] Samy Jelassi, David Brandfonbrener, Sham M Kakade, and Eran Malach. Repeat after me: Transformers are better than state space models at copying. arXiv preprint arXiv:2402.01032, 2024. [102] Li Jing, Caglar Gulcehre, John Peurifoy, Yichen Shen, Max Tegmark, Marin Soljacic, and Yoshua Bengio. Gated orthogonal recurrent units: On learning to forget. Neural computation, 31(4):765–783, 2019. [103] Alexia Jolicoeur-Martineau. Less is more: Recursive reasoning with tiny networks. arXiv preprint arXiv:2510.04871, 2025. [104] Nal Kalchbrenner, Lasse Espeholt, Karen Simonyan, Aaron van den Oord, Alex Graves, and Koray Kavukcuoglu. Neural machine translation in linear time. arXiv preprint arXiv:1610.10099, 2016. [105] Louis Kang and Taro Toyoizumi. Distinguishing examples while building concepts in hippocampal and artificial networks. Nature Communications, 15(1):647, 2024. [106] Mahdi Karami, Razvan Pascanu, and Vahab Mirrokni. Lattice: Learning to efficiently compress the memory. arXiv preprint arXiv:2504.05646, 2025. [107] Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are rnns: Fast autoregressive transformers with linear attention. In International conference on machine learning, pages 5156–5165. PMLR, 2020. [108] Menoua Keshishian, Hassan Akbari, Bahar Khalighinejad, Jose L Herrero, Ashesh D Mehta, and Nima Mesgarani. Estimating and interpreting nonlinear receptive field of sensory neural responses with deep neural network models. Elife, 9:e53445, 2020. [109] Hyunjik Kim, George Papamakarios, and Andriy Mnih. The lipschitz constant of self-attention. In International Conference on Machine Learning, pages 5562–5571. PMLR, 2021. [110] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014. [111] Christian Klos and Raoul-Martin Memmesheimer. Smooth exact gradient descent learning in spiking neural networks. Physical Review Letters, 134(2):027301, 2025. [112] Bernard O Koopman. Hamiltonian systems and transformation in hilbert space. Proceedings of the National Academy of Sciences, 17(5):315–318, 1931. 73
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
[113] Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009. [114] Dhireesha Kudithipudi, Catherine Schuman, Craig M Vineyard, Tej Pandit, Cory Merkel, Rajkumar Kubendran, James B Aimone, Garrick Orchard, Christian Mayr, Ryad Benosman, et al. Neuromorphic computing at scale. Nature, 637(8047):801–812, 2025. [115] Aakash Sunil Lahoti, Kevin Li, Berlin Chen, Caitlin Wang, Aviv Bick, Zico Kolter, Tri Dao, and Albert Gu. Mamba-3: Improved sequence modeling using state space principles. In International Conference on Learning Representations, volume 2026, pages 86173–86202, 2026. [116] James Lee-Thorp, Joshua Ainslie, Ilya Eckstein, and Santiago Ontanon. Fnet: Mixing tokens with fourier transforms. In Proceedings of the 2022 Conference of the north American chapter of the Association for Computational Linguistics: human language technologies, pages 4296–4313, 2022. [117] Yuhong Li, Tianle Cai, Yi Zhang, Deming Chen, and Debadeepta Dey. What makes convolutional models great on long sequence modeling? arXiv preprint arXiv:2210.09298, 2022. [118] Yi Heng Lim, Qi Zhu, Joshua Selfridge, and Muhammad Firmansyah. Parallelizing non-linear sequential models over the sequence length. In International Conference on Learning Representations, volume 2024, pages 55334–55360, 2024. [119] Scott Linderman, Annika Nichols, David Blei, Manuel Zimmer, and Liam Paninski. Hierarchical recurrent state space models reveal discrete and continuous dynamics of neural activity in c. elegans. BioRxiv, page 621540, 2019. [120] Scott W Linderman, Andrew C Miller, Ryan P Adams, David M Blei, Liam Paninski, and Matthew J Johnson. Recurrent switching linear dynamical systems. arXiv preprint arXiv:1610.08466, 2016. [121] Hao Liu and Pieter Abbeel. Blockwise parallel transformers for large context models. Advances in neural information processing systems, 36:8828–8844, 2023. [122] Xuezhe Ma, Chunting Zhou, Xiang Kong, Junxian He, Liangke Gui, Graham Neubig, Jonathan May, and Luke Zettlemoyer. Mega: Moving average equipped gated attention. arXiv preprint arXiv:2209.10655, 2022. [123] Wolfgang Maass. Networks of spiking neurons: the third generation of neural network models. Neural networks, 10(9):1659–1671, 1997. [124] Brianna Marsh, M Gabriela Navas-Zuloaga, Burke Q Rosen, Yury Sokolov, Jean Erik Delanois, Oscar C Gonzalez, Giri P Krishnan, Eric Halgren, and Maxim Bazhenov. Emergent effects of synaptic connectivity on the dynamics of global and local slow waves in a large-scale thalamocortical network model of the human brain. PLoS computational biology, 20(7):e1012245, 2024. [125] James Martens and Roger Grosse. Optimizing neural networks with kronecker-factored approximate curvature. In International conference on machine learning, pages 2408–2417. PMLR, 2015. [126] Stefano Massaroli, Michael Poli, Jinkyoo Park, Atsushi Yamashita, and Hajime Asama. Dissecting neural odes. Advances in neural information processing systems, 33:3952–3963, 2020. [127] Francesca Mastrogiuseppe and Srdjan Ostojic. Linking connectivity, dynamics, and computations in low-rank recurrent neural networks. Neuron, 99(3):609–623, 2018. [128] William Merrill and Ashish Sabharwal. The parallelism tradeoff: Limitations of log-precision transformers. Transactions of the Association for Computational Linguistics, 11:531–545, 2023. [129] William Merrill, Jackson Petty, and Ashish Sabharwal. The illusion of state in state-space models. arXiv preprint arXiv:2404.08819, 2024. [130] Balázs Mészáros, James C Knight, and Thomas Nowotny. Efficient event-based delay learning in spiking neural networks. Nature Communications, 16(1):10422, 2025. [131] Svea Marie Meyer, Philipp Weidel, Philipp Plank, Leobardo Campos-Macias, Sumit Bam Shreshta, Philipp Stratmann, Jonathan Timcheck, and Mathis Richter. A diagonal structured state space model on loihi 2 for efficient streaming sequence processing. In 2025 Neuro Inspired Computational Elements (NICE), pages 1–9. IEEE, 2025. [132] Mayank Mishra, Shawn Tan, Ion Stoica, Joseph Gonzalez, and Tri Dao. M2 RNN: Non-linear RNNs with matrix-valued states for scalable language modeling. arXiv preprint arXiv:2603.14360, 2026. [133] Todd Morrill, Christian Pehle, and Anthony Zador. Bullet trains: Parallelizing training of temporally precise spiking neural networks. arXiv preprint arXiv:2603.13283, 2026. 74
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
[134] Nicola Muca Cirone, Antonio Orvieto, Benjamin Walker, Cristopher Salvi, and Terry Lyons. Theoretical foundations of deep selective state-space models. Advances in Neural Information Processing Systems, 37: 127226–127272, 2024. [135] Farzan Nadim and Dirk Bucher. Neuromodulation of neurons and synapses. Current opinion in neurobiology, 29:48–56, 2014. [136] Emre O Neftci, Hesham Mostafa, and Friedemann Zenke. Surrogate gradient learning in spiking neural networks. IEEE Signal Processing Magazine, 36(6):51–63, 2019. [137] Sam V Norman-Haignere, Menoua Keshishian, Orrin Devinsky, Werner Doyle, Guy M McKhann, Catherine A Schevon, Adeen Flinker, and Nima Mesgarani. Temporal integration in human auditory cortex is predominantly yoked to absolute time. Nature Neuroscience, 28(11):2356–2365, 2025. [138] Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, et al. In-context learning and induction heads. arXiv preprint arXiv:2209.11895, 2022. [139] Garrick Orchard, E Paxon Frady, Daniel Ben Dayan Rubin, Sophia Sanborn, Sumit Bam Shrestha, Friedrich T Sommer, and Mike Davies. Efficient neuromorphic signal processing with loihi 2. In 2021 IEEE workshop on signal processing systems (SiPS), pages 254–259. IEEE, 2021. [140] Antonio Orvieto, Samuel L Smith, Albert Gu, Anushan Fernando, Caglar Gulcehre, Razvan Pascanu, and Soham De. Resurrecting recurrent neural networks for long sequences. In International conference on machine learning, pages 26670–26698. PMLR, 2023. [141] Mitchell Ostrow, Adam Eisen, Leo Kozachkov, and Ila Fiete. Beyond geometry: Comparing the temporal structure of computation in neural circuits with dynamical similarity analysis. Advances in Neural Information Processing Systems, 36:33824–33837, 2023. [142] Razvan Pascanu, Tomas Mikolov, and Yoshua Bengio. On the difficulty of training recurrent neural networks. In International conference on machine learning, pages 1310–1318. Pmlr, 2013. [143] Jens E Pedersen, Steven Abreu, Matthias Jobst, Gregor Lenz, Vittorio Fra, Felix Christian Bauer, Dylan Richard Muir, Peng Zhou, Bernhard Vogginger, Kade Heckel, et al. Neuromorphic intermediate representation: A unified instruction set for interoperable brain-inspired computing. Nature Communications, 15(1):8122, 2024. [144] Bo Peng, Eric Alcaide, Quentin Anthony, Alon Albalak, Samuel Arcadinho, Stella Biderman, Huanqi Cao, Xin Cheng, Michael Chung, Leon Derczynski, et al. Rwkv: Reinventing rnns for the transformer era. In Findings of the association for computational linguistics: EMNLP 2023, pages 14048–14077, 2023. [145] Michael Poli, Stefano Massaroli, Eric Nguyen, Daniel Y Fu, Tri Dao, Stephen Baccus, Yoshua Bengio, Stefano Ermon, and Christopher Ré. Hyena hierarchy: Towards larger convolutional language models. In International Conference on Machine Learning, pages 28043–28078. PMLR, 2023. [146] Michael Poli, Armin W Thomas, Eric Nguyen, Pragaash Ponnusamy, Björn Deiseroth, Kristian Kersting, Taiji Suzuki, Brian Hie, Stefano Ermon, Christopher Ré, et al. Mechanistic design and scaling of hybrid architectures. arXiv preprint arXiv:2403.17844, 2024. [147] Alexandre Queant, Ulysse Rançon, Benoit R Cottereau, and Timothée Masquelier. Delrec: Learning delays in recurrent spiking neural networks. arXiv preprint arXiv:2509.24852, 2025. [148] Hubert Ramsauer, Bernhard Schäfl, Johannes Lehner, Philipp Seidl, Michael Widrich, Thomas Adler, Lukas Gruber, Markus Holzleitner, Milena Pavlović, Geir Kjetil Sandve, et al. Hopfield networks is all you need. arXiv preprint arXiv:2008.02217, 2020. [149] Oliver Rhodes, Luca Peres, Andrew GD Rowley, Andrew Gait, Luis A Plana, Christian Brenninkmeijer, and Steve B Furber. Real-time cortical simulation on neuromorphic hardware. Philosophical Transactions of the Royal Society A, 378(2164):20190160, 2020. [150] David W Romero, Anna Kuzina, Erik J Bekkers, Jakub M Tomczak, and Mark Hoogendoorn. Ckconv: Continuous kernel convolution for sequential data. arXiv preprint arXiv:2102.02611, 2021. [151] Magdalena Sabat, Hortense Gouyette, Quentin Gaucher, Mateo López Espejo, Stephen V David, Sam NormanHaignere, and Yves Boubenec. Neurons in auditory cortex integrate information within a constrained and context-invariant temporal window. Current Biology, 35(24):6114–6125, 2025. [152] Darjan Salaj, Anand Subramoney, Ceca Kraisnikovic, Guillaume Bellec, Robert Legenstein, and Wolfgang Maass. Spike frequency adaptation supports network computations on temporally dispersed information. Elife, 10:e65459, 2021. 75
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
[153] Imanol Schlag, Kazuki Irie, and Jürgen Schmidhuber. Linear transformers are secretly fast weight programmers. In International conference on machine learning, pages 9355–9366. PMLR, 2021. [154] Mark Schöne, Neeraj Mohan Sushma, Jingyue Zhuge, Christian Mayr, Anand Subramoney, and David Kappel. Scalable event-by-event processing of neuromorphic sensory signals with deep state-space models. In 2024 International Conference on Neuromorphic Systems (ICONS), pages 124–131. IEEE, 2024. [155] Mark Schöne, Babak Rahmani, Heiner Kremer, Fabian Falck, Hitesh Ballani, and Jannes Gladrow. Implicit language models are rnns: balancing parallelization and expressivity. arXiv preprint arXiv:2502.07827, 2025. [156] Simon Schug, Seijin Kobayashi, Yassir Akram, João Sacramento, and Razvan Pascanu. Attention as a hypernetwork. In International Conference on Learning Representations, volume 2025, pages 68744–68770, 2025. [157] Claude Elwood Shannon. A mathematical theory of communication. The Bell system technical journal, 27(3): 379–423, 1948. [158] Shuaijie Shen, Chao Wang, Renzhuo Huang, Yan Zhong, Qinghai Guo, Zhichao Lu, Jianguo Zhang, and Luziwei Leng. Spikingssms: Learning long sequences with sparse and parallel spiking state space models. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 39, pages 20380–20388, 2025. [159] Yixin Shen. Kron-lora: Hybrid kronecker-lora adapters for scalable, sustainable fine-tuning. arXiv preprint arXiv:2508.01961, 2025. [160] Julien Siems, Timur Carstensen, Arber Zela, Frank Hutter, Massimiliano Pontil, and Riccardo Grazzi. Deltaproduct: Improving state-tracking in linear rnns via householder products. Advances in Neural Information Processing Systems, 38:153738–153782, 2026. [161] Wolf Singer and Andreea Lazar. Does the cerebral cortex exploit high-dimensional, non-linear dynamics for information processing? Frontiers in computational neuroscience, 10:99, 2016. [162] Aditi Singh, Nirmal Prakashbhai Patel, Abul Ehtesham, Saket Kumar, and Tala Talaei Khoei. A survey of sustainability in large language models: Applications, economics, and challenges. In 2025 IEEE 15th Annual Computing and Communication Workshop and Conference (CCWC), pages 00008–00014. IEEE, 2025. [163] Jimmy Smith, Scott Linderman, and David Sussillo. Reverse engineering recurrent neural networks with jacobian switching linear dynamical systems. Advances in Neural Information Processing Systems, 34:16700–16713, 2021. [164] Jimmy TH Smith, Andrew Warrington, and Scott W Linderman. Simplified state space layers for sequence modeling. International Conference on Learning Representations (ICLR), 2023. [165] Matei-Ioan Stan and Oliver Rhodes. Learning long sequences in spiking neural networks. Scientific Reports, 14 (1):21957, 2024. [166] Jianlin Su, Yu Lu, Shengfeng Pan, Ahmed Murtadha, Bo Wen, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. arXiv preprint arXiv:2104.09864, 2021. [167] Pengfei Sun, Jibin Wu, Paul Devos, and Dick Botteldooren. Towards parameter-free attentional spiking neural networks. Neural Networks, 185:107154, 2025. [168] David Sussillo and Omri Barak. Opening the black box: low-dimensional dynamics in high-dimensional recurrent neural networks. Neural computation, 25(3):626–649, 2013. [169] Yi Tay, Mostafa Dehghani, Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang, Sebastian Ruder, and Donald Metzler. Long range arena: A benchmark for efficient transformers. arXiv preprint arXiv:2011.04006, 2020. [170] Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. Efficient transformers: A survey. ACM Computing Surveys, 55(6):1–28, 2022. [171] Kimi Team, Yu Zhang, Zongyu Lin, Xingcheng Yao, Jiaxi Hu, Fanqing Meng, Chengyin Liu, Xin Men, Songlin Yang, Zhiyuan Li, et al. Kimi linear: An expressive, efficient attention architecture. arXiv preprint arXiv:2510.26692, 2025. [172] Aleksandar Terzic, Nicolas Menet, Michael Hersche, Thomas Hofmann, and Abbas Rahimi. Structured sparse transition matrices to enable state tracking in state-space models. Advances in Neural Information Processing Systems, 38:83072–83111, 2026. [173] Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13971, 2023. 76
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
[174] Trieu Trinh, Andrew Dai, Thang Luong, and Quoc Le. Learning longer-term dependencies in rnns with auxiliary losses. In International conference on machine learning, pages 4965–4974. PMLR, 2018. [175] Jonathan H Tu. Dynamic mode decomposition: Theory and applications. Princeton University, 2013. [176] Unconventional AI. Introducing Un-0: Generating images with coupled oscillators. Blog post, June 2026. URL https://unconv.ai/blog/introducing-un-0-generating-images-with-coupled-oscillators/. [177] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017. [178] D Vignesh, Shaobo He, and Santo Banerjee. A review on the complexities of brain activity: insights from nonlinear dynamics in neuroscience. Nonlinear Dyn, 113(5):4531–4552, 2025. [179] Aaron Voelker, Ivana Kajić, and Chris Eliasmith. Legendre memory units: Continuous-time representation in recurrent neural networks. Advances in neural information processing systems, 32, 2019. [180] Johannes Von Oswald, Christian Henning, Benjamin F Grewe, and João Sacramento. Continual learning with hypernetworks. arXiv preprint arXiv:1906.00695, 2019. [181] Johannes von Oswald, Nino Scherrer, Seijin Kobayashi, Luca Versari, Songlin Yang, Maximilian Schlegel, Kaitlin Maile, Yanick Schimpf, Oliver Sieberling, Alexander Meulemans, et al. Mesanet: Sequence modeling by locally optimal test-time training. arXiv preprint arXiv:2506.05233, 2025. [182] Jiaqi Wang, Liutao Yu, Xiongri Shen, Sihang Guo, Chenlin Zhou, Leilei Zhao, Yi Zhong, Zhiguo Zhang, and Zhengyu Ma. Spikcommander: A high-performance spiking transformer with multi-view learning for efficient speech command recognition. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 40, pages 2119–2127, 2026. [183] Samuel S-H Wang, Jennifer R Shultz, Mark J Burish, Kimberly H Harrison, Patrick R Hof, Lex C Towns, Matthew W Wagers, and Krysta D Wyatt. Functional trade-offs in white matter axonal scaling. Journal of neuroscience, 28(15):4047–4056, 2008. [184] Pete Warden. Speech commands: A dataset for limited-vocabulary speech recognition. arXiv preprint arXiv:1804.03209, 2018. [185] Carlo Wenig, Raoul-Martin Memmesheimer, and Christian Klos. Quadratic integrate-and-fire neurons exhibit less fragmented loss landscapes and outperform leaky integrate-and-fire neurons in spike-based gradient descent. arXiv preprint arXiv:2606.03935, 2026. [186] Paul J Werbos. Backpropagation through time: what it does and how to do it. Proceedings of the IEEE, 78(10): 1550–1560, 1990. [187] Bernard Widrow and Marcian E Hoff. Adaptive switching circuits. In Neurocomputing: foundations of research, pages 123–134. 1988. [188] Scott Wisdom, Thomas Powers, John Hershey, Jonathan Le Roux, and Les Atlas. Full-capacity unitary recurrent neural networks. Advances in neural information processing systems, 29, 2016. [189] Max A Woodbury. Inverting modified matrices. Department of Statistics, Princeton University, 1950. [190] Timo C Wunderlich and Christian Pehle. Event-based backpropagation can compute exact gradients for spiking neural networks. Scientific Reports, 11(1):12829, 2021. [191] Shang Xu, Jiayu Zhang, Ziming Wang, Runhao Jiang, Rui Yan, and Huajin Tang. Asrc-snn: Adaptive skip recurrent connection spiking neural network. In ICASSP 2026-2026 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pages 4786–4790. IEEE, 2026. [192] Songlin Yang, Bailin Wang, Yikang Shen, Rameswar Panda, and Yoon Kim. Gated linear attention transformers with hardware-efficient training. arXiv preprint arXiv:2312.06635, 2023. [193] Songlin Yang, Bailin Wang, Yu Zhang, Yikang Shen, and Yoon Kim. Parallelizing linear transformers with the delta rule over sequence length. Advances in neural information processing systems, 37:115491–115522, 2024. [194] Songlin Yang, Jan Kautz, and Ali Hatamizadeh. Gated delta networks: Improving mamba2 with delta rule. In International Conference on Learning Representations, volume 2025, pages 29687–29707, 2025. [195] Man Yao, Jiakui Hu, Zhaokun Zhou, Li Yuan, Yonghong Tian, Bo Xu, and Guoqi Li. Spike-driven transformer. Advances in neural information processing systems, 36:64043–64058, 2023. [196] Bojian Yin, Federico Corradi, and Sander M Bohté. Accurate and efficient time-domain classification with adaptive spiking recurrent neural networks. Nature Machine Intelligence, 3(10):905–913, 2021. 77
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
[197] Annan Yu and N Benjamin Erichson. Block-biased mamba for long-range sequence processing. Advances in Neural Information Processing Systems, 38:120541–120577, 2026. [198] Anthony Zador, Ari S Benjamin, Kyle Daruwalla, Christian Pehle, and Abdul-Malik Zekri. Walking the weight manifold: a topological approach to conditioning inspired by neuromodulation. arXiv, 2025. [199] Shuangfei Zhai, Walter Talbott, Nitish Srivastava, Chen Huang, Hanlin Goh, Ruixiang Zhang, and Josh Susskind. An attention free transformer. arXiv preprint arXiv:2105.14103, 2021. [200] Zhaokun Zhou, Yuesheng Zhu, Chao He, Yaowei Wang, Shuicheng Yan, Yonghong Tian, and Li Yuan. Spikformer: When spiking neural network meets transformer. arXiv preprint arXiv:2209.15425, 2022. [201] Rui-Jie Zhu, Yu Zhang, Steven Abreu, Ethan Sifferman, Tyler Sheaves, Yiqiao Wang, Dustin Richmond, Sumit Bam Shrestha, Peng Zhou, and Jason K Eshraghian. Scalable matmul-free language modeling. arXiv preprint arXiv:2406.02528, 2024. [202] Rui-Jie Zhu, Zixuan Wang, Kai Hua, Tianyu Zhang, Ziniu Li, Haoran Que, Boyi Wei, Zixin Wen, Fan Yin, He Xing, et al. Scaling latent reasoning via looped language models. arXiv preprint arXiv:2510.25741, 2025.
Appendices A
Numerically Stable Lyapunov Exponent Algorithm
Computing the Lyapunov spectrum from the long-term Jacobian of a dynamical system may encounter numerical stability issues due to vanishing or exploding values in repeated matrix products. This study uses Algorithm 2 from Engelken et al. [50] and Engelken [47] to compute exponents accurately. Algorithm 2 Numerically Stable Lyapunov Spectrum Computation Require: T > 0, B, inputs∈ RBatch Size×D×T , initial state ∈ RBatch Size×D , states∈ RBatch Size×D×T , niters Ensure: T %B = 0 Nblocks ← T /B input blocks ← inputs.chunk(Nblocks , dim = −1) states blocks ← states.chunk(Nblocks , dim = −1) states out ← empty list i←0 while i < nblocks do j←0 states guess ← state blocks[i] while j < niters do states guess ← DEER(states guess, initial guess, input blocks[i]) j ←j+1 end while i←i+1 initial state = states guess[..., -1] states out.append(states guess) end while return cat(states out, dim=-1)
B
Code for Forward DEER MAC Estimate class ForwardDEERMACs(nn.Module): def __init__(self): super().__init__() def forward(self, U, c, V, Lambda): L, R = U @ c, V 78
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
LR_diag = L[..., 0] * R[..., 0] + L[..., 1] * R[..., 1] RLL = R.matmul((L.mT * Lambda).matmul(L))# RLLR_diag = RLL[..., 0] * R[..., 0] + RLL[..., 1] * R[..., 1] diag_approx = Lambda + 2 * Lambda * LR_diag + RLLR_diag return diag_approx
C
Training Setup for State Tracking
L = 8 sweep s5_L8_*_beta_sparsity
L = 32 sweep s5_L32_ln_relu_cdot
Task and data Group Training sequences Test sequences Training length Evaluation lengths
S5 , 120 elements, non-solvable; identity monoid, p = 0 8,388,608 (512 × 16384), uint8 4,096 per evaluation length 8 32 8, 16, 32 (1/2/4×) 32, 64, 128 (1/2/4×)
Model Layers Width Heads Normalisation Dropout Output projection Positional encoding Input gate Recurrence solver
1 512 1 batch norm 0.05 GLU none (RoPE off) on, per element; gate trainable Sequential Simulation
Optimisation Optimiser Learning rate Weight decay β1 , β 2 ϵ Gradient clipping Schedule Batch size Epochs Seeds Model selection Evaluation interval
AdamW (adamw_torch_fused) 2 × 10−3 0 0.9, 0.95 10−8 0.25 linear, 10% warmup (1,638 steps) 512 1 (16,384 steps) 0, 1, 2 best eval_loss, loaded at end every 1,000 steps, all lengths
Swept Spectrum init β Sparsity Negative half k/v LayerNorm relu_past_out Controlled dot Early stopping Configurations
s4d-real 0, 0.00125, 0.0125, 0.125, 0.5, 1, 4 0%, 75% yes / no off off off none 7 × 2 × 2 = 28
Table 16
79
continuous_pm (fixed) 4 (fixed) 75% (fixed) yes (fixed) swept swept swept, dot_value = 0.5 stop when L=64 is perfect see results table
ADPTNet: ADaptive with Prescriptive Timescales Non-Linear SSM for Sequence Modelling
D
Training Setup for Selective Copy Baseline
ADPTNet
Vocabulary Tokens to copy Training samples Test samples Layers Eig. Init. Batch size Epochs
16 16 10,000 per epoch, regenerated 1,000 2 S4D-Real 64 200
LR WD No. Seeds
10−3 , 5 × 10−3 , 10−2 0, 0.1 3
Table 17
dmodel nstate Layers Eig. Init. β
64 4 2 S4D-Inv 0.0125
Epochs LR WD
100 5 × 10−4 0
Table 18
E
MAD Benchmark Setup
F
Sequential CIFAR-10 Setup
G
Spiking Speech Commands Setup
80
10−3 0.1 1