Optimized Matrix mLSTM Block

• The optimized mLSTM block is a recurrent unit that generalizes classic LSTM using matrix-valued memory states and exponential/sigmoid gating to enable efficient long-range sequence processing.
• It incorporates full sequence-parallelism and tiled flash linear attention to significantly reduce memory usage and boost kernel throughput on modern GPUs.
• Empirical results show mLSTM variants achieving competitive language modeling performance while enabling scalable residual architectures and multi-head integration for diverse pattern recognition.

An optimized mLSTM (matrix Long Short-Term Memory) block is a modern recurrent sequence modeling component that generalizes classic LSTM by introducing matrix-valued memory states, exponential or sigmoid gating, full sequence-parallelism, and highly efficient GPU implementations. These innovations enable mLSTM to compete with and in some cases outperform self-attention and state space models on large-scale language modeling and long-range sequence tasks (Beck et al., 2024, Beck et al., 18 Mar 2025, Lawan et al., 1 Jul 2025).

1. Matrix LSTM Memory and Gate Structure

The core of the optimized mLSTM block is its matrix-valued memory, with cell state $M_t \in \mathbb{R}^{d \times d}$ (for a single head) updated via an outer-product "covariance" rule. Gating functions—input $i_t$, forget $f_t$, output $o_t$—control writing, decay, and exposure of memory, and may utilize exponential or sigmoid activations.

At each timestep $t$ with input $x_t$:

• Compute query $q_t$, key $k_t$, value $v_t$:

$$\begin{aligned} q_t &= W_q x_t + b_q \ \tilde{k}_t &= W_k x_t + b_k,\quad k_t = \tilde{k}_t/\sqrt{d} \ v_t &= W_v x_t + b_v \end{aligned}$$

• Gate pre-activations:

$i_t$0

• Gates:
  • Exponential variant: $i_t$1
  • Sigmoid variant: $i_t$2
  • $i_t$3

The matrix memory and a normalizer vector $i_t$4 update as:

$i_t$5

The output is retrieved via

$i_t$6

2. Exponential Gating and Stabilization

Standard LSTMs utilize bounded $i_t$7 sigmoid gating. The mLSTM generalizes this:

• Exponential input/forget gates allow unbounded magnitudes:
  • Amplifies memory updates or decays memory arbitrarily fast.
  • Risks numerical overflow.

To maintain stability:

• Stabilizer state $i_t$8
• Stabilized gates: $i_t$9, $f_t$0
• These stabilized gates substitute all $f_t$1 in the forward pass; the output and gradients remain unchanged (Beck et al., 2024).

With sigmoid-gated mLSTM ("mLSTMsig" (Beck et al., 18 Mar 2025)), the input gate is simply $f_t$2, making overflow impossible and removing the need for normalizer $f_t$3 or $f_t$4, reducing kernel complexity.

3. Full Parallelism and Tiled Flash Linear Attention (TFLA)

Unlike standard RNNs, optimized mLSTM blocks lack hidden-to-hidden matrix recurrency, making the entire sequence-parallel:

• All time steps along a sequence can be processed simultaneously.
• Construct gate-matrix $f_t$5 with $f_t$6 (forget decay) and $f_t$7 (input write) encoding the sequential dependencies via element-wise products.

TFLA (Beck et al., 18 Mar 2025) implements two-level tiling:

• Level 1: Chunk input sequence into large blocks for memory efficiency.
• Level 2: Further tile matrix multiplications within each chunk to align with on-chip SRAM/tensor-core optimality on GPUs.

This achieves both reduced HBM footprint and maximized arithmetic intensity, leading to empirical performance:

• $f_t$8 speedup over previous approaches (e.g., Flash Attention, Mamba) on long-sequence benchmarks.
• HBM usage for mLSTM is consistently lower than attention-based models at large context lengths.

4. Block Architecture and Residual Integration

