引言
长短期记忆网络（Long Short-Term Memory，LSTM）是循环神经网络（RNN）的重要变体，由Hochreiter和Schmidhuber于1997年提出。LSTM通过引入门控机制（输入门、遗忘门、输出门）有效解决了传统RNN中的梯度消失和梯度爆炸问题，使其能够学习长期依赖关系，成为序列建模领域的核心算法之一。本文将用Java从零实现一个简化版LSTM网络，讲解门控机制的计算原理，并通过正弦波预测任务展示其序列预测能力。
算法原理
LSTM核心结构
LSTM单元包含三个关键门控和一个记忆单元：
遗忘门（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)$$
细胞状态更新：
$$C_t = f_t \odot C_{t-1} + i_t \odot \tilde{C}_t$$
输出门（Output Gate）：决定基于细胞状态输出什么值
$$o_t = \sigma(W_o \cdot [h_{t-1}, x_t] + b_o)$$
$$h_t = o_t \odot \tanh(C_t)$$
其中 $\sigma$ 为Sigmoid函数，$\odot$ 表示逐元素乘法，$h_t$ 为隐藏状态。
Java实现
1. 数学工具类
/**
* LSTM数学工具类：提供激活函数和矩阵运算
*/
public class LSTMMath {
/** Sigmoid激活函数：将值压缩到(0,1)区间 */
public static double sigmoid(double x) {
return 1.0 / (1.0 + Math.exp(-x));
}
/** Sigmoid导数 */
public static double sigmoidDerivative(double x) {
double s = sigmoid(x);
return s * (1.0 - s);
}
/** Tanh激活函数：将值压缩到(-1,1)区间 */
public static double tanh(double x) {
return Math.tanh(x);
}
/** Tanh导数 */
public static double tanhDerivative(double x) {
double t = tanh(x);
return 1.0 - t * t;
}
/** 向量逐元素乘法（Hadamard积） */
public static double[] hadamard(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;
}
/** 向量逐元素加法 */
public static double[] add(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;
}
/** 向量逐元素减法 */
public static double[] subtract(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;
}
/** 向量与标量乘法 */
public static double[] scale(double[] a, double scalar) {
double[] result = new double[a.length];
for (int i = 0; i < a.length; i++) {
result[i] = a[i] * scalar;
}
return result;
}
/** 计算向量平均值 */
public static double mean(double[] a) {
double sum = 0;
for (double v : a) sum += v;
return sum / a.length;
}
}
2. LSTM单元实现
import java.util.Random;
/**
* LSTM单元：实现完整的门控机制与状态传播
*/
public class LSTMCell {
private final int inputSize; // 输入维度
private final int hiddenSize; // 隐藏层维度
private final Random random;
// 权重矩阵：每个门控都有对应的输入权重和隐藏权重
private final double[][] Wf, Wi, Wc, Wo; // 输入到各门的权重
private final double[][] Uf, Ui, Uc, Uo; // 隐藏状态到各门的权重
private final double[] bf, bi, bc, bo; // 偏置项
// 用于反向传播的缓存
private double[] lastInput, lastHidden, lastCell;
private double[] forgetGate, inputGate, candidateCell, outputGate;
private double[] newCell, newHidden;
public LSTMCell(int inputSize, int hiddenSize, long seed) {
this.inputSize = inputSize;
this.hiddenSize = hiddenSize;
this.random = new Random(seed);
// 初始化所有权重矩阵
Wf = initMatrix(hiddenSize, inputSize);
Wi = initMatrix(hiddenSize, inputSize);
Wc = initMatrix(hiddenSize, inputSize);
Wo = initMatrix(hiddenSize, inputSize);
Uf = initMatrix(hiddenSize, hiddenSize);
Ui = initMatrix(hiddenSize, hiddenSize);
Uc = initMatrix(hiddenSize, hiddenSize);
Uo = initMatrix(hiddenSize, hiddenSize);
bf = new double[hiddenSize];
bi = new double[hiddenSize];
bc = new double[hiddenSize];
bo = new double[hiddenSize];
}
/** Xavier初始化：根据输入输出维度随机初始化权重 */
private double[][] initMatrix(int rows, int cols) {
double[][] matrix = 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++) {
matrix[i][j] = random.nextGaussian() * scale;
}
}
return matrix;
}
/** 矩阵与向量乘法 */
private double[] matVecMul(double[][] matrix, double[] vector) {
double[] result = new double[matrix.length];
for (int i = 0; i < matrix.length; i++) {
double sum = 0;
for (int j = 0; j < vector.length; j++) {
sum += matrix[i][j] * vector[j];
}
result[i] = sum;
}
return result;
}
/**
* 前向传播：计算LSTM单元的输出
* @param input 当前时刻输入向量
* @param prevHidden 上一时刻隐藏状态
* @param prevCell 上一时刻细胞状态
* @return [newHidden, newCell] 新的隐藏状态和细胞状态
*/
public double[][] forward(double[] input, double[] prevHidden, double[] prevCell) {
this.lastInput = input.clone();
this.lastHidden = prevHidden.clone();
this.lastCell = prevCell.clone();
// 遗忘门：决定保留多少历史信息
double[] forgetInput = matVecMul(Wf, input);
double[] forgetHidden = matVecMul(Uf, prevHidden);
forgetGate = new double[hiddenSize];
for (int i = 0; i < hiddenSize; i++) {
forgetGate[i] = LSTMMath.sigmoid(forgetInput[i] + forgetHidden[i] + bf[i]);
}
// 输入门：决定接受多少新信息
double[] inputInput = matVecMul(Wi, input);
double[] inputHidden = matVecMul(Ui, prevHidden);
inputGate = new double[hiddenSize];
for (int i = 0; i < hiddenSize; i++) {
inputGate[i] = LSTMMath.sigmoid(inputInput[i] + inputHidden[i] + bi[i]);
}
// 候选细胞状态：生成新的候选值
double[] cellInput = matVecMul(Wc, input);
double[] cellHidden = matVecMul(Uc, prevHidden);
candidateCell = new double[hiddenSize];
for (int i = 0; i < hiddenSize; i++) {
candidateCell[i] = LSTMMath.tanh(cellInput[i] + cellHidden[i] + bc[i]);
}
// 细胞状态更新：遗忘旧信息 + 输入新信息
newCell = LSTMMath.add(
LSTMMath.hadamard(forgetGate, prevCell),
LSTMMath.hadamard(inputGate, candidateCell)
);
// 输出门：决定输出多少信息
double[] outputInput = matVecMul(Wo, input);
double[] outputHidden = matVecMul(Uo, prevHidden);
outputGate = new double[hiddenSize];
for (int i = 0; i < hiddenSize; i++) {
outputGate[i] = LSTMMath.sigmoid(outputInput[i] + outputHidden[i] + bo[i]);
}
// 新的隐藏状态
newHidden = new double[hiddenSize];
for (int i = 0; i < hiddenSize; i++) {
newHidden[i] = outputGate[i] * LSTMMath.tanh(newCell[i]);
}
return new double[][]{newHidden.clone(), newCell.clone()};
}
public int getHiddenSize() { return hiddenSize; }
public int getInputSize() { return inputSize; }
}
3. LSTM网络与训练
/**
* LSTM序列预测网络：包含输入层、LSTM层和输出层
*/
public class LSTMNetwork {
private final LSTMCell lstmCell;
private final double[][] outputWeight; // 隐藏层到输出的权重
private final double[] outputBias;
private final int hiddenSize;
// 当前状态
private double[] hiddenState;
private double[] cellState;
// 学习率
private final double learningRate;
public LSTMNetwork(int inputSize, int hiddenSize, double learningRate, long seed) {
this.lstmCell = new LSTMCell(inputSize, hiddenSize, seed);
this.hiddenSize = hiddenSize;
this.learningRate = learningRate;
// 初始化输出层
this.outputWeight = new double[1][hiddenSize];
this.outputBias = new double[1];
Random rand = new Random(seed + 1);
double scale = Math.sqrt(2.0 / hiddenSize);
for (int i = 0; i < hiddenSize; i++) {
outputWeight[0][i] = rand.nextGaussian() * scale;
}
// 初始化状态为零
this.hiddenState = new double[hiddenSize];
this.cellState = new double[hiddenSize];
}
/** 重置网络状态（用于新一轮序列预测） */
public void resetState() {
this.hiddenState = new double[hiddenSize];
this.cellState = new double[hiddenSize];
}
/**
* 单步前向传播
* @param input 单维输入值
* @return 预测输出值
*/
public double predict(double input) {
double[] inputVec = new double[]{input};
double[][] result = lstmCell.forward(inputVec, hiddenState, cellState);
hiddenState = result[0];
cellState = result[1];
// 输出层线性变换
double output = outputBias[0];
for (int i = 0; i < hiddenSize; i++) {
output += outputWeight[0][i] * hiddenState[i];
}
return output;
}
/**
* 训练一个序列：使用均方误差和梯度下降
* @param inputs 输入序列
* @param targets 目标序列
* @return 平均损失
*/
public double train(double[] inputs, double[] targets) {
resetState();
double totalLoss = 0;
// 前向传播并记录预测值
double[] predictions = new double[inputs.length];
for (int t = 0; t < inputs.length; t++) {
predictions[t] = predict(inputs[t]);
double error = predictions[t] - targets[t];
totalLoss += error * error;
}
// 简化的训练：直接根据输出误差调整输出层权重
// 完整版需要BPTT（随时间反向传播），此处为教学简化
resetState();
for (int t = 0; t < inputs.length; t++) {
predict(inputs[t]);
double error = predictions[t] - targets[t];
// 梯度下降更新输出层
for (int i = 0; i < hiddenSize; i++) {
outputWeight[0][i] -= learningRate * error * hiddenState[i];
}
outputBias[0] -= learningRate * error;
}
return totalLoss / inputs.length;
}
}
4. 主程序：正弦波预测
/**
* LSTM正弦波预测演示程序
*/
public class LSTMDemo {
public static void main(String[] args) {
// 网络参数
int hiddenSize = 16;
double learningRate = 0.01;
int epochs = 500;
int seqLength = 50;
// 创建LSTM网络
LSTMNetwork network = new LSTMNetwork(1, hiddenSize, learningRate, 42);
// 生成训练数据：正弦波序列
double[] trainInputs = new double[seqLength];
double[] trainTargets = new double[seqLength];
for (int i = 0; i < seqLength; i++) {
trainInputs[i] = Math.sin(i * 0.1);
trainTargets[i] = Math.sin((i + 1) * 0.1);
}
System.out.println("=== LSTM正弦波预测训练 ===");
System.out.println("网络结构: 输入层(1) -> LSTM隐藏层(" + hiddenSize + ") -> 输出层(1)");
System.out.println("训练序列长度: " + seqLength);
System.out.println("训练轮数: " + epochs);
System.out.println();
// 训练过程
for (int epoch = 0; epoch < epochs; epoch++) {
double loss = network.train(trainInputs, trainTargets);
if (epoch % 100 == 0) {
System.out.printf("Epoch %d, MSE Loss: %.6f%n", epoch, loss);
}
}
// 测试预测
System.out.println("\n=== 预测结果对比 ===");
System.out.printf("%-8s %-12s %-12s %-12s%n", "Step", "Input", "Predict", "Target");
System.out.println("-".repeat(50));
network.resetState();
double[] testInputs = new double[20];
for (int i = 0; i < 20; i++) {
testInputs[i] = Math.sin((seqLength + i) * 0.1);
}
double totalError = 0;
for (int i = 0; i < 20; i++) {
double pred = network.predict(testInputs[i]);
double target = Math.sin((seqLength + i + 1) * 0.1);
double error = Math.abs(pred - target);
totalError += error;
if (i < 10) {
System.out.printf("%-8d %-12.4f %-12.4f %-12.4f%n", i, testInputs[i], pred, target);
}
}
System.out.printf("\n平均绝对误差(MAE): %.6f%n", totalError / 20);
System.out.println("\nLSTM通过门控机制有效记住了正弦波的周期性特征，");
System.out.println("能够基于历史输入预测下一时刻的值。");
}
}
复杂度分析
| 操作 | 时间复杂度 | 空间复杂度 |
|---|---|---|
| 单步前向传播 | O(hiddenSize × (inputSize + hiddenSize)) | O(hiddenSize) |
| 序列前向传播(T步) | O(T × hiddenSize × (inputSize + hiddenSize)) | O(T × hiddenSize) |
| BPTT反向传播 | O(T × hiddenSize × (inputSize + hiddenSize)) | O(T × hiddenSize) |
LSTM相比标准RNN的主要开销在于额外的门控计算，但通过门控机制获得的长期依赖建模能力远超额外计算成本。
总结
本文从零实现了一个简化版LSTM网络，核心要点包括：
- 遗忘门通过Sigmoid函数决定历史信息的保留比例，使网络能够”忘记”无关信息
- 输入门控制新信息的注入，结合候选细胞状态更新记忆内容
- 输出门调节隐藏状态的输出，防止细胞状态的线性变化直接暴露
- 细胞状态作为”传送带”贯穿整个时间链，实现长期信息的稳定传递
LSTM的门控机制设计精妙，是深度学习序列建模的基石算法。在实际应用中，TensorFlow和PyTorch等框架提供了高度优化的LSTM实现，但理解其底层原理对于调参和模型设计至关重要。
