PyTorch教程-11.6. 自注意力和位置編碼

在深度學(xué)習(xí)中，我們經(jīng)常使用 CNN 或 RNN 對(duì)序列進(jìn)行編碼?，F(xiàn)在考慮到注意力機(jī)制，想象一下將一系列標(biāo)記輸入注意力機(jī)制，這樣在每個(gè)步驟中，每個(gè)標(biāo)記都有自己的查詢、鍵和值。在這里，當(dāng)在下一層計(jì)算令牌表示的值時(shí)，令牌可以（通過其查詢向量）參與每個(gè)其他令牌（基于它們的鍵向量進(jìn)行匹配）。使用完整的查詢鍵兼容性分?jǐn)?shù)集，我們可以通過在其他標(biāo)記上構(gòu)建適當(dāng)?shù)募訖?quán)和來為每個(gè)標(biāo)記計(jì)算表示。因?yàn)槊總€(gè)標(biāo)記都關(guān)注另一個(gè)標(biāo)記（不同于解碼器步驟關(guān)注編碼器步驟的情況），這種架構(gòu)通常被描述為自注意力模型 （Lin等。, 2017 年， Vaswani等人。, 2017 )，以及其他地方描述的內(nèi)部注意力模型 ( Cheng et al. , 2016 , Parikh et al. , 2016 , Paulus et al. , 2017 )。在本節(jié)中，我們將討論使用自注意力的序列編碼，包括使用序列順序的附加信息。

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

import math
from mxnet import autograd, np, npx
from mxnet.gluon import nn
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

No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)

import numpy as np
import tensorflow as tf
from d2l import tensorflow as d2l

11.6.1。自注意力

給定一系列輸入標(biāo)記 x1,…,xn任何地方 xi∈Rd(1≤i≤n), 它的self-attention輸出一個(gè)相同長(zhǎng)度的序列 y1,…,yn， 在哪里

(11.6.1)yi=f(xi,(x1,x1),…,(xn,xn))∈Rd

根據(jù) (11.1.1)中attention pooling的定義。使用多頭注意力，以下代碼片段計(jì)算具有形狀（批量大小、時(shí)間步數(shù)或標(biāo)記中的序列長(zhǎng)度， d). 輸出張量具有相同的形狀。

num_hiddens, num_heads = 100, 5
attention = d2l.MultiHeadAttention(num_hiddens, num_heads, 0.5)
batch_size, num_queries, valid_lens = 2, 4, torch.tensor([3, 2])
X = torch.ones((batch_size, num_queries, num_hiddens))
d2l.check_shape(attention(X, X, X, valid_lens),
        (batch_size, num_queries, num_hiddens))

num_hiddens, num_heads = 100, 5
attention = d2l.MultiHeadAttention(num_hiddens, num_heads, 0.5)
attention.initialize()

batch_size, num_queries, valid_lens = 2, 4, np.array([3, 2])
X = np.ones((batch_size, num_queries, num_hiddens))
d2l.check_shape(attention(X, X, X, valid_lens),
        (batch_size, num_queries, num_hiddens))

num_hiddens, num_heads = 100, 5
attention = d2l.MultiHeadAttention(num_hiddens, num_heads, 0.5)

batch_size, num_queries, valid_lens = 2, 4, jnp.array([3, 2])
X = jnp.ones((batch_size, num_queries, num_hiddens))
d2l.check_shape(attention.init_with_output(d2l.get_key(), X, X, X, valid_lens,
                      training=False)[0][0],
        (batch_size, num_queries, num_hiddens))

num_hiddens, num_heads = 100, 5
attention = d2l.MultiHeadAttention(num_hiddens, num_hiddens, num_hiddens,
                  num_hiddens, num_heads, 0.5)

batch_size, num_queries, valid_lens = 2, 4, tf.constant([3, 2])
X = tf.ones((batch_size, num_queries, num_hiddens))
d2l.check_shape(attention(X, X, X, valid_lens, training=False),
        (batch_size, num_queries, num_hiddens))

11.6.2。比較 CNN、RNN 和自注意力

讓我們比較一下映射一系列的架構(gòu)n標(biāo)記到另一個(gè)等長(zhǎng)序列，其中每個(gè)輸入或輸出標(biāo)記由一個(gè)d維向量。具體來說，我們將考慮 CNN、RNN 和自注意力。我們將比較它們的計(jì)算復(fù)雜度、順序操作和最大路徑長(zhǎng)度。請(qǐng)注意，順序操作會(huì)阻止并行計(jì)算，而序列位置的任意組合之間的較短路徑可以更容易地學(xué)習(xí)序列內(nèi)的遠(yuǎn)程依賴關(guān)系 （Hochreiter等人，2001 年）。