Optimized mLSTM is embedded within deep residual stacks ("xLSTM"). Each xLSTM block comprises:

1. LayerNorm
2. Up-projection (expanding dimension)
3. 1D causal convolution (with ReLU$f_t$9 on gates)
4. Learnable skip-connection from input
5. mLSTM sequence mixing (parallel, as above)
6. GroupNorm (head-wise LayerNorm)
7. Residual addition
8. Output gate application
9. Down-projection (restoring dimension) 10. Final residual add to block input

xLSTM architectures stack mostly mLSTM blocks with periodic scalar sLSTM blocks to enrich nonlinearity and mixing. Each block’s output is pre-normalized before entering the next (Beck et al., 2024).

5. Variants and Applicability

Optimized mLSTM can be specialized:

• Exponential vs. sigmoid gating: Exponential gates increase expressivity and update flexibility but require stabilization; sigmoid gates offer safety and speed at marginal expressivity cost (Beck et al., 18 Mar 2025).
• Multi-head architectures: Multiple mLSTM "heads" each with independent memory matrices, paralleling multihead attention, improve pattern diversity and performance at scale (Beck et al., 2024).
• Bi-directional and cross-fusion variants: In tasks requiring nuanced dependencies (e.g., aspect-based sentiment analysis), architectures such as MEGA (Lawan et al., 1 Jul 2025) use forward mLSTM, PF-mLSTM (partially-flipped for local context), and multihead exponential-gated fusion (MECGAF) to efficiently combine global and local context within $o_t$0 cost.

6. Training, Hyperparameters, and Efficiency Benchmarks

Canonical hyperparameters for language modeling:

• Head dimension $o_t$1; typically $o_t$2 independent mLSTM heads
• $o_t$3–$o_t$4 residual blocks; context length $o_t$5
• Optimizer: AdamW ($o_t$6), weight decay $o_t$7, batch size $o_t$8–$o_t$9
• Training uses mixed precision (bfloat16), large batch sizes, and learning rate warm-up plus cosine decay schedules (Beck et al., 2024).
• TFLA-optimized mLSTMsig kernels deliver $t$0 forward+backward throughput of Mamba and outperform FlashAttention on sequence lengths up to $t$1, with $t$2–$t$3 GB HBM compared to $t$4–$t$5 GB for attention (Beck et al., 18 Mar 2025).

Variant	Gate Type	Stabilization Needed	HBM Usage	Kernel Speed (8192 seq)
mLSTMexp	exponential	Yes	9–11 GB	18.2 ms
mLSTMsig	sigmoid	No	9–11 GB	18.2 ms ($t$630% faster forward)
FlashAttn 3	softmax	N/A	12–14 GB	58.7 ms
Mamba 2	N/A	N/A	16 GB	41.2 ms

Values from (Beck et al., 18 Mar 2025), NVIDIA H100, bfloat16, 65K token context.

7. Empirical Context and Research Significance

Optimized mLSTM blocks form the backbone of xLSTM architectures, enabling LLMs at billion-parameter scale to perform competitively with state-of-the-art Transformer and State Space models. The ability to train these models efficiently on modern accelerators—achieving full sequence-parallelism and high arithmetic intensity—is directly attributable to the innovations in gating, matrix memory, and kernel-level optimization. Bi-directional variants and exponential-gated cross-fusions extend the flexibility for complex structured sequence tasks (Beck et al., 2024, Lawan et al., 1 Jul 2025).

A plausible implication is that as attention models face scaling and memory bottlenecks in the long-context regime, optimized mLSTM (with TFLA or related techniques) becomes increasingly attractive for tasks where $t$7 memory/bandwidth and gradient flow are critical.

Empirical results on language modeling (Beck et al., 2024, Beck et al., 18 Mar 2025) and aspect-level sentiment (Lawan et al., 1 Jul 2025) demonstrate mLSTM’s utility across tasks requiring both long-range dependence and efficient high-throughput computation.