PyTorch教程-10.1. 長短期記憶 (LSTM)

在使用反向傳播（Elman，1990）訓練第一個 Elman 式 RNN 后不久，學習長期依賴性（由于梯度消失和爆炸）的問題變得突出，Bengio 和 Hochreiter 討論了這個問題 （Bengio等人， 1994 年，Hochreiter等人，2001 年）. Hochreiter 早在他 1991 年的碩士論文中就闡明了這個問題，盡管結果并不廣為人知，因為論文是用德語寫的。雖然梯度裁剪有助于梯度爆炸，但處理消失的梯度似乎需要更精細的解決方案。Hochreiter 和 Schmidhuber ( 1997 )提出的長短期記憶 (LSTM) 模型是解決梯度消失問題的第一個也是最成功的技術之一。LSTM 類似于標準的遞歸神經(jīng)網(wǎng)絡，但這里每個普通的遞歸節(jié)點都被一個記憶單元取代。每個存儲單元包含一個內部狀態(tài)，即具有固定權重 1 的自連接循環(huán)邊的節(jié)點，確保梯度可以跨越多個時間步而不會消失或爆炸。

“長短期記憶”一詞來自以下直覺。簡單的遞歸神經(jīng)網(wǎng)絡具有權重形式的長期記憶。權重在訓練過程中緩慢變化，對數(shù)據(jù)的一般知識進行編碼。它們還具有短暫激活形式的短期記憶，從每個節(jié)點傳遞到連續(xù)的節(jié)點。LSTM 模型通過記憶單元引入了一種中間類型的存儲。存儲單元是一個復合單元，由具有特定連接模式的較簡單節(jié)點構成，并包含新的乘法節(jié)點。

import torch
from torch import nn
from d2l import torch as d2l

from mxnet import np, npx
from mxnet.gluon import rnn
from d2l import mxnet as d2l

npx.set_np()

import jax
from flax import linen as nn
from jax import numpy as jnp
from d2l import jax as d2l

import tensorflow as tf
from d2l import tensorflow as d2l

10.1.1。門控存儲單元

每個存儲單元都配備了一個內部狀態(tài)和多個乘法門，用于確定 (i) 給定的輸入是否應該影響內部狀態(tài)（輸入門），(ii) 內部狀態(tài)是否應該被刷新到0（遺忘門），以及 (iii) 應該允許給定神經(jīng)元的內部狀態(tài)影響細胞的輸出（輸出門）。

10.1.1.1。門控隱藏狀態(tài)

普通 RNN 和 LSTM 之間的主要區(qū)別在于后者支持隱藏狀態(tài)的門控。這意味著我們有專門的機制來確定何時應該更新隱藏狀態(tài)以及何時應該重置它。這些機制是學習的，它們解決了上面列出的問題。例如，如果第一個標記非常重要，我們將學習在第一次觀察后不更新隱藏狀態(tài)。同樣，我們將學會跳過不相關的臨時觀察。最后，我們將學習在需要時重置潛在狀態(tài)。我們將在下面詳細討論。

10.1.1.2。輸入門、遺忘門和輸出門

進入 LSTM 門的數(shù)據(jù)是當前時間步的輸入和前一時間步的隱藏狀態(tài)，如圖 10.1.1所示。三個具有 sigmoid 激活函數(shù)的全連接層計算輸入門、遺忘門和輸出門的值。作為 sigmoid 激活的結果，三個門的所有值都在范圍內(0,1). 此外，我們需要一個 輸入節(jié)點，通常使用tanh激活函數(shù)計算。直觀上，輸入門決定了輸入節(jié)點的多少值應該添加到當前存儲單元的內部狀態(tài)。遺忘 門決定是保留內存的當前值還是刷新內存。而輸出門決定了記憶單元是否應該影響當前時間步的輸出。

圖 10.1.1計算 LSTM 模型中的輸入門、遺忘門和輸出門。

