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

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

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

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。門控存儲單元

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

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

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

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

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

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

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

(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 是權(quán)重參數(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é)點(diǎn)

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

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

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

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

圖 10.1.2計(jì)算 LSTM 模型中的輸入節(jié)點(diǎn)。

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

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

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

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

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

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

10.1.1.5。隱藏狀態(tài)

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

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

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

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

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

10.1.2。從零開始實(shí)施

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

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

接下來，我們需要定義和初始化模型參數(shù)。如前所述，超參數(shù)num_hiddens決定了隱藏單元的數(shù)量。我們按照標(biāo)準(zhǔn)偏差為 0.01 的高斯分布初始化權(quán)重，并將偏差設(shè)置為 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

實(shí)際模型的定義如上所述，由三個門和一個輸入節(jié)點(diǎn)組成。請注意，只有隱藏狀態(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。訓(xùn)練和預(yù)測

讓我們通過實(shí)例化第 9.5 節(jié)RNNLMScratch中介紹的類來訓(xùn)練 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。簡潔的實(shí)現(xiàn)

使用高級 API，我們可以直接實(shí)例化 LSTM 模型。這封裝了我們在上面明確說明的所有配置細(xì)節(jié)。該代碼明顯更快，因?yàn)樗褂镁幾g運(yùn)算符而不是 Python 來處理我們之前闡明的許多細(xì)節(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)提出了許多其變體，例如，多層、殘差連接、不同類型的正則化。然而，由于序列的長程依賴性，訓(xùn)練 LSTM 和其他序列模型（例如 GRU）的成本非常高。稍后我們會遇到在某些情況下可以使用的替代模型，例如 Transformers。

10.1.4。概括

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

10.1.5。練習(xí)

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

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

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

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

為時間序列預(yù)測而不是字符序列預(yù)測實(shí)施 LSTM 模型。