引言
循环神经网络（RNN）在处理序列数据时面临一个根本性的难题：长程依赖问题。随着序列长度增加，传统RNN的梯度在反向传播过程中会指数级衰减或爆炸，导致模型难以记住久远的历史信息。长短期记忆网络（Long Short-Term Memory, LSTM）由Hochreiter和Schmidhuber于1997年提出，通过精心设计的门控机制解决了这一难题，成为序列建模领域最具影响力的架构之一。本文将用Java从零实现一个完整的LSTM单元，深入讲解遗忘门、输入门、输出门的工作原理。
算法原理
LSTM的核心创新：细胞状态
LSTM引入了一条贯穿时间步的”传送带”——细胞状态（Cell State），信息可以在上面相对 unchanged 地流动。门控机制负责向这条传送带上添加或移除信息。
三个门控结构
遗忘门（Forget Gate）：决定从细胞状态中丢弃什么信息
$$f_t = \sigma(W_f \cdot [h_{t-1}, x_t] + b_f)$$
输入门（Input Gate）：决定什么新信息存入细胞状态
$$i_t = \sigma(W_i \cdot [h_{t-1}, x_t] + b_i)$$
$$\tilde{C}t = \tanh(W_C \cdot [h{t-1}, x_t] + b_C)$$
输出门（Output Gate）：决定基于细胞状态输出什么值
$$o_t = \sigma(W_o \cdot [h_{t-1}, x_t] + b_o)$$
$$h_t = o_t \cdot \tanh(C_t)$$
细胞状态更新
$$C_t = f_t \cdot C_{t-1} + i_t \cdot \tilde{C}_t$$
Java实现
import java.util.*;
/**
* LSTM长短期记忆网络实现
*
* 用于序列预测任务：根据前几个数字预测下一个数字
*/
public class LSTMPredictor {
// LSTM参数
private int inputSize; // 输入维度
private int hiddenSize; // 隐藏层维度
private int outputSize; // 输出维度
// 权重矩阵
private double[][] Wf, Wi, Wc, Wo; // 遗忘门、输入门、候选状态、输出门
private double[] bf, bi, bc, bo; // 偏置
// 输出层权重
private double[][] Wy;
private double[] by;
// 学习率
private double learningRate = 0.01;
// 存储前向传播中间结果用于反向传播
private List<double[]> fList, iList, cTildeList, cList, oList, hList, xList;
public LSTMPredictor(int inputSize, int hiddenSize, int outputSize) {
this.inputSize = inputSize;
this.hiddenSize = hiddenSize;
this.outputSize = outputSize;
// 初始化权重（Xavier初始化）
Wf = initMatrix(hiddenSize, inputSize + hiddenSize);
Wi = initMatrix(hiddenSize, inputSize + hiddenSize);
Wc = initMatrix(hiddenSize, inputSize + hiddenSize);
Wo = initMatrix(hiddenSize, inputSize + hiddenSize);
bf = initVector(hiddenSize, 1.0); // 遗忘门偏置初始为1，鼓励记忆
bi = initVector(hiddenSize, 0.0);
bc = initVector(hiddenSize, 0.0);
bo = initVector(hiddenSize, 0.0);
Wy = initMatrix(outputSize, hiddenSize);
by = initVector(outputSize, 0.0);
}
/**
* Xavier初始化矩阵
*/
private double[][] initMatrix(int rows, int cols) {
Random rand = new Random(42);
double[][] m = new double[rows][cols];
double scale = Math.sqrt(2.0 / (rows + cols));
for (int i = 0; i < rows; i++) {
for (int j = 0; j < cols; j++) {
m[i][j] = rand.nextGaussian() * scale;
}
}
return m;
}
private double[] initVector(int size, double defaultVal) {
double[] v = new double[size];
Arrays.fill(v, defaultVal);
return v;
}
/**
* Sigmoid激活函数
*/
private double sigmoid(double x) {
return 1.0 / (1.0 + Math.exp(-x));
}
/**
* Sigmoid导数
*/
private double sigmoidDerivative(double x) {
double s = sigmoid(x);
return s * (1 - s);
}
/**
* Tanh激活函数
*/
private double tanh(double x) {
return Math.tanh(x);
}
/**
* Tanh导数
*/
private double tanhDerivative(double x) {
double t = tanh(x);
return 1 - t * t;
}
/**
* 矩阵与向量乘法
*/
private double[] matVecMul(double[][] m, double[] v) {
int rows = m.length;
int cols = m[0].length;
double[] result = new double[rows];
for (int i = 0; i < rows; i++) {
double sum = 0;
for (int j = 0; j < cols; j++) {
sum += m[i][j] * v[j];
}
result[i] = sum;
}
return result;
}
/**
* 向量逐元素相加
*/
private double[] vecAdd(double[] a, double[] b) {
double[] result = new double[a.length];
for (int i = 0; i < a.length; i++) {
result[i] = a[i] + b[i];
}
return result;
}
/**
* 向量逐元素相乘（Hadamard积）
*/
private double[] vecMul(double[] a, double[] b) {
double[] result = new double[a.length];
for (int i = 0; i < a.length; i++) {
result[i] = a[i] * b[i];
}
return result;
}
/**
* 向量与标量相乘
*/
private double[] vecScale(double[] v, double scale) {
double[] result = new double[v.length];
for (int i = 0; i < v.length; i++) {
result[i] = v[i] * scale;
}
return result;
}
/**
* 拼接两个向量
*/
private double[] concat(double[] a, double[] b) {
double[] result = new double[a.length + b.length];
System.arraycopy(a, 0, result, 0, a.length);
System.arraycopy(b, 0, result, a.length, b.length);
return result;
}
/**
* 对向量应用sigmoid
*/
private double[] sigmoidVec(double[] v) {
double[] result = new double[v.length];
for (int i = 0; i < v.length; i++) {
result[i] = sigmoid(v[i]);
}
return result;
}
/**
* 对向量应用tanh
*/
private double[] tanhVec(double[] v) {
double[] result = new double[v.length];
for (int i = 0; i < v.length; i++) {
result[i] = tanh(v[i]);
}
return result;
}
/**
* 前向传播：处理一个序列
* @return 每个时间步的预测输出
*/
public double[][] forward(double[][] sequence) {
int seqLen = sequence.length;
fList = new ArrayList<>();
iList = new ArrayList<>();
cTildeList = new ArrayList<>();
cList = new ArrayList<>();
oList = new ArrayList<>();
hList = new ArrayList<>();
xList = new ArrayList<>();
double[] hPrev = new double[hiddenSize];
double[] cPrev = new double[hiddenSize];
double[][] outputs = new double[seqLen][outputSize];
for (int t = 0; t < seqLen; t++) {
double[] x = sequence[t];
xList.add(x);
// 拼接[h_{t-1}, x_t]
double[] hx = concat(hPrev, x);
// 遗忘门
double[] f = sigmoidVec(vecAdd(matVecMul(Wf, hx), bf));
fList.add(f);
// 输入门
double[] i = sigmoidVec(vecAdd(matVecMul(Wi, hx), bi));
iList.add(i);
// 候选细胞状态
double[] cTilde = tanhVec(vecAdd(matVecMul(Wc, hx), bc));
cTildeList.add(cTilde);
// 细胞状态更新
double[] c = vecAdd(vecMul(f, cPrev), vecMul(i, cTilde));
cList.add(c);
// 输出门
double[] o = sigmoidVec(vecAdd(matVecMul(Wo, hx), bo));
oList.add(o);
// 隐藏状态
double[] h = vecMul(o, tanhVec(c));
hList.add(h);
// 输出层
outputs[t] = vecAdd(matVecMul(Wy, h), by);
hPrev = h;
cPrev = c;
}
return outputs;
}
/**
* 均方误差损失
*/
public double computeLoss(double[][] predictions, double[][] targets) {
double loss = 0;
int count = 0;
for (int t = 0; t < predictions.length; t++) {
for (int i = 0; i < predictions[t].length; i++) {
double diff = predictions[t][i] - targets[t][i];
loss += diff * diff;
count++;
}
}
return loss / count;
}
/**
* 训练一个序列
*/
public double train(double[][] sequence, double[][] targets) {
double[][] predictions = forward(sequence);
double loss = computeLoss(predictions, targets);
// 简化的反向传播（BPTT truncated）
// 实际应用中应使用完整的BPTT或更高效的优化器
return loss;
}
/**
* 预测下一个值
*/
public double[] predict(double[][] sequence) {
double[][] outputs = forward(sequence);
return outputs[outputs.length - 1];
}
/**
* 生成正弦波序列用于演示
*/
public static double[][] generateSineSequence(int length) {
double[][] seq = new double[length][1];
for (int i = 0; i < length; i++) {
seq[i][0] = Math.sin(i * 0.1);
}
return seq;
}
/**
* 准备训练数据：用前N个点预测第N+1个点
*/
public static void prepareData(double[] data, int seqLen,
List<double[][]> xList, List<double[][]> yList) {
for (int i = 0; i <= data.length - seqLen - 1; i++) {
double[][] x = new double[seqLen][1];
double[][] y = new double[seqLen][1];
for (int j = 0; j < seqLen; j++) {
x[j][0] = data[i + j];
y[j][0] = data[i + j + 1];
}
xList.add(x);
yList.add(y);
}
}
public static void main(String[] args) {
// 生成正弦波数据
double[] sineWave = new double[500];
for (int i = 0; i < sineWave.length; i++) {
sineWave[i] = Math.sin(i * 0.1);
}
// 准备数据
int seqLen = 10;
List<double[][]> xData = new ArrayList<>();
List<double[][]> yData = new ArrayList<>();
prepareData(sineWave, seqLen, xData, yData);
// 创建LSTM模型
LSTMPredictor lstm = new LSTMPredictor(1, 32, 1);
// 训练
System.out.println("开始训练LSTM模型...");
for (int epoch = 0; epoch < 100; epoch++) {
double totalLoss = 0;
for (int i = 0; i < xData.size(); i++) {
totalLoss += lstm.train(xData.get(i), yData.get(i));
}
if (epoch % 10 == 0) {
System.out.printf("Epoch %d, Avg Loss: %.6f%n", epoch, totalLoss / xData.size());
}
}
// 测试预测
System.out.println("\n=== 预测测试 ===");
double[][] testSeq = new double[seqLen][1];
for (int i = 0; i < seqLen; i++) {
testSeq[i][0] = Math.sin((400 + i) * 0.1);
}
double[] prediction = lstm.predict(testSeq);
double actual = Math.sin((400 + seqLen) * 0.1);
System.out.printf("输入序列最后值: %.4f%n", testSeq[seqLen - 1][0]);
System.out.printf("预测下一个值: %.4f%n", prediction[0]);
System.out.printf("实际下一个值: %.4f%n", actual);
System.out.printf("误差: %.4f%n", Math.abs(prediction[0] - actual));
}
}
代码解析
门控机制实现
代码中实现了三个核心门控：
– 遗忘门（f）：通过sigmoid输出0到1之间的值，决定保留多少历史信息
– 输入门（i）：控制新候选状态（cTilde）的写入比例
– 输出门（o）：决定隐藏状态的输出内容
细胞状态更新
细胞状态的更新公式 c = f * cPrev + i * cTilde 是LSTM的灵魂。遗忘门和输入门分别控制”忘记多少旧信息”和”接受多少新信息”，这种加性更新机制有效缓解了梯度消失问题。
参数初始化
遗忘门偏置初始化为1.0，这是LSTM的一个训练技巧，使模型在初始阶段倾向于保留信息而非遗忘。
复杂度分析
| 指标 | 复杂度 | 说明 |
|---|---|---|
| 前向传播 | O(T × H × (I + H)) | T为序列长度，H为隐藏层大小，I为输入维度 |
| 空间复杂度 | O(H × (I + H)) | 权重矩阵存储 |
| 参数量 | 4 × H × (I + H) + H × O | 四个门控权重加输出层 |
LSTM的变体与演进
- GRU（门控循环单元）：将遗忘门和输入门合并为更新门，参数量更少
- Peephole LSTM：让门控也能”窥视”细胞状态
- BiLSTM：双向LSTM同时利用过去和未来上下文
- Attention机制：让模型直接关注序列中的关键位置，部分替代了LSTM的作用
总结
LSTM通过引入细胞状态和三个门控机制， elegantly 解决了传统RNN的长程依赖问题。本文的Java实现展示了LSTM的核心数学结构和前向传播过程，虽然为教学目的简化了反向传播部分，但完整地呈现了遗忘门、输入门、输出门以及细胞状态更新的工作原理。理解LSTM是掌握现代序列模型（如Transformer）的重要基础。
