动手学深度学习 Part5

注意力机制#

注意力提示#

用于决定你要注意什么东西的一种东西, 例如你面前有十个杯子, 其中一个会发光, 你就会更容易注意到发光的杯子, 这就是一种注意力提示. 对模型来说, 在翻译 I ate an apple 中的 apple 时, 模型需要在整句英文中寻找哪些词和当前要生成的中文词最相关, 这也是一种注意力提示.

  • 非自主性提示: 信息本身的特征决定了注意力, 例如其中一个杯子突然爆炸了.
  • 自主性提示: 主动寻找某个东西 (Query), 例如我手机没了我在找手机.

只使用非自主性提示时, 只用简单的全连接层, 汇聚层等即可实现将选择偏向于感官输入, 因此关键在于自主性提示.

在深度学习中, 自主性提示被抽象为 Query, Key, Value. 给定一个 query, 注意力机制通过注意力汇聚将选择引导至感官输入 (value).

注意力汇聚#

注意力汇聚是根据 query, 把所有已知信息根据其和查询的相关程度, 加权汇总出一个输出:

Attention(q,K,V)=iα(q,ki)vi\text{Attention}(q, K, V) = \sum_{i} \alpha(q, k_i)v_i

其中 q 是 query, k 是 key, 即这条信息描述的内容是什么, value 是这条信息的实际内容.

也就是说, q 和 k 决定权重, 然后用权重对 v 加权求和.

Nadaraya-Watson 核回归#

例如有训练数据 (x1,y1),,(xn,yn)(x_1, y_1), \dots, (x_n, y_n) , 需要预测 xx 对应的 yy , 则最简单的想法是: 离 xx 越近的 xix_i, 对应的 yiy_i 越值得参考.

于是:

f(x)=iK(xxi)yiiK(xxi)f(x) = \frac{\sum_i K(x - x_i)y_i}{\sum_i K(x - x_i)}

这就是 N 一大串-W 一大串核回归.

KK 称为核函数 kernel, 用于衡量 xxxix_i 有多接近, 例如高斯核:

K(u)=exp(u22)K(u) = \exp\left(-\frac{u^2}{2}\right)

计算时将 xxix-x_i 代入 uu.

import torch
from torch import nn
import matplotlib.pyplot as plt

torch.manual_seed(66666666666666)

n_train = 150
x_train = torch.sort(torch.rand(n_train) * 5)[0]

def f(x):
return 2.333*torch.sin(x) + x**0.666

y_train = f(x_train) + torch.randn(n_train) * 0.5

x_test = torch.arange(0, 5, 0.1)
y_true = f(x_test)

# print(x_train.shape)
# print(y_train.shape)
# print(x_test.shape)

plt.scatter(x_train, y_train, label="train data")
plt.plot(x_test, y_true, label="true function")

plt.xlabel("x")
plt.ylabel("y")
plt.legend()
plt.show()

class NWKernelRegression(nn.Module):
def __init__(self):
super().__init__()
self.w = nn.Parameter(torch.tensor(1.0))

def forward(self, queries, keys, values):
dis = queries.unsqueeze(1) - keys
scores = -0.5 * (self.w * dis) ** 2

attention_weights = torch.softmax(scores, dim=1)
y_hat = torch.sum(attention_weights * values, dim=1)
return y_hat

n_train = len(x_train)
mask = ~torch.eye(n_train, dtype=torch.bool)
keys_train = x_train.repeat(n_train, 1)[mask].reshape(
n_train,
n_train - 1
)
values_train = y_train.repeat(n_train, 1)[mask].reshape(
n_train,
n_train - 1
)



model = NWKernelRegression()
loss_fn = nn.MSELoss()
optimizer = torch.optim.SGD(
model.parameters(),
lr=0.5
)

epochs = 20
for epoch in range(epochs):
y_hat = model(x_train, keys_train, values_train)
loss = loss_fn(y_hat, y_train)
optimizer.zero_grad()
loss.backward()
optimizer.step()
print(
f"epoch {epoch + 1:02d}, "
f"loss = {loss.item():.4f}, "
f"w = {model.w.item():.4f}"
)

keys_test = x_train.repeat(
len(x_test),
1
)

values_test = y_train.repeat(
len(x_test),
1
)

with torch.no_grad():
y_hat = model(
x_test,
keys_test,
values_test
)

plt.scatter(x_train, y_train, alpha=0.5, label="train data")

plt.plot(x_test, y_true, label="true function")

plt.plot(x_test, y_hat, label="NW prediction")

plt.xlabel("x")
plt.ylabel("y")
plt.legend()

plt.show()

注意力评分函数#

用于计算 query 和每个 key 的匹配程度的函数, 前面所说的核函数属于注意力评分函数, 但注意力评分函数的概念更广.

缩放点积注意力#

最简单的想法是直接进行点积, 但点积会随着维度增大而增大, 所以我们有缩放点积注意力:

a(q,k)=qTkda(q,k)=\frac{q^Tk}{\sqrt d}

完整公式:

O=softmax(QKTd)VO= \operatorname{softmax} \left( \frac{QK^T}{\sqrt d} \right)V

形状为:

(m,n)(n,dv)=(m,dv)(m,n)(n,d_v) = (m,d_v)

考虑一句话 I love deep learning, embedding 维度为 8, 则有:

Q: (4, 8)
K: (4, 8)
V: (4, 8)

那么计算 Q @ K.T 得到的结果形状为 (4, 4), 注意力矩阵类似:

A=[0.50.20.20.10.10.20.50.20.10.10.30.50.20.20.40.2]A= \begin{bmatrix} 0.5&0.2&0.2&0.1\\ 0.1&0.2&0.5&0.2\\ 0.1&0.1&0.3&0.5\\ 0.2&0.2&0.4&0.2 \end{bmatrix}

