mLSTM with Sigmoid Input Gate

• mLSTM with Sigmoid Input Gate is a gated recurrent neural network that leverages a multiplicative intermediate state and sigmoid activation to enable dynamic, context-sensitive memory updates.
• The architecture simplifies recurrence by eliminating auxiliary states, reducing computation and memory overhead while achieving performance speedups of up to 30%.
• Practical evaluations demonstrate that mLSTM variants excel in long-context language modeling, delivering competitive results in benchmarks with efficient compute and memory usage.

The multiplicative Long Short-Term Memory (mLSTM) architecture with sigmoid input gate is a gated recurrent neural network that synthesizes the expressive hidden-state transitions of multiplicative RNNs with the robust memory control mechanisms of LSTMs. Distinguished by its use of input-dependent multiplicative paths and a sigmoid-activated input gate, mLSTM—and its modern variants such as mLSTMsig—demonstrate enhanced modeling capabilities for sequence data, while enabling efficient implementation and kernel optimization for long-context scenarios (Krause et al., 2016, Beck et al., 18 Mar 2025).

1. Mathematical Formulation

The mLSTM family is characterized by the introduction of an intermediate multiplicative state $m_t$ that modulates both the candidate cell update and all gating signals. The classic recurrence equations for a single-layer mLSTM with sigmoid input gate are: $m_t = (W_{mx}\,x_t) \odot (W_{mh}\,h_{t-1})$

$\hat{h}_t = W_{hx}\,x_t + W_{hm}\,m_t$

$i_t = \sigma(W_{ix}\,x_t + W_{im}\,m_t)$

$f_t = \sigma(W_{fx}\,x_t + W_{fm}\,m_t)$

$o_t = \sigma(W_{ox}\,x_t + W_{om}\,m_t)$

$$c_t = f_t \odot c_{t-1} + i_t \odot \tanh(\hat{h}_t)$$

$h_t = o_t \odot \tanh(c_t)$

with $x_t \in \mathbb{R}^N$, $h_t, c_t, m_t \in \mathbb{R}^H$, and $\sigma(\cdot)$ the elementwise sigmoid. All gates are thus functions of both the input and a multiplicative interaction between the input and previous hidden state (Krause et al., 2016).

Modern efficient mLSTM variants such as mLSTMsig reformulate the cell state as a learned matrix memory $C_t$: $$f_t = \sigma(\tilde R^f_t), \quad i_t = \sigma(\tilde R^i_t)$$

$C_t = f_t C_{t-1} + i_t k_t v_t^\top$

$\tilde h_t = C_t^\top q_t$

$$h_t = \sigma(\tilde o_t) \odot \mathrm{NORM}(\tilde h_t)$$

Here, $k_t, q_t, v_t$ are linearly projected from $x_t$, and the plain sigmoid gating removes the need for auxiliary stabilizer states (Beck et al., 18 Mar 2025).

2. Distinguishing Features and Theoretical Implications

Unlike standard LSTMs, which drive gating by additive combinations of input and previous hidden state, mLSTM employs an elementwise product, imparting each input symbol the potential to induce a unique hidden-to-hidden transition. This expands the effective capacity of the model's state space. The dependence of the input, forget, and output gates on $m_t$ introduces strong input-state interactions, which permit context-dependent modulation of memory writing and retention.

Sigmoid input gating constrains the valid update range to $(0,1)$, ensuring numerically stable input modulation and obviating the need for normalization or max-rescaling routines present in earlier exponential-gated architectures.

3. Comparison with Exponential-Gated and Other Linear RNNs

The transition from exponential-gated input (mLSTMexp) to sigmoid-gated input (mLSTMsig) eliminates two auxiliary states: the numerical normalizer and max state. This simplification yields several practical consequences:

• Reduces the forward compute and memory footprint by 20–30% per chunk.
• Simplifies the recurrence, removing the need for rescaling and auxiliary tracking in long-context runs.
• Yields kernel designs that halve the number of global memory barriers and significantly decrease non-tensor-core compute (Beck et al., 18 Mar 2025).

Relative to linear RNNs omitting input gates (e.g., FlashLinearAttention), mLSTM with sigmoid input gate maintains a richer gating structure while retaining competitive computational properties.

4. Practical Performance and Implementation

Empirical results establish mLSTM and its variants as strong sequence modelers. Key benchmarks are:

• On text8, character-level bits per char: $\approx1.27$.
• On Hutter Prize: $\approx1.24$ bits/char (with $\approx$46M parameters).
• On WikiText-2: character-level entropy $1.26$ bits/char (byte-level), word-level perplexity $88.8$—on par with optimized word-level LSTMs (Krause et al., 2016).
• In next-token prediction tasks (DCLM) up to 1.4B parameters, mLSTMsig and mLSTMexp match perplexities to within $0.1$ PPL for all head configs (Beck et al., 18 Mar 2025).

The resource-optimized TFLA kernel for mLSTMsig achieves:

• Inference (65k tokens): $\sim30\%$ faster than mLSTMexp, up to $2\times$ faster than FlashAttention 3 for $T\geq4096$.
• Training: $\sim2\times$ faster than Mamba 2 for all $T$, and $20$–$30\%$ faster than FlashAttention 3 at $T\geq4096$.
• Memory–runtime trade-off: at $T=8192$, batch 8, embedding 4096, optimal chunking yields $25\%$ lower GPU memory use than competing kernels, consistent with roofline/runtime analysis (Beck et al., 18 Mar 2025).

5. Kernel Design and Optimization

The Tiled Flash Linear Attention (TFLA) kernel for mLSTMsig adopts a two-level sequence parallelism:

• Sequence is chunked (Level 1); each chunk's initial state is computed and materialized by a recurrent kernel.
• Within the chunk (Level 2), matrix-multiply operations—$QK^\top$ with $V$ and chunk-initial $C$—are block-tiled for efficient tensor-core utilization and SRAM management.

The absence of auxiliary normalizers/max-states in mLSTMsig enables fusion of intra- and inter-chunk computations. In Triton implementation, a thread block loads and processes blocks of $Q$, $K$, and $V$, applying sigmoid gates, and accumulates both attention-derived and memory contributions in a fully fused forward pass with fewer memory barriers (Beck et al., 18 Mar 2025).

6. Training Stability, Hyperparameters, and Best Practices

Experimental best practices for mLSTM recurrent architectures, including mLSTMsig:

• Hidden dimension $H=1900$ (or up to $H=2800$ for large-scale models).
• Embedding layer of size $400$ preceding first-layer projections.
• Adam optimizer, learning rate decaying from $0.001$ to $\approx0.00005$.
• Scaled orthogonal initialization for recurrent matrices.
• Initial forget-gate bias $+3$ (classic) or input-gate bias $-10$ (for mLSTMsig, to suppress early gradient spikes).
• Truncated BPTT: length $200$–$250$.
• Variational dropout $0.2$–$0.5$ on embeddings and hidden paths; weight normalization on all recurrent matrices (Krause et al., 2016, Beck et al., 18 Mar 2025).

Both classical and sigmoid-gated mLSTM variants exhibit stable optimization and non-degrading accuracy up to large parameter and context scales.

7. Applications and Modeling Power

mLSTM with sigmoid input gate excels in character-level and byte-level language modeling due to the flexibility of input-dependent gating and the capacity to instantiate complex, context-sensitive transition functions. The architecture’s capacity for long-range dependency modeling is maintained through its extended gating structure. The expressiveness introduced by the multiplicative input path, without incurring additional instabilities of exponential gating, renders mLSTMsig a preferred primitive for efficient, scalable sequence modeling, particularly in long-context and high-throughput deployment scenarios (Krause et al., 2016, Beck et al., 18 Mar 2025).