xLSTM v2: Improved Extended Long Short-Term Memory.

Implements improvements from the xLSTM 7B scaling paper, building on the original xLSTM architecture with enhanced matrix memory and normalization.

Key Improvements over xLSTM v1

1. Block-diagonal matrix memory: Reduces mLSTM parameters by partitioning the memory matrix into independent blocks. Each block operates on a subset of dimensions, reducing per-head memory from O(d^2) to O(d^2/B) where B is the number of blocks.

2. Improved normalizer with learnable bias: The normalizer n_t = f_t * n_{t-1} + i_t gains a learnable bias term for better gradient flow: h_t = o_t * (c_t / max(|n_t + bias|, 1))

3. Pre-norm + post-norm hybrid: Combines pre-LayerNorm for stable training with post-LayerNorm for better representation quality.

Architecture

Input [batch, seq_len, embed_dim]
      |
      v
+-------------------------------------+
|     xLSTM v2 Block                   |
|  PreNorm -> mLSTM v2 -> PostNorm     |
|       + Residual                     |
|  PreNorm -> FFN -> PostNorm          |
|       + Residual                     |
+-------------------------------------+
      | (repeat for num_layers)
      v
Output [batch, hidden_size]

Usage

model = XLSTMv2.build(
  embed_dim: 287,
  hidden_size: 256,
  num_layers: 4,
  num_heads: 4,
  num_blocks: 2
)

References

• Beck et al., "xLSTM: Extended Long Short-Term Memory" (NeurIPS 2024)
• xLSTM 7B scaling paper improvements