圖 11.6.1比較 CNN（省略填充標(biāo)記）、RNN 和自注意力架構(gòu)。

考慮一個(gè)卷積層，其內(nèi)核大小為k. 我們將在后面的章節(jié)中提供有關(guān)使用 CNN 進(jìn)行序列處理的更多詳細(xì)信息?，F(xiàn)在，我們只需要知道，因?yàn)樾蛄虚L(zhǎng)度是n，輸入和輸出通道的數(shù)量都是 d, 卷積層的計(jì)算復(fù)雜度為 O(knd2). 如圖11.6.1 所示，CNN 是分層的，因此有O(1) 順序操作和最大路徑長(zhǎng)度是 O(n/k). 例如，x1和 x5位于圖 11.6.1中內(nèi)核大小為 3 的雙層 CNN 的接受域內(nèi)。

在更新 RNN 的隱藏狀態(tài)時(shí)，乘以 d×d權(quán)重矩陣和d維隱藏狀態(tài)的計(jì)算復(fù)雜度為O(d2). 由于序列長(zhǎng)度為n，循環(huán)層的計(jì)算復(fù)雜度為O(nd2). 根據(jù) 圖 11.6.1，有O(n) 不能并行化的順序操作，最大路徑長(zhǎng)度也是O(n).

在自注意力中，查詢、鍵和值都是 n×d矩陣?？紤](11.3.6)中的縮放點(diǎn)積注意力，其中n×d矩陣乘以d×n矩陣，然后是輸出 n×n矩陣乘以n×d矩陣。因此，self-attention 有一個(gè)O(n2d) 計(jì)算復(fù)雜度。正如我們?cè)趫D 11.6.1中看到的 ，每個(gè)標(biāo)記都通過自注意力直接連接到任何其他標(biāo)記。因此，計(jì)算可以與O(1)順序操作和最大路徑長(zhǎng)度也是O(1).

總而言之，CNN 和 self-attention 都享有并行計(jì)算，并且 self-attention 具有最短的最大路徑長(zhǎng)度。然而，關(guān)于序列長(zhǎng)度的二次計(jì)算復(fù)雜度使得自注意力對(duì)于非常長(zhǎng)的序列來說非常慢。

11.6.3。位置編碼

與循環(huán)一個(gè)接一個(gè)地處理序列標(biāo)記的 RNN 不同，self-attention 摒棄順序操作以支持并行計(jì)算。但是請(qǐng)注意，self-attention 本身并不能保持序列的順序。如果模型知道輸入序列到達(dá)的順序真的很重要，我們?cè)撛趺崔k？

保留有關(guān)標(biāo)記順序的信息的主要方法是將其表示為與每個(gè)標(biāo)記相關(guān)聯(lián)的附加輸入的模型。這些輸入稱為位置編碼。它們可以被學(xué)習(xí)或先驗(yàn)固定。我們現(xiàn)在描述一種基于正弦和余弦函數(shù)的固定位置編碼的簡(jiǎn)單方案（Vaswani等人，2017 年）。

假設(shè)輸入表示 X∈Rn×d包含 d-維度嵌入n序列的標(biāo)記。位置編碼輸出X+P使用位置嵌入矩陣 P∈Rn×d形狀相同，其元素在ith行和 (2j)th或者(2j+1)th專欄是

(11.6.2)pi,2j=sin?(i100002j/d),pi,2j+1=cos?(i100002j/d).

乍一看，這種三角函數(shù)設(shè)計(jì)看起來很奇怪。在解釋這個(gè)設(shè)計(jì)之前，讓我們先在下面的 PositionalEncoding類中實(shí)現(xiàn)它。

class PositionalEncoding(nn.Module): #@save
  """Positional encoding."""
  def __init__(self, num_hiddens, dropout, max_len=1000):
    super().__init__()
    self.dropout = nn.Dropout(dropout)
    # Create a long enough P
    self.P = torch.zeros((1, max_len, num_hiddens))
    X = torch.arange(max_len, dtype=torch.float32).reshape(
      -1, 1) / torch.pow(10000, torch.arange(
      0, num_hiddens, 2, dtype=torch.float32) / num_hiddens)
    self.P[:, :, 0::2] = torch.sin(X)
    self.P[:, :, 1::2] = torch.cos(X)

  def forward(self, X):
    X = X + self.P[:, :X.shape[1], :].to(X.device)
    return self.dropout(X)

class PositionalEncoding(nn.Block): #@save
  """Positional encoding."""
  def __init__(self, num_hiddens, dropout, max_len=1000):
    super().__init__()
    self.dropout = nn.Dropout(dropout)
    # Create a long enough P
    self.P = np.zeros((1, max_len, num_hiddens))
    X = np.arange(max_len).reshape(-1, 1) / np.power(
      10000, np.arange(0, num_hiddens, 2) / num_hiddens)
    self.P[:, :, 0::2] = np.sin(X)
    self.P[:, :, 1::2] = np.cos(X)

  def forward(self, X):
    X = X + self.P[:, :X.shape[1], :].as_in_ctx(X.ctx)
    return self.dropout(X)

