H3: Hungry Hungry Hippos.

Implements the H3 architecture from "Hungry Hungry Hippos: Towards Language Modeling with State Space Models" (Fu et al., ICLR 2023). H3 combines two SSM layers with a short convolution and multiplicative gating to close the gap between SSMs and Transformers on language modeling.

Key Innovation: Two-SSM + Short Conv

H3 interleaves two types of SSMs with multiplicative interaction:

Branch 1 (Shift SSM): Captures local dependencies via diagonal SSM
Branch 2 (Diag SSM):  Captures broader patterns via diagonal SSM
Short Conv:           Models very local (1-4 token) patterns

Output = ShortConv(ShiftSSM(x) * DiagSSM(x))

Architecture

Input [batch, seq_len, embed_dim]
      |
      v
+-----------------------+
| Input Projection      |
+-----------------------+
      |
      v
+-----------------------+
| H3 Block x N          |
|  +-- ShiftSSM(x) --+  |
|  |                  |  |
|  +-- DiagSSM(x) ---+  |
|  |                  |  |
|  +--- multiply ----+  |
|  |                     |
|  ShortConv + OutProj   |
|  Residual + FFN        |
+-----------------------+
      |
      v
[batch, hidden_size]    (last timestep)

Usage

model = H3.build(
  embed_dim: 287,
  hidden_size: 256,
  state_size: 64,
  conv_size: 4,
  num_layers: 4
)

Reference

• Paper: "Hungry Hungry Hippos: Towards Language Modeling with State Space Models"
• arXiv: https://arxiv.org/abs/2212.14052