在數(shù)學上，假設有h隱藏單元，批量大小為n，輸入的數(shù)量是d. 因此，輸入是Xt∈Rn×d上一個時間步的隱藏狀態(tài)是 Ht?1∈Rn×h. 相應地，時間步長的門t定義如下：輸入門是It∈Rn×h, 遺忘門是 Ft∈Rn×h，輸出門是 Ot∈Rn×h. 它們的計算方式如下：

(10.1.1)It=σ(XtWxi+Ht?1Whi+bi),Ft=σ(XtWxf+Ht?1Whf+bf),Ot=σ(XtWxo+Ht?1Who+bo),

在哪里 Wxi,Wxf,Wxo∈Rd×h 和 Whi,Whf,Who∈Rh×h 是權重參數(shù)和 bi,bf,bo∈R1×h 是偏置參數(shù)。請注意，廣播（請參閱 第 2.1.4 節(jié)）是在求和期間觸發(fā)的。我們使用 sigmoid 函數(shù)（如第 5.1 節(jié)中介紹的）將輸入值映射到區(qū)間(0,1).

10.1.1.3。輸入節(jié)點

接下來我們設計存儲單元。由于我們還沒有具體說明各個門的作用，所以先介紹一下輸入節(jié)點 C~t∈Rn×h. 它的計算類似于上述三個門的計算，但使用tanh取值范圍為(?1,1)作為激活函數(shù)。這導致以下時間步等式t:

(10.1.2)C~t=tanh(XtWxc+Ht?1Whc+bc),

在哪里Wxc∈Rd×h和 Whc∈Rh×h是權重參數(shù)和bc∈R1×h是偏置參數(shù)。

輸入節(jié)點的快速圖示如圖 10.1.2所示 。

圖 10.1.2計算 LSTM 模型中的輸入節(jié)點。

10.1.1.4。存儲單元內部狀態(tài)

在 LSTM 中，輸入門It管理我們通過以下方式考慮新數(shù)據(jù)的程度C~t和遺忘門Ft解決舊單元格內部狀態(tài)的多少Ct?1∈Rn×h我們保留。使用 Hadamard（按元素）乘積運算符⊙我們得出以下更新方程：

(10.1.3)Ct=Ft⊙Ct?1+It⊙C~t.

如果遺忘門始終為 1，輸入門始終為 0，則存儲單元內部狀態(tài)Ct?1將永遠保持不變，不變地傳遞到每個后續(xù)時間步。然而，輸入門和遺忘門使模型能夠靈活地學習何時保持此值不變以及何時擾動它以響應后續(xù)輸入。在實踐中，這種設計緩解了梯度消失問題，導致模型更容易訓練，尤其是在面對序列長度較長的數(shù)據(jù)集時。

這樣我們就得到了圖 10.1.3中的流程圖。

圖 10.1.3計算 LSTM 模型中的存儲單元內部狀態(tài)。

10.1.1.5。隱藏狀態(tài)

最后，我們需要定義如何計算記憶單元的輸出，即隱藏狀態(tài)Ht∈Rn×h，正如其他層所見。這就是輸出門發(fā)揮作用的地方。在 LSTM 中，我們首先應用tanh到存儲單元內部狀態(tài)，然后應用另一個逐點乘法，這次是輸出門。這確保了值Ht總是在區(qū)間(?1,1):

(10.1.4)Ht=Ot⊙tanh?(Ct).

每當輸出門接近 1 時，我們允許記憶單元內部狀態(tài)不受抑制地影響后續(xù)層，而對于接近 0 的輸出門值，我們防止當前記憶在當前時間步影響網(wǎng)絡的其他層。請注意，一個記憶單元可以在不影響網(wǎng)絡其余部分的情況下跨越多個時間步長積累信息（只要輸出門的值接近于 0），然后一旦輸出門突然在后續(xù)時間步長影響網(wǎng)絡從接近 0 的值翻轉到接近 1 的值。

圖 10.1.4有數(shù)據(jù)流的圖形說明。

