Sessa: Selective State Space Attention — A Deep Dive into the New Architecture
Background and Motivation
Modern sequence modeling is dominated by two architecture families: Transformers and structured state space models (SSMs). Transformers can access any element in context via self-attention, but when attention becomes diffuse—spreading uniformly over many past tokens—a single token's influence is diluted, typically decaying as O(1/ℓ) with distance. Structured SSMs propagate information through an explicit recurrent state, but even selective SSMs like Mamba suffer long-range sensitivity that decays exponentially with lag when they cannot maintain "frozen time."
Both mechanisms therefore struggle at long context: Transformers lose distant information through attention dilution, while SSMs lose long-range dependencies through state forgetting. Existing architectures either read the past through a single attention read or propagate it through a single feedback chain—lacking the ability to aggregate multiple paths within a single layer.
Core Innovation: Attention in the Feedback Path
Sessa's central innovation is placing attention inside the feedback path. It injects input-dependent attention routing into the recurrent feedback loop, combining direct reads of the input with aggregation through the feedback state.
This enables multi-path aggregation within a single layer: past tokens can influence future states through multiple attention paths, rather than a single attention read or a single recurrent chain. Concretely, Sessa's decoder uses a single gated MLP-style module per layer that combines a causal forward attention signal with a feedback term. The feedback term mixes past outputs through a lower-triangular linear system (a causal convolution), enabling multi-hop routing. This "feedback solve" mechanism is the key to Sessa's long-range memory properties.
Theoretical Analysis: Long-Range Memory and Selective Retrieval
Power-Law Memory Tail
Under given assumptions and matching conditions, Sessa achieves a power-law memory tail: influence decays as O(ℓ^-β) with lag ℓ, where 0 < β < 1. This is substantially slower than the Transformer's O(1/ℓ) or Mamba's exponential decay when frozen time fails.
Mechanistically, Sessa's feedback forms a strictly lower-triangular matrix B_fb; the induced transmission kernel K = (I - B_fb)^-1 is lower-triangular with unit diagonal. With feedback routing coefficients constrained to an appropriate range, |K_{t,k}| is bounded by C_K · (t-k)^(-β_tail) for targets far from the source, with 0 < β_tail < 1. In the extreme cases of uniform routing and constant positive feedback, this power-law decay is shown to be tight: influence decays exactly as Θ(ℓ^-β_tail), so the tail cannot decay faster.
Flexible Selective Retrieval
Selective retrieval means a model can dynamically choose which past information to attend to, rather than uniformly decaying all history. By accumulating information across multiple routing depths and temporal paths, Sessa can preserve certain distant information without decay (non-decaying configurations) while letting other influences fade as needed. Theoretical analysis shows that only Sessa among the compared model classes achieves both power-law decay and non-decaying configurations—even when attention routing is highly diffuse and lacks any clear "frozen" regime.
Comparison with Transformer and Mamba
- Memory decay rate: Transformer influence decays as
O(1/ℓ); Mamba decays exponentiallyO(e^(-λℓ))when frozen time fails. Sessa achievesO(ℓ^-β)with0 < β < 1—the slowest decay for large ℓ. - Selective retrieval: Transformers dilute all old tokens uniformly; Mamba retains information over short horizons via selective gating but rapidly loses long-range sensitivity outside frozen-time regimes. Sessa accumulates influence over multiple paths even under diffuse attention, selectively preserving specific information.
- Mechanism: Transformers use global self-attention with quadratic cost in sequence length; Mamba processes sequences in linear time but needs extra machinery for long-range retention. Sessa combines global attention (direct input reads) with local feedback (aggregating past information) within a single layer, achieving linear time complexity with memory properties beyond both.
Experimental Evaluation
Experiments validate the theoretical claims. On long-context benchmarks, Sessa achieves the strongest performance under matched architecture and training budgets. On short-context language modeling, it remains competitive with Transformer and Mamba-style baselines—indicating Sessa does not trade short-term performance for long-range memory, but performs well on both.
Architecture Details and Implementation
Sessa is designed as a plug-and-play replacement for standard decoder blocks. Rather than alternating attention and MLP blocks, it wraps a recurrent mixer in a single gated MLP-style block. The mixer combines standard causal forward attention with a feedback term implemented via lower-triangular matrix multiplication for multi-hop routing.
To bound feedback gain and ensure stability, Sessa introduces a hyperparameter γ_max < 1, requiring the feedback gain |γ_t| to remain below γ_max at all times. This guarantees stability of the feedback system and prevents gradient explosion. The implementation is based on PyTorch with optional FlashAcceleration via FlashAttention. Simply replacing existing decoder layers with Sessa layers yields stronger long-range information retention.
Conclusion
By embedding attention in the feedback path, Sessa creatively addresses the long-context limitations of both Transformers and SSMs. It demonstrates the feasibility of multi-path aggregation within a single layer, and shows theoretically and empirically slower long-range decay and more flexible selective retrieval. The architecture opens new directions for information retention and retrieval in sequence models, with potential impact on long-context NLP, time-series forecasting, and beyond.