長短期記憶網路(Long Short-Term Memory，LSTM)是一種專為處理序列資料所設計的遞迴神經網路(Recurrent Neural Network, RNN)，並被廣泛應用於時間序列預測、語音辨識、自然語言處理等需要捕捉時序依賴的任務。標準 RNN 在理論上可以記住任意長度的歷史資訊，但在實際訓練時，梯度會隨著時間步的增加而指數性地縮小或放大，導致難以學習長距離的依賴關係，此現象稱為梯度消失(vanishing gradient)或梯度爆炸(exploding gradient)問題。LSTM 透過引入細胞狀態(cell state)與三個閘門(gate)的設計，讓網路可以選擇性地記憶或遺忘資訊，從而有效緩解上述問題。

Long Short-Term Memory (LSTM) is a recurrent neural network (RNN) designed for sequence data. It is widely used in tasks that need to capture time dependencies, such as time series prediction, speech recognition, and natural language processing. In theory, a standard RNN can remember history of any length. But in real training, the gradient shrinks or grows exponentially as the number of time steps increases, so it is hard to learn long-distance dependencies. This is called the vanishing gradient or exploding gradient problem. LSTM adds a cell state and three gates, so the network can choose what to remember and what to forget. This eases the problem above.

在每個時間步 t，LSTM 接收當前輸入 xt 與上一步的狀態 ht-1、ct-1，並計算出 ht 以及 ct，即：

At each time step t, the LSTM takes the current input xt and the previous states ht-1 and ct-1, and computes ht and ct, that is:

(ht, ct) = LSTM(xt, ht-1, ct-1)

而其內部之詳細運算步驟，可分解如下：

The detailed steps inside can be broken down as follows:

ft = σ(Wf[ht-1, xt] + bf)

it = σ(Wi[ht-1, xt] + bi)

ot = σ(Wo[ht-1, xt] + bo)

c̃t = tanh(Wc[ht-1, xt] + bc)

ct = ft⊙ct-1 + it⊙c̃t

ht = ot⊙tanh(ct)

其中 W* 與 b* 為對應的權重與 bias，方括號表示矩陣並排，σ 為 sigmoid function，⊙ 為 element-wise product；ct-1 代表上一步的記憶，乘上 ft 代表上一步的記憶要保留多少；c̃t 代表這一步新產生的候選記憶，乘上 it 代表它要貢獻多少到這一步的新記憶當中。h0 與 c0 實務上通常使用 zero vector。

Here, W* and b* are the matching weights and biases, the square brackets mean putting the vectors side by side (concatenation), σ is the sigmoid function, and ⊙ is the element-wise product. ct-1 is the memory from the previous step, and multiplying it by ft decides how much of the previous memory to keep. c̃t is the new candidate memory made at this step, and multiplying it by it decides how much it adds to the new memory of this step. In practice, h0 and c0 are usually zero vectors.

上述的詳細運算步驟，若要用 PyTorch class 的方式表示，則寫法如下：

If we write the detailed steps above as a PyTorch class, it looks like this:

class LSTMCell(nn.Module):
	def __init__(self, input_size, hidden_size):
		super().__init__()
		self.W_f = nn.Linear(input_size + hidden_size, hidden_size)
		self.W_i = nn.Linear(input_size + hidden_size, hidden_size)
		self.W_o = nn.Linear(input_size + hidden_size, hidden_size)
		self.W_c = nn.Linear(input_size + hidden_size, hidden_size)

	def forward(self, x, h_prev, c_prev):
		xh = torch.cat([h_prev, x], dim=-1)
		f = torch.sigmoid(self.W_f(xh))
		i = torch.sigmoid(self.W_i(xh))
		o = torch.sigmoid(self.W_o(xh))
		c_tilde = torch.tanh(self.W_c(xh))
		c = f * c_prev + i * c_tilde
		h = o * torch.tanh(c)
		return h, c