class PositionalEncoding(nn.Module): #@save
  """Positional encoding."""
  num_hiddens: int
  dropout: float
  max_len: int = 1000

  def setup(self):
    # Create a long enough P
    self.P = jnp.zeros((1, self.max_len, self.num_hiddens))
    X = jnp.arange(self.max_len, dtype=jnp.float32).reshape(
      -1, 1) / jnp.power(10000, jnp.arange(
      0, self.num_hiddens, 2, dtype=jnp.float32) / self.num_hiddens)
    self.P = self.P.at[:, :, 0::2].set(jnp.sin(X))
    self.P = self.P.at[:, :, 1::2].set(jnp.cos(X))

  @nn.compact
  def __call__(self, X, training=False):
    # Flax sow API is used to capture intermediate variables
    self.sow('intermediates', 'P', self.P)
    X = X + self.P[:, :X.shape[1], :]
    return nn.Dropout(self.dropout)(X, deterministic=not training)

class PositionalEncoding(tf.keras.layers.Layer): #@save
  """Positional encoding."""
  def __init__(self, num_hiddens, dropout, max_len=1000):
    super().__init__()
    self.dropout = tf.keras.layers.Dropout(dropout)
    # Create a long enough P
    self.P = np.zeros((1, max_len, num_hiddens))
    X = np.arange(max_len, dtype=np.float32).reshape(
      -1,1)/np.power(10000, np.arange(
      0, num_hiddens, 2, dtype=np.float32) / num_hiddens)
    self.P[:, :, 0::2] = np.sin(X)
    self.P[:, :, 1::2] = np.cos(X)

  def call(self, X, **kwargs):
    X = X + self.P[:, :X.shape[1], :]
    return self.dropout(X, **kwargs)

在位置嵌入矩陣中P，行對(duì)應(yīng)于序列中的位置，列代表不同的位置編碼維度。在下面的示例中，我們可以看到6th和7th位置嵌入矩陣的列具有比 8th和9th列。之間的偏移量6th和 7th（同樣的8th和 9th) 列是由于正弦和余弦函數(shù)的交替。

encoding_dim, num_steps = 32, 60
pos_encoding = PositionalEncoding(encoding_dim, 0)
X = pos_encoding(torch.zeros((1, num_steps, encoding_dim)))
P = pos_encoding.P[:, :X.shape[1], :]
d2l.plot(torch.arange(num_steps), P[0, :, 6:10].T, xlabel='Row (position)',
     figsize=(6, 2.5), legend=["Col %d" % d for d in torch.arange(6, 10)])

encoding_dim, num_steps = 32, 60
pos_encoding = PositionalEncoding(encoding_dim, 0)
pos_encoding.initialize()
X = pos_encoding(np.zeros((1, num_steps, encoding_dim)))
P = pos_encoding.P[:, :X.shape[1], :]
d2l.plot(np.arange(num_steps), P[0, :, 6:10].T, xlabel='Row (position)',
     figsize=(6, 2.5), legend=["Col %d" % d for d in np.arange(6, 10)])

encoding_dim, num_steps = 32, 60
pos_encoding = PositionalEncoding(encoding_dim, 0)
params = pos_encoding.init(d2l.get_key(), jnp.zeros((1, num_steps, encoding_dim)))
X, inter_vars = pos_encoding.apply(params, jnp.zeros((1, num_steps, encoding_dim)),
                  mutable='intermediates')
P = inter_vars['intermediates']['P'][0] # retrieve intermediate value P
P = P[:, :X.shape[1], :]
d2l.plot(jnp.arange(num_steps), P[0, :, 6:10].T, xlabel='Row (position)',
     figsize=(6, 2.5), legend=["Col %d" % d for d in jnp.arange(6, 10)])

encoding_dim, num_steps = 32, 60
pos_encoding = PositionalEncoding(encoding_dim, 0)
X = pos_encoding(tf.zeros((1, num_steps, encoding_dim)), training=False)
P = pos_encoding.P[:, :X.shape[1], :]
d2l.plot(np.arange(num_steps), P[0, :, 6:10].T, xlabel='Row (position)',
     figsize=(6, 2.5), legend=["Col %d" % d for d in np.arange(6, 10)])

11.6.3.1。絕對(duì)位置信息

為了了解沿編碼維度單調(diào)降低的頻率與絕對(duì)位置信息的關(guān)系，讓我們打印出的二進(jìn)制表示0,1,…,7. 正如我們所看到的，最低位、第二低位和第三低位分別在每個(gè)數(shù)字、每?jī)蓚€(gè)數(shù)字和每四個(gè)數(shù)字上交替出現(xiàn)。