圖 10.1.4計算 LSTM 模型中的隱藏狀態(tài)。

10.1.2。從零開始實施

現(xiàn)在讓我們從頭開始實現(xiàn) LSTM。與9.5 節(jié)的實驗一樣 ，我們首先加載時間機器數(shù)據(jù)集。

10.1.2.1。初始化模型參數(shù)

接下來，我們需要定義和初始化模型參數(shù)。如前所述，超參數(shù)num_hiddens決定了隱藏單元的數(shù)量。我們按照標準偏差為 0.01 的高斯分布初始化權重，并將偏差設置為 0。

class LSTMScratch(d2l.Module):
  def __init__(self, num_inputs, num_hiddens, sigma=0.01):
    super().__init__()
    self.save_hyperparameters()

    init_weight = lambda *shape: nn.Parameter(torch.randn(*shape) * sigma)
    triple = lambda: (init_weight(num_inputs, num_hiddens),
             init_weight(num_hiddens, num_hiddens),
             nn.Parameter(torch.zeros(num_hiddens)))
    self.W_xi, self.W_hi, self.b_i = triple() # Input gate
    self.W_xf, self.W_hf, self.b_f = triple() # Forget gate
    self.W_xo, self.W_ho, self.b_o = triple() # Output gate
    self.W_xc, self.W_hc, self.b_c = triple() # Input node

實際模型的定義如上所述，由三個門和一個輸入節(jié)點組成。請注意，只有隱藏狀態(tài)會傳遞到輸出層。

@d2l.add_to_class(LSTMScratch)
def forward(self, inputs, H_C=None):
  if H_C is None:
    # Initial state with shape: (batch_size, num_hiddens)
    H = torch.zeros((inputs.shape[1], self.num_hiddens),
           device=inputs.device)
    C = torch.zeros((inputs.shape[1], self.num_hiddens),
           device=inputs.device)
  else:
    H, C = H_C
  outputs = []
  for X in inputs:
    I = torch.sigmoid(torch.matmul(X, self.W_xi) +
            torch.matmul(H, self.W_hi) + self.b_i)
    F = torch.sigmoid(torch.matmul(X, self.W_xf) +
            torch.matmul(H, self.W_hf) + self.b_f)
    O = torch.sigmoid(torch.matmul(X, self.W_xo) +
            torch.matmul(H, self.W_ho) + self.b_o)
    C_tilde = torch.tanh(torch.matmul(X, self.W_xc) +
              torch.matmul(H, self.W_hc) + self.b_c)
    C = F * C + I * C_tilde
    H = O * torch.tanh(C)
    outputs.append(H)
  return outputs, (H, C)

class LSTMScratch(d2l.Module):
  def __init__(self, num_inputs, num_hiddens, sigma=0.01):
    super().__init__()
    self.save_hyperparameters()

    init_weight = lambda *shape: np.random.randn(*shape) * sigma
    triple = lambda: (init_weight(num_inputs, num_hiddens),
             init_weight(num_hiddens, num_hiddens),
             np.zeros(num_hiddens))
    self.W_xi, self.W_hi, self.b_i = triple() # Input gate
    self.W_xf, self.W_hf, self.b_f = triple() # Forget gate
    self.W_xo, self.W_ho, self.b_o = triple() # Output gate
    self.W_xc, self.W_hc, self.b_c = triple() # Input node

The actual model is defined as described above, consisting of three gates and an input node. Note that only the hidden state is passed to the output layer.