if __name__ == '__main__':
	N = 2
	T = 5
	input_size = 8
	hidden_size = 16

	cell = LSTMCell(input_size, hidden_size)

	x = torch.rand(N, T, input_size)
	h = torch.zeros(N, hidden_size)
	c = torch.zeros(N, hidden_size)

	outputs = []
	for t in range(T):
		h, c = cell(x[:, t, :], h, c) # Shapes: (N, hidden_size)
		outputs.append(h)

	out = torch.stack(outputs, dim=1)
	print('Input shape:', x.shape)
	print('Output shape:', out.shape)
	print('Last h shape:', h.shape)
	print('Last c shape:', c.shape)

當然，實務上我們不必用自己撰寫的 LSTM cell 來搭建網路，而是使用 PyTorch 的 nn.LSTMCell 甚至 nn.LSTM，這兩者的差別在於，你是否需要自己撰寫迴圈，把每個時間步的資料一步一步一步地餵進模型。一個以 sine wave 對後者使用方式的示範如下：

Of course, in practice we don't need to build a network from our own LSTM cell. We use PyTorch's nn.LSTMCell, or even nn.LSTM. The difference between the two is whether you need to write the loop yourself, to feed the data into the model one time step at a time. Here is an example of using the latter with a sine wave:

import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

t_train = np.linspace(0, 100, 10001)
wave_train = np.sin(t_train).astype(np.float32)

seq_len = 50
X_train = np.array([wave_train[i:i + seq_len] for i in range(len(wave_train) - seq_len)])
y_train = wave_train[seq_len:]

X_train = torch.tensor(X_train).unsqueeze(-1) # (N, seq_len, 1)
y_train = torch.tensor(y_train).unsqueeze(-1) # (N, 1)

print('Data shapes:', X_train.shape, y_train.shape)

train_set = TensorDataset(X_train, y_train)
train_loader = DataLoader(train_set, batch_size=128, shuffle=True)


class LSTMModel(nn.Module):
	def __init__(self):
		super().__init__()
		self.lstm = nn.LSTM(input_size=1, hidden_size=64, batch_first=True)
		self.linear = nn.Linear(64, 1)

	def forward(self, x):
		out, _ = self.lstm(x)
		return self.linear(out[:, -1, :])


model = LSTMModel().to(device)
optim = torch.optim.Adam(model.parameters())
criterion = nn.MSELoss()

model.train()
for i in range(20):
	print(f'Epoch {i+1}/20')
	for X_batch, y_batch in train_loader:
		loss = criterion(model(X_batch.to(device)), y_batch.to(device))
		optim.zero_grad()
		loss.backward()
		optim.step()

t_end = 100 + 2 * np.pi
t_step = t_train[1] - t_train[0]
t_test = np.arange(100 + t_step, t_end + t_step, t_step, dtype=np.float32)

window = wave_train[-seq_len:].tolist()
preds = []

model.eval()
with torch.no_grad():
	for _ in t_test:
		x_in = torch.tensor(window[-seq_len:], dtype=torch.float32).unsqueeze(0).unsqueeze(-1).to(device)
		pred = model(x_in).item()
		preds.append(pred)
		window.append(pred)

y_true = np.sin(t_test)

plt.plot(t_test, y_true, label='Ground Truth')
plt.plot(t_test, preds, label='Prediction', linestyle='--')
plt.xlabel('t')
plt.ylabel('sin(t)')
plt.title('LSTM Rolling Prediction')
plt.legend()
plt.show()

在上述範例中：

In the example above:

值得一提的是，這種在預測時，把模型的輸出回饋為下一步輸入的方式，稱為自迴歸(autoregressive)預測。這類預測方式，除了用在 LSTM 以外，也常被用於 Transformer 或是大語言模型(Large Language Model, LLM)等處。

It is worth mentioning that, when predicting, feeding the model's output back in as the next input is called autoregressive prediction. Besides LSTM, this kind of prediction is also commonly used in Transformers and large language models (LLMs).

雖然 LSTM 實際上的應用，主要是跟時間序列有關的問題，但我們也可以將影像的其中一個軸視為時間軸，來套用 LSTM。以下是用 LSTM 來進行 MNIST 數字分類的範例：

