[Sequence Modeling 02] LSTM

01Vanilla RNN 的缺陷：为什么我们需要 LSTM？#

在上一篇中我们提到，RNN 的核心公式是。这种设计在处理长序列时会遇到两个工程上的致命问题：

梯度消失与梯度爆炸（Vanishing/Exploding Gradients）
- 在反向传播计算梯度时，误差需要沿着时间步反向传递。这会导致权重矩阵被连乘次。根据线性代数原理，如果的最大特征值小于 1，连乘后梯度会呈指数级衰减趋近于 0（梯度消失）；如果大于 1，则会指数级放大（梯度爆炸）。梯度消失意味着网络根本无法学习到长距离的依赖关系
信息覆盖（Information Overwrite）
- RNN 只有一个隐藏状态。在每一个时间步，新的输入都会强制与历史信息混合。没有任何机制能够保护早期非常重要但最近没有出现的信息。这就好比一个容量有限的栈，新数据不断涌入，旧数据很快就被冲刷掉了

02LSTM 的核心思想：分离状态与引入门控机制#

为了解决上述问题，LSTM 将内部状态拆分为两条相互配合的通路，并通过门控信号控制信息如何在通路之间流动。

细胞状态 (Cell State)：位于上方的主干通路，负责携带长期记忆。从到的更新主要由逐元素乘法和加法构成，因此为梯度跨时间步传播提供了更顺畅的路径。
隐藏状态 (Hidden State)：位于下方的输出通路，既参与当前时间步的门控计算，也作为当前时间步的输出，并传递给下一时间步。
门控机制 (Gating Mechanism)：将上一时刻的隐藏状态与当前输入拼接后，分别计算遗忘门、输入门、输出门和候选记忆。其中三个门使用 Sigmoid，将每个维度的控制值限制在；候选记忆使用生成待写入的内容。

一个时间步内部的数据流可以按图中的顺序理解：遗忘门决定保留多少，输入门决定写入多少候选记忆，二者相加得到新的细胞状态；随后，经过压缩，并由输出门决定其中多少信息暴露为隐藏状态。

03数学工作流：LSTM 的四个核心步骤#

对于第个 token，LSTM 内部执行以下运算。为了方便，我们通常将前一个隐藏状态和当前输入拼接在一起计算。

这里的拼接写法与代码中的两层线性变换是等价的。若将拼接后的向量记为，可以写成一个大的线性变换；代码则把这个矩阵按输入来源拆成两块，分别计算和，再将结果相加。因此与不是同一个矩阵，但共同完成对拼接向量的投影。

第一步：遗忘门 (Forget Gate) —— 决定丢弃什么历史信息

是一个与维度相同的向量。它的值在 0 到 1 之间，决定了长期记忆中哪些维度的数据应该被清空。

第二步：输入门 (Input Gate) —— 决定写入什么新信息 我们需要计算两部分：哪些维度需要更新（），以及新的候选内容是什么（）。

充当开关，是当前时刻提取出的新特征。

第三步：更新细胞状态 (Update Cell State)

这里的表示逐元素乘法，即 Hadamard Product

这是 LSTM 最精妙的一步：旧记忆乘以遗忘门（按比例遗忘），然后加上新特征乘以输入门（按比例写入）。这是一个纯线性的加法过程，梯度可以完美地沿着加号回传。

第四步：输出门 (Output Gate) —— 决定输出什么

根据当前输入和历史决定输出开关，然后将长期记忆通过压缩到之间，乘以输出开关，得到当前时刻的隐藏状态。

04代码与工作流解析#

以下是标准 LSTM 单元的伪代码实现，展示了参数是如何定义以及前向传播是如何完成的。代码没有为四个门各建一个线性层，而是用两次 matmul 一次算出全部 pre-activation，再切成四块：

python

import torch
import torch.nn as nn

class LSTMCell(nn.Module):
    def __init__(self, input_size: int, hidden_size: int):
        super().__init__()
        self.hidden_size = hidden_size
        
        # Input-to-hidden projection: maps x_t to all four gate pre-activations.
        self.W_ih = nn.Linear(input_size, 4 * hidden_size)

        # Hidden-to-hidden projection: maps h_{t-1} to all four gate pre-activations.
        self.W_hh = nn.Linear(hidden_size, 4 * hidden_size)

        # These two projections are equivalent to one linear layer on [h_{t-1}, x_t].

    def forward(self, x_t: torch.Tensor, states: tuple[torch.Tensor, torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]:
        # x_t shape: (batch_size, input_size)
        # states = (h_{t-1}, c_{t-1}), each shape: (batch_size, hidden_size)
        h_prev, c_prev = states
        
        # Step 1: Compute all gate pre-activations simultaneously
        # gates shape: (batch_size, 4 * hidden_size)
        gates = self.W_ih(x_t) + self.W_hh(h_prev)
        
        # Split the fused tensor into 4 chunks for f, i, o, and c_candidate
        # Each chunk shape: (batch_size, hidden_size)
        forget_gate, input_gate, output_gate, cell_candidate = gates.chunk(4, dim=1)
        
        # Step 2: Apply non-linear activations
        f_t = torch.sigmoid(forget_gate)
        i_t = torch.sigmoid(input_gate)
        o_t = torch.sigmoid(output_gate)
        c_tilde = torch.tanh(cell_candidate)
        
        # Step 3: Update Cell State (The core gradient superhighway)
        # Element-wise operations
        c_t = (f_t * c_prev) + (i_t * c_tilde)
        
        # Step 4: Compute new Hidden State
        h_t = o_t * torch.tanh(c_t)
        
        # Return the updated states for the next time step or final output
        return h_t, c_t

class LSTM(nn.Module):
    def __init__(self, input_size: int, hidden_size: int):
        super().__init__()
        self.cell = LSTMCell(input_size, hidden_size)
        
    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
        # x shape: (batch_size, seq_len, input_size)
        batch_size, seq_len, _ = x.size()
        
        # Initialize h_0 and c_0 with zeros
        h_t = torch.zeros(batch_size, self.cell.hidden_size, device=x.device)
        c_t = torch.zeros(batch_size, self.cell.hidden_size, device=x.device)
        
        outputs = []
        
        # Unroll over time (Sequential nature remains)
        for t in range(seq_len):
            x_t = x[:, t, :]
            h_t, c_t = self.cell(x_t, (h_t, c_t))
            outputs.append(h_t)
            
        final_outputs = torch.stack(outputs, dim=1)
        return final_outputs, (h_t, c_t)

05总结#

LSTM 通过引入复杂的门控机制和跨越时间步的残差连接（即的加法结构），极大地缓解了 RNN 的梯度消失问题，使得模型能够捕获更长距离的依赖。但正如代码中 for t in range(seq_len) 所示，它的本质仍然是基于马尔可夫链的时序迭代，这意味着它依旧无法像 Transformer 那样在训练阶段实现高度的并行计算。