这其中, 第一行表示 token I 分别应该关注 I / love / deep / learning 多少, 以此类推.

Attention 矩阵本质上就是「谁关注谁」的关系矩阵.

加性注意力#

query 和 key 维度不同时无法点积, 因此将两者投影到一个共同空间:

a(q,k)=wvTtanh(Wqq+Wkk)a(\mathbf q,\mathbf k) = \mathbf w_v^T \tanh( W_q\mathbf q + W_k\mathbf k )

其中 WqW_q 把 query 投影到一个隐藏空间, WkW_k 把 key 投影到同一个隐藏空间, 先分别线性变换再相加. 注意这里的 wvw_v 不是 value, 只是一个可学习参数.

Bahdanau 注意力#

在早期模型中, Encoder 会把整个句子压缩成一个固定长度向量, 但这显然不能完整保存一个很长的句子.

Bahdanau 提出, Encoder 本来就产生了一系列隐藏状态:

h1,h2,,hTh_1,h_2,\cdots,h_T

分别对应输入序列的不同位置, 则 Decoder 在生成每个单词时, 可以直接动态计算此时应该关注哪个 hih_i.

此时 Query 是 Decoder 当前的隐藏状态 st1s_{t-1} , Key 和 Value 都是hih_i.

使用的是加性注意力, 公式:

et,i=vaTtanh(Wsst1+Whhi)e_{t,i} = v_a^T \tanh( W_s s_{t-1} + W_h h_i )

多头注意力#

对句子:

The animal didn't cross the street because it was too tired.

我们需要理解 it 指代的是谁, 也就是理解指代关系. 但一个 token 可能还需要同时关注语法关系, 时间关系, 位置关系等, 而一个 attention 只能形成一套注意力权重, 于是 transformer 的思路就是同时做很多套 attention.

例如, 一个 512 维的模型, 有 h=8h=8 个注意力头, 那么每个 head 在自己的 64 维空间里做 attention, 可能 head1 关注语法关系, head2 关注代词等等.

对 Q, K, V 分别进行不同线性变换:

Qi=QWiQKi=KWiKVi=VWiV\begin{aligned} Q_i &= QW_i^Q \\ K_i &= KW_i^K \\ V_i &= VW_i^V \end{aligned}

对第 i 个 head:

headi=Attention(QWiQ,KWiK,VWiV)head_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V)

Attention 一般是:

headi=softmax(QiKiTdk)Vihead_i = \text{softmax} \left( \frac{Q_i K_i^T}{\sqrt{d_k}} \right) V_i

然后把所有 head 拼起来, 最后乘一个输出矩阵:

Concat(head1,,headh)MultiHead(Q,K,V)=Concat(head1,,headh)WO\text{Concat}(head_1, \dots, head_h) \\ \text{MultiHead}(Q, K, V) = \text{Concat}(head_1, \dots, head_h) W^O

这里不同头看到的是输入向量的不同投影.

WOW^O 可以重新融合不同 head 得到的信息.

自注意力和位置编码#

例如「我喜欢玩原神」, 我们希望让「玩」也包含「我」, 「喜欢」, 「原神」的意思. 于是我们可以对于第 i 个词, 让 xix_i 作为 query, 查询整个序列中的所有词, 这就是自注意力.

这里需要三个线性变换:

Q=XWQK=XWKV=XWVQ = XW_Q \\ K = XW_K \\ V = XW_V

Query 表示这个词想找什么信息, Key 表示这个词可以被通过什么特征找到, Value 表示这个词真正应该被拿走的信息.

计算注意力评分 qikjq_i^\top k_j , 然后做 softmax:

αij=eqikj/dkleqikl/dk\alpha_{ij} = \frac{e^{q_i^\top k_j / \sqrt{d_k}}}{\sum_l e^{q_i^\top k_l / \sqrt{d_k}}}

最终输出:

oi=jαijvjo_i = \sum_{j} \alpha_{ij} v_j

注意力矩阵是一个 n×nn \times n 的矩阵, 表示序列中每一个 token 对每一个 token 的注意力评分.

自注意力自身没有序列顺序的概念, 因此需要引入位置编码, 另外制造一个位置矩阵 P, 然后 X=X+PX = X+P'.

正弦位置编码:

Pi,2j=sin(i100002j/d)Pi,2j+1=cos(i100002j/d)P_{i,2j} = \sin\left( \frac{i}{10000^{2j/d}} \right) \\ P_{i,2j+1} = \cos\left( \frac{i}{10000^{2j/d}} \right)

不仅允许学习绝对位置信息, 还允许学习相对位置信息.

不同维度使用不同频率.

Transformer#

在传统 RNN/GRU/LSTM 的 Encoder-Decoder 中, 输入通常要按顺序处理, 例如 I→love→cats, 同时隐藏状态不断传递:

ht=f(ht1,xt)h_t=f(h_{t−1},x_t)

这样有两个问题:

  • 不能充分并行
  • 长距离依赖的路径很长

于是 Transformer 通过自注意力让第 100 个词可以直接看到第 1 个词.

Transformer的架构

输入首先经过 embedding, 每个 token 变为一个向量, 然后加入位置编码.

Encoder 的一个 Block 中, Add & Norm 进行残差连接和归一化, 让深层网络更容易训练, 稳定梯度, 保留原始信息.

逐位前馈网络 (Position-wise Feed-Forward Network, FFN) 是一个两层全连接网络, 对序列里的每个 token 分别做处理, 对每个 token 进一步非线性加工. 一般会先升维, 再降维.

掩蔽多头注意力即 Masked Multi-Head Attention, 主要在 decoder 中使用, 防止模型在当前词汇中看到未来的词.