Although LSTM is mainly used for problems related to time series, we can also treat one axis of an image as the time axis and use LSTM on it. Below is an example of using LSTM to classify MNIST digits:

import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

train_set = datasets.MNIST(root='./data', train=True, download=True, transform=transforms.ToTensor())
test_set = datasets.MNIST(root='./data', train=False, download=True, transform=transforms.ToTensor())
print('Data shapes:', train_set.data.shape, test_set.data.shape)

train_loader = DataLoader(train_set, batch_size=128, shuffle=True)
test_loader = DataLoader(test_set, batch_size=128)


class LSTMModel(nn.Module):
	def __init__(self):
		super().__init__()
		self.lstm = nn.LSTM(input_size=28, hidden_size=64, batch_first=True)
		self.linear = nn.Linear(64, 10)

	def forward(self, x):
		x = x.squeeze(1) # (N, 1, 28, 28) -> (N, 28, 28)
		out, _ = self.lstm(x)
		return self.linear(out[:, -1, :])


model = LSTMModel().to(device)
optim = torch.optim.Adam(model.parameters())
criterion = nn.CrossEntropyLoss()

model.train()
for i in range(10):
	print(f'Epoch {i+1}/10')
	for X_batch, y_batch in train_loader:
		loss = criterion(model(X_batch.to(device)), y_batch.to(device))
		optim.zero_grad()
		loss.backward()
		optim.step()

model.eval()
correct = sum(
	(model(X.to(device)).argmax(dim=1) == y.to(device)).sum().item()
	for X, y in test_loader
)
print('Accuracy: {:.2f}%'.format(100 * correct / len(test_set)))

在上述範例中，原始影像的 shape 是 (N, C=1, H=28, W=28)，我們用 squeeze(1) 來拿掉 channel 軸，並將 height 軸視為時間軸，餵給 LSTM。

In the example above, the shape of the original image is (N, C=1, H=28, W=28). We use squeeze(1) to remove the channel axis, treat the height axis as the time axis, and feed it into the LSTM.

既然 LSTM 的後面可以疊其他東西，那前面也可以。一個常見的作法是，把圖片等資料先用 CNN 處理過後，再餵給 LSTM 處理。若以基於上一個範例做修改來說，你只要把 class 替換為以下的版本即可：

Since we can stack other things after an LSTM, we can also put things before it. A common way is to process the image (or other data) with a CNN first, and then feed the result into the LSTM. If you modify the previous example, you only need to replace the class with the following version:

class CNNLSTMModel(nn.Module):
	def __init__(self):
		super().__init__()
		self.cnn = nn.Sequential(
			nn.Conv2d(1, 16, kernel_size=3, padding=1),
			nn.ReLU(),
			nn.MaxPool2d(2),
			nn.Conv2d(16, 32, kernel_size=3, padding=1),
			nn.ReLU(),
			nn.MaxPool2d(2)
		)
		self.lstm = nn.LSTM(input_size=32*7, hidden_size=64, batch_first=True)
		self.linear = nn.Linear(64, 10)

	def forward(self, x):
		x = self.cnn(x) # (N, ch=1, h=28, w=28) to (N, ch=32, h=7, w=7)
		x = x.permute(0, 3, 1, 2) # (N, ch=32, h=7, w=7) to (N, w=7, ch=32, h=7)
		x = x.flatten(2) # (N, w=7, ch=32, h=7) to (N, 7, 32*7)
		out, _ = self.lstm(x)
		return self.linear(out[:, -1, :])

在上述的 class 中，你需要把 CNN 處理完的結果中，適合當時間軸的維度移到前面來，才能餵給 LSTM。此外，LSTM 的輸入大小必須事先計算，因此你需要相當確定原始的輸入大小。

In the class above, you need to move the dimension that is suitable as a time axis to the front, after the CNN finishes processing, so that you can feed it into the LSTM. Also, the input size of the LSTM must be calculated in advance, so you need to be quite sure about the size of the original input.

關於前述內容中尚未示範的變體，包含但不限於以下：

Here are some variants that were not shown above, including but not limited to: