GSS: Gated State Space Model.

Implements the Gated State Space model from "Long Range Language Modeling via Gated State Spaces" (Mehta et al., 2023). GSS simplifies S4 by using fixed (learned but not input-dependent) A, B, C matrices combined with multiplicative gating for non-linearity.

Key Innovation: Fixed SSM + Multiplicative Gating

Unlike Mamba (where B, C, dt are input-dependent), GSS uses:

• Fixed diagonal A, B, C matrices (learned via Axon.param)
• Gating for input-dependent non-linearity: gate = sigmoid(W_g * x)
• Result: simpler than Mamba, more expressive than vanilla S4

Equations

# SSM with fixed parameters:
h_t = A * h_{t-1} + B * x_t      # A, B are learned parameters (not input-dependent)
y_t = C * h_t                      # C is a learned parameter

# Gating for non-linearity:
gate_t = sigmoid(W_g * x_t + b_g)
output_t = gate_t * y_t

Architecture

Input [batch, seq_len, embed_dim]
      |
      v
+-------------------------------------+
|         GSS Block                    |
|  LayerNorm -> [SSM path, Gate path]  |
|    SSM: linear -> scan(A,B) -> C*h   |
|    Gate: linear -> sigmoid           |
|  output = SSM * Gate                 |
|  -> project -> residual              |
|  LayerNorm -> FFN -> residual        |
+-------------------------------------+
      | (repeat for num_layers)
      v
Output [batch, hidden_size]

Compared to Other SSMs

Usage

model = GSS.build(
  embed_dim: 287,
  hidden_size: 256,
  state_size: 16,
  num_layers: 4
)

References

• Mehta et al., "Long Range Language Modeling via Gated State Spaces" (2023)
• https://arxiv.org/abs/2206.13947