@d2l.add_to_class(LSTMScratch)
def forward(self, inputs, H_C=None):
  if H_C is None:
    # Initial state with shape: (batch_size, num_hiddens)
    H = np.zeros((inputs.shape[1], self.num_hiddens),
           ctx=inputs.ctx)
    C = np.zeros((inputs.shape[1], self.num_hiddens),
           ctx=inputs.ctx)
  else:
    H, C = H_C
  outputs = []
  for X in inputs:
    I = npx.sigmoid(np.dot(X, self.W_xi) +
            np.dot(H, self.W_hi) + self.b_i)
    F = npx.sigmoid(np.dot(X, self.W_xf) +
            np.dot(H, self.W_hf) + self.b_f)
    O = npx.sigmoid(np.dot(X, self.W_xo) +
            np.dot(H, self.W_ho) + self.b_o)
    C_tilde = np.tanh(np.dot(X, self.W_xc) +
              np.dot(H, self.W_hc) + self.b_c)
    C = F * C + I * C_tilde
    H = O * np.tanh(C)
    outputs.append(H)
  return outputs, (H, C)

class LSTMScratch(d2l.Module):
  num_inputs: int
  num_hiddens: int
  sigma: float = 0.01

  def setup(self):
    init_weight = lambda name, shape: self.param(name,
                           nn.initializers.normal(self.sigma),
                           shape)
    triple = lambda name : (
      init_weight(f'W_x{name}', (self.num_inputs, self.num_hiddens)),
      init_weight(f'W_h{name}', (self.num_hiddens, self.num_hiddens)),
      self.param(f'b_{name}', nn.initializers.zeros, (self.num_hiddens)))

    self.W_xi, self.W_hi, self.b_i = triple('i') # Input gate
    self.W_xf, self.W_hf, self.b_f = triple('f') # Forget gate
    self.W_xo, self.W_ho, self.b_o = triple('o') # Output gate
    self.W_xc, self.W_hc, self.b_c = triple('c') # Input node

The actual model is defined as described above, consisting of three gates and an input node. Note that only the hidden state is passed to the output layer. A long for-loop in the forward method will result in an extremely long JIT compilation time for the first run. As a solution to this, instead of using a for-loop to update the state with every time step, JAX has jax.lax.scan utility transformation to achieve the same behavior. It takes in an initial state called carry and an inputs array which is scanned on its leading axis. The scan transformation ultimately returns the final state and the stacked outputs as expected.

@d2l.add_to_class(LSTMScratch)
def forward(self, inputs, H_C=None):
  # Use lax.scan primitive instead of looping over the
  # inputs, since scan saves time in jit compilation.
  def scan_fn(carry, X):
    H, C = carry
    I = jax.nn.sigmoid(jnp.matmul(X, self.W_xi) + (
      jnp.matmul(H, self.W_hi)) + self.b_i)
    F = jax.nn.sigmoid(jnp.matmul(X, self.W_xf) +
            jnp.matmul(H, self.W_hf) + self.b_f)
    O = jax.nn.sigmoid(jnp.matmul(X, self.W_xo) +
            jnp.matmul(H, self.W_ho) + self.b_o)
    C_tilde = jnp.tanh(jnp.matmul(X, self.W_xc) +
              jnp.matmul(H, self.W_hc) + self.b_c)
    C = F * C + I * C_tilde
    H = O * jnp.tanh(C)
    return (H, C), H # return carry, y

  if H_C is None:
    batch_size = inputs.shape[1]
    carry = jnp.zeros((batch_size, self.num_hiddens)), 
        jnp.zeros((batch_size, self.num_hiddens))
  else:
    carry = H_C

  # scan takes the scan_fn, initial carry state, xs with leading axis to be scanned
  carry, outputs = jax.lax.scan(scan_fn, carry, inputs)
  return outputs, carry

class LSTMScratch(d2l.Module):
  def __init__(self, num_inputs, num_hiddens, sigma=0.01):
    super().__init__()
    self.save_hyperparameters()

    init_weight = lambda *shape: tf.Variable(tf.random.normal(shape) * sigma)
    triple = lambda: (init_weight(num_inputs, num_hiddens),
             init_weight(num_hiddens, num_hiddens),
             tf.Variable(tf.zeros(num_hiddens)))

    self.W_xi, self.W_hi, self.b_i = triple() # Input gate
    self.W_xf, self.W_hf, self.b_f = triple() # Forget gate
    self.W_xo, self.W_ho, self.b_o = triple() # Output gate
    self.W_xc, self.W_hc, self.b_c = triple() # Input node