for i in range(8):
  print(f'{i} in binary is {i:>03b}')

0 in binary is 000
1 in binary is 001
2 in binary is 010
3 in binary is 011
4 in binary is 100
5 in binary is 101
6 in binary is 110
7 in binary is 111

for i in range(8):
  print(f'{i} in binary is {i:>03b}')

0 in binary is 000
1 in binary is 001
2 in binary is 010
3 in binary is 011
4 in binary is 100
5 in binary is 101
6 in binary is 110
7 in binary is 111

for i in range(8):
  print(f'{i} in binary is {i:>03b}')

0 in binary is 000
1 in binary is 001
2 in binary is 010
3 in binary is 011
4 in binary is 100
5 in binary is 101
6 in binary is 110
7 in binary is 111

for i in range(8):
  print(f'{i} in binary is {i:>03b}')

0 in binary is 000
1 in binary is 001
2 in binary is 010
3 in binary is 011
4 in binary is 100
5 in binary is 101
6 in binary is 110
7 in binary is 111

在二進(jìn)制表示中，較高位的頻率比較低位低。類似地，如下面的熱圖所示，位置編碼通過使用三角函數(shù)降低編碼維度上的頻率。由于輸出是浮點(diǎn)數(shù)，因此這種連續(xù)表示比二進(jìn)制表示更節(jié)省空間。

P = P[0, :, :].unsqueeze(0).unsqueeze(0)
d2l.show_heatmaps(P, xlabel='Column (encoding dimension)',
         ylabel='Row (position)', figsize=(3.5, 4), cmap='Blues')

P = np.expand_dims(np.expand_dims(P[0, :, :], 0), 0)
d2l.show_heatmaps(P, xlabel='Column (encoding dimension)',
         ylabel='Row (position)', figsize=(3.5, 4), cmap='Blues')

P = jnp.expand_dims(jnp.expand_dims(P[0, :, :], axis=0), axis=0)
d2l.show_heatmaps(P, xlabel='Column (encoding dimension)',
         ylabel='Row (position)', figsize=(3.5, 4), cmap='Blues')

P = tf.expand_dims(tf.expand_dims(P[0, :, :], axis=0), axis=0)
d2l.show_heatmaps(P, xlabel='Column (encoding dimension)',
         ylabel='Row (position)', figsize=(3.5, 4), cmap='Blues')

11.6.3.2。相對(duì)位置信息

除了捕獲絕對(duì)位置信息外，上述位置編碼還允許模型輕松學(xué)習(xí)相對(duì)位置的注意。這是因?yàn)閷?duì)于任何固定位置偏移δ, 位置的位置編碼i+δ可以用位置的線性投影表示i.

這個(gè)投影可以用數(shù)學(xué)來解釋。表示 ωj=1/100002j/d, 任何一對(duì) (pi,2j,pi,2j+1)在 (11.6.2)中可以線性投影到 (pi+δ,2j,pi+δ,2j+1)對(duì)于任何固定偏移 δ:

(11.6.3)[cos?(δωj)sin?(δωj)?sin?(δωj)cos?(δωj)][pi,2jpi,2j+1]=[cos?(δωj)sin?(iωj)+sin?(δωj)cos?(iωj)?sin?(δωj)sin?(iωj)+cos?(δωj)cos?(iωj)]=[sin?((i+δ)ωj)cos?((i+δ)ωj)]=[pi+δ,2jpi+δ,2j+1],

在哪里2×2投影矩陣不依賴于任何位置索引i.

11.6.4。概括

在自我關(guān)注中，查詢、鍵和值都來自同一個(gè)地方。CNN 和 self-attention 都享有并行計(jì)算，并且 self-attention 具有最短的最大路徑長(zhǎng)度。然而，關(guān)于序列長(zhǎng)度的二次計(jì)算復(fù)雜度使得自注意力對(duì)于非常長(zhǎng)的序列來說非常慢。要使用序列順序信息，我們可以通過向輸入表示添加位置編碼來注入絕對(duì)或相對(duì)位置信息。

11.6.5。練習(xí)

假設(shè)我們?cè)O(shè)計(jì)了一個(gè)深度架構(gòu)來表示一個(gè)序列，通過使用位置編碼堆疊自注意力層。可能是什么問題？

你能設(shè)計(jì)一個(gè)可學(xué)習(xí)的位置編碼方法嗎？

我們能否根據(jù)在自注意力中比較的查詢和鍵之間的不同偏移來分配不同的學(xué)習(xí)嵌入？提示：你可以參考相對(duì)位置嵌入 （Huang et al. , 2018 , Shaw et al. , 2018）。