Skip to content

LinOSS

Introduced in Oscillatory State-Space Models.

The Linear Oscillatory State-Space model (LinOSS) introduces forced linear second-order ODEs (Differential Equation#Ordinary Differential Equation (ODE)), inspired by the oscillatory dynamics of cortical neurons. Compared to the ordinary State Space Model, it achieves more stable dynamics and better models long-range interactions with a simpler and less restrictive parameterization.

Idea

The State Space Model (SSM) is a dynamic system, meaning its recurrent state (continuous memory) is evolving over time. A common issue in these models are vanishing or exploding gradients, where this internal states values explode or vanish in size. The ordinary SSM uses the strict [[HiPPO]] parameterization to initialize the transition matrix, which guarantee bounded growth over time. Finding less strict parameterizations allows for more easily trainable and more expressive models. The LinOSS models the inner states transitions using insights from linear oscillations, which allows for very simple parameterizations of the transition matrix (non-negative diagonal matrix) to guarantee stable inner state dynamics.

Oscillatory State Space

At the core of LinOSS is in how the inner states dynamics are modeled over time. The ordinary SSM uses the first-order ODE \(y^{\prime}(t)=Ay(t)+Bu(t)\) to model the dynamics of the state space \(y(t)\). In contrast, LinOSS uses the second-order ODE to model the dynamics. This is the functional form of a Harmonic Oscillator#Forced Harmonic Oscillator, an approximation of neural oscillation. Neural oscillations represent synchronizations in neural spike trains (Neural Spiking#Oscillations) and are the foundation of many neural functions like information processing, perception, temporal decoding, memory, sleep and more.

Formulation

The model is defined via the second-order ODE

\[ y^{\prime\prime}(t)=-Ay(t)+Bu(t)+b \]

with a linear read-out

\[ x(t)=Cy(t)+Du(t) \]

where - \(y(t)\in \mathbb{R}^{m}\) is \(m\)-dim hidden-state - \(x(t)\in \mathbb{R}^{q}\) is \(q\)-dim output state - \(u(t)\in \mathbb{R}^{p}\) is \(p\)-dim input state - \(A\in \mathbb{R}^{m\times m}\), \(B \in \mathbb{R}^{m\times p}\), \(C \in \mathbb{R}^{q \times m}\), \(D\in \mathbb{R}^{q\times p}\), and \(b \in \mathbb{R}^{m}\)

We see that weight \(B\) projects input \(u(t)\) into state-space dimension \(m\), weight \(A\) is a transition matrix in the state-space dimension, weight \(C\) projects state-space to output space dimension, and weight \(D\) carries over information from input space to output space.

In the single-dimensional case (\(p=m=1\)) with no bias or input dependency (\(B=b=0\)), the model describes simple harmonic motion with frequency \(A\):

\[y^{\prime\prime}=-Ay \Leftrightarrow y^{\prime\prime} + Ay=0\]

Setting \(B\neq 0\) introduces an external force to the system, proportional to the input \(u(t)\) and modulated by weight \(B\). Increasing dimensionality (\(p,m>1\)) represents an uncoupled system of forced harmonic oscillators.

Discretization

To solve the ODE system described above as fast as possible, two discretization schemes are proposed. Both rely on the substitution \(z=y^{\prime}\), which lets us rewrite the above formulation (without bias) as two first-order systems:

\[ \begin{align} z(t)^{\prime}&=-Ay(t)+Bu(t)\\ y(t)^{\prime}&=z(t) \end{align} \]

For discretization, either an implicit or explicit scheme can be chosen: - The implicit scheme uses the slope at the unkown time-step to find the value of the next time-step: \(y_{n+1}=y_{n}+\Delta t \cdot y^{\prime}_{n+1}\). - The explicit scheme uses the slope at the known time-step: \(y_{n+1}=y_{n}+\Delta t \cdot y_{n}^{\prime}(t)\).

Implicit Time Integration

The implicit scheme defines the following discretization:

\[ \begin{align} z_{n} &= z_{n-1}+\Delta t (-Ay_{n}+Bu_{n}),\\ y_{n} &= y_{n-1}+\Delta t \cdot z_{n} \end{align} \]

Note that the introduction of \(z\) allows us to perform two first-order implicit discretizations. We can bring this formulation into a single matrix equation by introducing \(x_{n}=[z_{n},y_{n}]^{\top}\):

\[ Mx_{n}=x_{n-1}+F_{n} \]

where

\[ \begin{align} M &= \begin{bmatrix} I & \Delta t A\\ -\Delta t I & I \end{bmatrix}\\ F &= \begin{bmatrix} \Delta t Bu_{n}\\ 0 \end{bmatrix} \end{align} \]

We can now solve for \(x_{n}\) by multiplying with the inverse of matrix \(M\). Inverting the matrix

Ideas

  • Hidden state
    • Analyse internal dynamics of model (eigenvalue-based frequency decomposition?) => Does model approach Gamma and Theta frequencies?
    • Can initialization of model be optimized based on knowledge about cortical auditory processing? E.g. choose transition matrix which creates oscillations close to Theta/Gamma frequencies
  • BioOSS outperforms LinOSS with hierarchical periodic patterns, might be interesting comparison to LinOSS for SE