The actual model is defined as described above, consisting of three gates and an input node. Note that only the hidden state is passed to the output layer.

@d2l.add_to_class(LSTMScratch)
def forward(self, inputs, H_C=None):
  if H_C is None:
    # Initial state with shape: (batch_size, num_hiddens)
    H = tf.zeros((inputs.shape[1], self.num_hiddens))
    C = tf.zeros((inputs.shape[1], self.num_hiddens))
  else:
    H, C = H_C
  outputs = []
  for X in inputs:
    I = tf.sigmoid(tf.matmul(X, self.W_xi) +
            tf.matmul(H, self.W_hi) + self.b_i)
    F = tf.sigmoid(tf.matmul(X, self.W_xf) +
            tf.matmul(H, self.W_hf) + self.b_f)
    O = tf.sigmoid(tf.matmul(X, self.W_xo) +
            tf.matmul(H, self.W_ho) + self.b_o)
    C_tilde = tf.tanh(tf.matmul(X, self.W_xc) +
              tf.matmul(H, self.W_hc) + self.b_c)
    C = F * C + I * C_tilde
    H = O * tf.tanh(C)
    outputs.append(H)
  return outputs, (H, C)

10.1.2.2。訓練和預測

讓我們通過實例化第 9.5 節(jié)RNNLMScratch中介紹的類來訓練 LSTM 模型。

data = d2l.TimeMachine(batch_size=1024, num_steps=32)
lstm = LSTMScratch(num_inputs=len(data.vocab), num_hiddens=32)
model = d2l.RNNLMScratch(lstm, vocab_size=len(data.vocab), lr=4)
trainer = d2l.Trainer(max_epochs=50, gradient_clip_val=1, num_gpus=1)
trainer.fit(model, data)

data = d2l.TimeMachine(batch_size=1024, num_steps=32)
lstm = LSTMScratch(num_inputs=len(data.vocab), num_hiddens=32)
model = d2l.RNNLMScratch(lstm, vocab_size=len(data.vocab), lr=4)
trainer = d2l.Trainer(max_epochs=50, gradient_clip_val=1, num_gpus=1)
trainer.fit(model, data)

data = d2l.TimeMachine(batch_size=1024, num_steps=32)
lstm = LSTMScratch(num_inputs=len(data.vocab), num_hiddens=32)
model = d2l.RNNLMScratch(lstm, vocab_size=len(data.vocab), lr=4)
trainer = d2l.Trainer(max_epochs=50, gradient_clip_val=1, num_gpus=1)
trainer.fit(model, data)

data = d2l.TimeMachine(batch_size=1024, num_steps=32)
with d2l.try_gpu():
  lstm = LSTMScratch(num_inputs=len(data.vocab), num_hiddens=32)
  model = d2l.RNNLMScratch(lstm, vocab_size=len(data.vocab), lr=4)
trainer = d2l.Trainer(max_epochs=50, gradient_clip_val=1)
trainer.fit(model, data)

10.1.3。簡潔的實現(xiàn)

使用高級 API，我們可以直接實例化 LSTM 模型。這封裝了我們在上面明確說明的所有配置細節(jié)。該代碼明顯更快，因為它使用編譯運算符而不是 Python 來處理我們之前闡明的許多細節(jié)。

class LSTM(d2l.RNN):
  def __init__(self, num_inputs, num_hiddens):
    d2l.Module.__init__(self)
    self.save_hyperparameters()
    self.rnn = nn.LSTM(num_inputs, num_hiddens)

  def forward(self, inputs, H_C=None):
    return self.rnn(inputs, H_C)

lstm = LSTM(num_inputs=len(data.vocab), num_hiddens=32)
model = d2l.RNNLM(lstm, vocab_size=len(data.vocab), lr=4)
trainer.fit(model, data)

model.predict('it has', 20, data.vocab, d2l.try_gpu())

'it has the the the the the'

