Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

LSTM Gate Equations › The Combined Weight Matrix Formulation

Michael BrenndoerferUpdated May 11, 20254 min read

Language AI HandbookAll chapters

Part 11 of 67Recurrent Neural Networks

Chapter 78LSTM Gate Equations › The Combined Weight Matrix Formulation

Reading section 13 of 17 · The Combined Weight Matrix Formulation · 4 minSection 13 of 17 · 4 min

Section 13 of 17 · 4 min read

LSTM Gate Equations › The Combined Weight Matrix FormulationOne section · Sections in this chapter13 / 17 ↓

Sections in this chapter17

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

In practice, frameworks like PyTorch don't compute each gate separately. Instead, they concatenate all weight matrices and compute all gates in a single matrix multiplication, which is more efficient on GPUs.

The combined formulation stacks the four weight matrices vertically into single large matrices:

\mathbf{W} = \begin{bmatrix} \mathbf{W}_i \\ \mathbf{W}_f \\ \mathbf{W}_c \\ \mathbf{W}_o \end{bmatrix} \in \mathbb{R}^{4h \times d} \quad\quad \mathbf{U} = \begin{bmatrix} \mathbf{U}_i \\ \mathbf{U}_f \\ \mathbf{U}_c \\ \mathbf{U}_o \end{bmatrix} \in \mathbb{R}^{4h \times h} \quad\quad \mathbf{b} = \begin{bmatrix} \mathbf{b}_i \\ \mathbf{b}_f \\ \mathbf{b}_c \\ \mathbf{b}_o \end{bmatrix} \in \mathbb{R}^{4h}

where:

$\mathbf{W} \in \mathbb{R}^{4h \times d}$ : combined input weight matrix, stacking all four gate input weights
$\mathbf{U} \in \mathbb{R}^{4h \times h}$ : combined recurrent weight matrix, stacking all four gate hidden weights
$\mathbf{b} \in \mathbb{R}^{4h}$ : combined bias vector, concatenating all four gate biases
The subscripts $i$ , $f$ , $c$ , $o$ denote input gate, forget gate, candidate, and output gate respectively

Then all gates are computed in one step:

\begin{bmatrix} \mathbf{i}_t \\ \mathbf{f}_t \\ \tilde{\mathbf{c}}_t \\ \mathbf{o}_t \end{bmatrix} = \begin{bmatrix} \sigma \\ \sigma \\ \tanh \\ \sigma \end{bmatrix} \left( \mathbf{W} \mathbf{x}_t + \mathbf{U} \mathbf{h}_{t-1} + \mathbf{b} \right)

The notation $\begin{bmatrix} \sigma \\ \sigma \\ \tanh \\ \sigma \end{bmatrix}$ indicates that different activation functions are applied to different portions of the result: sigmoid to the first $h$ elements (input gate), sigmoid to the next $h$ elements (forget gate), tanh to the next $h$ elements (candidate), and sigmoid to the final $h$ elements (output gate).

This produces a $4h$ -dimensional vector that is split into four chunks of size $h$ each. The efficiency gain comes from performing two large matrix multiplications instead of eight smaller ones.

In[21]:

Code

class EfficientLSTMCell:
    """LSTM cell using combined weight matrices for efficiency."""

    def __init__(self, input_dim, hidden_dim):
        self.input_dim = input_dim
        self.hidden_dim = hidden_dim

        # Combined weight matrices
        scale_x = np.sqrt(2.0 / (input_dim + hidden_dim))
        scale_h = np.sqrt(2.0 / (hidden_dim + hidden_dim))

        self.W = np.random.randn(4 * hidden_dim, input_dim) * scale_x
        self.U = np.random.randn(4 * hidden_dim, hidden_dim) * scale_h
        self.b = np.zeros(4 * hidden_dim)

        # Initialize forget gate bias to 1
        self.b[hidden_dim : 2 * hidden_dim] = 1.0

    def forward(self, x_t, h_prev, c_prev):
        """Forward pass using combined matrices."""
        h = self.hidden_dim

        # Single matrix multiplication for all gates
        gates = self.W @ x_t + self.U @ h_prev + self.b

        # Split and apply activations
        i_t = sigmoid(gates[0:h])
        f_t = sigmoid(gates[h : 2 * h])
        c_tilde = tanh(gates[2 * h : 3 * h])
        o_t = sigmoid(gates[3 * h : 4 * h])

        # Cell state and hidden state updates
        c_t = f_t * c_prev + i_t * c_tilde
        h_t = o_t * tanh(c_t)

        return h_t, c_t

Let's compare the performance of both implementations:

In[22]:

Code

import time

efficient_lstm = EfficientLSTMCell(input_dim=100, hidden_dim=256)
standard_lstm = LSTMCell(input_dim=100, hidden_dim=256)

x = np.random.randn(100)
h = np.zeros(256)
c = np.zeros(256)

# Time standard implementation
start = time.time()
for _ in range(1000):
    standard_lstm.forward(x, h, c)
standard_time = time.time() - start

# Time efficient implementation
start = time.time()
for _ in range(1000):
    efficient_lstm.forward(x, h, c)
efficient_time = time.time() - start

speedup = standard_time / efficient_time

Out[23]:

Console

Performance comparison (1000 forward passes):
  Standard implementation: 93.8 ms
  Efficient implementation: 71.2 ms
  Speedup: 1.32x

The efficient implementation is faster because it performs fewer matrix multiplications. The standard implementation does 8 matrix multiplications (2 per gate), while the efficient version does only 2 (one for inputs, one for hidden states). On GPUs, this difference is even more pronounced because large matrix multiplications are highly optimized.