Context-Selective State Space Models: Feedback is All You Need

TL;DR

Context-Selective State Space Models (COFFEE) introduce state feedback to enable context-based selectivity in time-varying SSMs, addressing long-range dependency modeling with improved efficiency over attention-based methods. By linearizing dynamics, applying state-driven gating, and removing redundancy via a change of basis, COFFEE achieves a lean parameterization with a diagonal Jacobian, enabling scalable parallel training. Empirically, COFFEE significantly outperforms the S6 module on induction-head tasks and MNIST, achieving near-perfect accuracy with far fewer parameters and training samples, and maintaining performance across varied sequences. The results suggest state feedback as a powerful mechanism for scalable, interpretable sequence modeling and motivate integrating COFFEE into larger architectures like Mamba for broader long-range sequence understanding.

Abstract

Transformers, powered by the attention mechanism, are the backbone of most foundation models, yet they suffer from quadratic complexity and difficulties in dealing with long-range dependencies in the input sequence. Recent work has shown that state space models (SSMs) provide an efficient alternative, with the S6 module at the core of the Mamba architecture achieving state-of-the-art results on long-sequence benchmarks. In this paper, we introduce the COFFEE (COntext From FEEdback) model, a novel time-varying SSM that incorporates state feedback to enable context-dependent selectivity, while still allowing for parallel implementation. Whereas the selectivity mechanism of S6 only depends on the current input, COFFEE computes it from the internal state, which serves as a compact representation of the sequence history. This shift allows the model to regulate its dynamics based on accumulated context, improving its ability to capture long-range dependencies. In addition to state feedback, we employ an efficient model parametrization that removes redundancies present in S6 and leads to a more compact and trainable formulation. On the induction head task, COFFEE achieves near-perfect accuracy with two orders of magnitude fewer parameters and training sequences compared to S6. On MNIST, COFFEE largely outperforms S6 within the same architecture, reaching 97% accuracy with only 3585 parameters. These results showcase the role of state feedback as a key mechanism for building scalable and efficient sequence models.

Figures (3)

• Figure 1: Block diagram of the $i$-th feature subsytem of COFFEE. Orange denotes the model signals, blue the learnable parameters, and green the selectivity or gating parameter, which is obtained from the past state through (\ref{['eq:Delta-feedback']}). By feeding this parameter back into the dynamics, the model forms a closed-loop mechanism that provides context-based selectivity, enabling identical tokens to be processed differently depending on their role in the sequence.
• Figure 2: Embeddings (in grey) and trajectories of the state for the reduced model solving the IH0 problem, with input $1\|2\|3\|1$ and $3\|1\|2\|1$ (in blue and red, respectively). The last trigger does not move significantly the state since the sigmoids in the input gates are not activated. The colored regions indicate the sections of the plane corresponding to each possible choice of the final output symbol (e.g final states in Region 2 lead to output "2").
• Figure 3: Proposed architecture.

Theorems & Definitions (7)

• Remark 1: Parameter efficiency
• Proposition 1: Redundancy Elimination by Change of Basis
• Remark 2: Scalable parallel training
• Remark 3
• Definition 1
• Remark 4
• Remark 5