class LSTM(d2l.RNN):
  def __init__(self, num_hiddens):
    d2l.Module.__init__(self)
    self.save_hyperparameters()
    self.rnn = rnn.LSTM(num_hiddens)

  def forward(self, inputs, H_C=None):
    if H_C is None: H_C = self.rnn.begin_state(
      inputs.shape[1], ctx=inputs.ctx)
    return self.rnn(inputs, H_C)

lstm = LSTM(num_hiddens=32)
model = d2l.RNNLM(lstm, vocab_size=len(data.vocab), lr=4)
trainer.fit(model, data)

model.predict('it has', 20, data.vocab, d2l.try_gpu())

'it has and the time travel'

class LSTM(d2l.RNN):
  num_hiddens: int

  @nn.compact
  def __call__(self, inputs, H_C=None, training=False):
    if H_C is None:
      batch_size = inputs.shape[1]
      H_C = nn.OptimizedLSTMCell.initialize_carry(jax.random.PRNGKey(0),
                            (batch_size,),
                            self.num_hiddens)

    LSTM = nn.scan(nn.OptimizedLSTMCell, variable_broadcast="params",
            in_axes=0, out_axes=0, split_rngs={"params": False})

    H_C, outputs = LSTM()(H_C, inputs)
    return outputs, H_C

lstm = LSTM(num_hiddens=32)
model = d2l.RNNLM(lstm, vocab_size=len(data.vocab), lr=4)
trainer.fit(model, data)

model.predict('it has', 20, data.vocab, trainer.state.params)

'it has it it it it it it i'

class LSTM(d2l.RNN):
  def __init__(self, num_hiddens):
    d2l.Module.__init__(self)
    self.save_hyperparameters()
    self.rnn = tf.keras.layers.LSTM(
        num_hiddens, return_sequences=True,
        return_state=True, time_major=True)

  def forward(self, inputs, H_C=None):
    outputs, *H_C = self.rnn(inputs, H_C)
    return outputs, H_C

lstm = LSTM(num_hiddens=32)
with d2l.try_gpu():
  model = d2l.RNNLM(lstm, vocab_size=len(data.vocab), lr=4)
trainer.fit(model, data)

model.predict('it has', 20, data.vocab)

'it has the the the the the'

LSTM 是具有非平凡狀態(tài)控制的原型潛變量自回歸模型。多年來已經(jīng)提出了許多其變體，例如，多層、殘差連接、不同類型的正則化。然而，由于序列的長程依賴性，訓練 LSTM 和其他序列模型（例如 GRU）的成本非常高。稍后我們會遇到在某些情況下可以使用的替代模型，例如 Transformers。

10.1.4。概括

雖然 LSTM 于 1997 年發(fā)布，但隨著 2000 年代中期在預測競賽中的一些勝利，它們變得更加突出，并從 2011 年成為序列學習的主要模型，直到最近隨著 Transformer 模型的興起，從 2017 年開始。甚至變形金剛他們的一些關鍵想法歸功于 LSTM 引入的架構設計創(chuàng)新。LSTM 具有三種類型的門：輸入門、遺忘門和控制信息流的輸出門。LSTM 的隱藏層輸出包括隱藏狀態(tài)和記憶單元內部狀態(tài)。只有隱藏狀態(tài)被傳遞到輸出層，而記憶單元內部狀態(tài)完全是內部的。LSTM 可以緩解梯度消失和爆炸。

10.1.5。練習

調整超參數(shù)并分析它們對運行時間、困惑度和輸出序列的影響。

您需要如何更改模型以生成正確的單詞而不是字符序列？

比較給定隱藏維度的 GRU、LSTM 和常規(guī) RNN 的計算成本。特別注意訓練和推理成本。

由于候選記憶單元確保取值范圍介于?1和1通過使用tanh功能，為什么隱藏狀態(tài)需要使用tanh函數(shù)再次確保輸出值范圍介于?1和 1？

為時間序列預測而不是字符序列預測實施 LSTM 模型。