4. 注意力机制
4.1 注意力机制
循环神经网络的一个主要局限是不能很好地建模长距离依赖，即使像长短期记忆这样的变体也只是改善而不是完全解决了长距离依赖的问题。其根本原因在于，如果序列中的第i个词需要对第j个词(假设j>i)产生影响，需经过j-i个计算步骤， 而随着步数增加， 第i个词的信息会很快衰减，被两个词之间其他词的信息所淹没。从另一个角度来看， 每一步用来预测下一个词的隐状态都需要包含这个词左边所有词的信息，但隐状态的维度有限，因而所能表达的信息容量也有限，从而形成了信息瓶颈， 阻碍了前置词信息的准确表示和传递。
为了更好地建模长距离依赖,我们引入注意力机制( attention mechanism)。在每一步,我们直接在历史状态和当前状态之间建立联系。由于历史状态可能很多，在重要性和相关性上会有区别，因此我们希望模型能够自动预测这种重要性和相关性，这类似于人类的注意力。
具体而言， 注意力机制根据当前状态计算查询( query)， 根据每一个历史隐状态计算键( key), 进而计算查询与键的匹配程度,即注意力分数( attention score)。注意力分数越高,意味着对应的历史隐状态对当前时刻的预测越重要。对所有历史隐状态的注意力分数进行归一化， 我们就得到了对历史的注意力分布。接下来，以注意力分布的值作为权重，将所有历史隐状态所计算出的值( value) 向量进行加权平均，得到最终的注意力输出向量，用于代替当前状态向量来预测下一个词。
注意力分数有多种计算方式，下面给出的注意力机制的公式将查询和键的内积作为注意力分数, 即点乘注意力( dot- product attention)。假设某一组查询、键、值分别为q∈、
、
,其中查询和键的维度与值的维度可以是不同的，但是为了简单起见， 实现时可以令这些向量维度相等：
。
对于一个查询和整个序列的键、值来说，当前查询对应的注意力输出向量是：
其中， K和V分别是将整个序列的键向量和值向量堆叠而成的矩阵：。
为查询q在第i个位置的注意力分数，
为归一化项。
输入序列的每一步都需要进行上述注意力机制的计算。可以将输入序列所有位置上的注意力计算合并， 即将序列中所有步骤的查询堆叠为Q=[q_{1}, \cdots ,q_{n}],由此得到注意力计算的矩阵形式：
Attention(Q,K,V)= softmax(Q)V
需要注意的是，在前面的讲解中，一个查询会对整个序列的所有位置计算注意力，但是对于语言模型，第t步的查询应当只能看到该步及该步之前的输入。因此，需要引入注意力掩码( attention mask), 将每一步的查询对该步之后位置的注意力分数置为- inf。
这里所讲解的注意力机制，每一步的隐状态既用于计算当前步的查询， 又用于计算其他查询的键和值。也就是说，不考虑注意力掩码的话，我们是在输入序列的所有位置两两之间计算注意力， 即输入序列对于自身的注意力。
因此，这种特殊的注意力结构又称作自注意力( self attention)。区别于自注意力， 注意力机制本身更加通用，也适用于查询、键和值对应不同元素的场景.
基于点乘的注意力分数计算有一个潜在的问题，即随着查询和键的维度dₖ的增大，不同的键所计算的内积的数值范围也会逐渐增大，由此会带来 softmax函数的数值稳定性问题。为了解决这个问题,可以采用缩放点乘注意力( scaled dot- product attention), 即
下面实现一个带有缩放点乘注意力的循环神经网络，并用其训练语言模型。
python
class AttentionRNN(nn.Module):
def __init__(self, input_size, hidden_size):
super(AttentionRNN, self).__init__()
self.input_size = input_size
self.hidden_size = hidden_size
# 循环神经网络参数
self.W_xh = nn.Parameter(normal((input_size, hidden_size)))
self.W_hh = nn.Parameter(normal((hidden_size, hidden_size)))
self.b_h = nn.Parameter(torch.zeros(hidden_size))
def init_rnn_state(self, batch_size, hidden_size):
return (torch.zeros((batch_size, hidden_size),\
dtype=torch.float),)
# 缩放点乘注意力
def attention(self, query, keys, values):
"""
query: batch_size * hidden_size
keys/values: batch_size * prev_len * hidden_size
"""
# batch_size * 1 * hidden_size
query = torch.unsqueeze(query, 1)
# batch_size * hidden_size * prev_len
keys = torch.permute(keys, (0, 2, 1))
# batch_size * 1 * prev_len
attention_scores = torch.bmm(query, keys) / np.sqrt(\
self.hidden_size)
# batch_size * 1 * prev_len
attention_weights = F.softmax(attention_scores, dim=1)
# batch_size * hidden_size
attention_state = torch.squeeze(torch.bmm(attention_weights,\
values))
return attention_state
def forward(self, inputs, states):
seq_len, batch_size, _ = inputs.shape
hidden_state, = states
hiddens = []
attention_hiddens = []
for step in range(seq_len):
xh = torch.mm(inputs[step], self.W_xh)
hh = torch.mm(hidden_state, self.W_hh)
hidden_state = xh + hh + self.b_h
hidden_state = torch.tanh(hidden_state)
if step > 0:
# batch_size * hidden_size
query = hidden_state
# batch_size * prev_len * hidden_size
keys = values = torch.permute(torch.stack(hiddens,\
dim=0), (1, 0, 2))
attention_state = self.attention(query, keys, values)
attention_hiddens.append(attention_state)
else:
# 第0步，历史隐状态为空，无法进行注意力运算，
# 直接用隐状态填充
attention_hiddens.append(hidden_state)
hiddens.append(hidden_state)
return torch.stack(attention_hiddens, dim=0), \
(attention_state,)
data_loader = DataLoader(torch.tensor(sent_tokens, dtype=torch.long),
batch_size=16, shuffle=True)
attention_rnn = AttentionRNN(128, 128)
train_rnn_lm(data_loader, attention_rnn, vocab_size, hidden_size=128,
epochs=200, learning_rate=1e-3)
#%% md
python
epoch-199, loss=1.1505: 100%|█| 200/200 [17:54<00:00, 5.37s
值得一提的是，注意力机制在每一步都需要查看整个历史，所以当序列很长时，注意力机制的计算代价就会很高。这种情况下一般会设置一个固定大小的上下文窗口，将注意力局限于窗口之内。
4.2 多头注意力
普通的注意力只允许不同的词之间通过一种方式进行交互，这可能会限制模型的表达能力。一个改进方案是多头注意力( multi- head attention)， 即允许词之间通过多种不同方式进行交互， 具体做法如下。
将Q、K、V映射到m个不同的低维空间中。对i=1,2,...,m, 分别计算:
在每个低维空间中独立使用注意力机制：
=Attention(Qᵢ,Kᵢ,Vᵢ)
将不同低维空间的注意力输出向量拼接起来做一个线性变换， 其中表示将不同头的注意力输出向量拼接起来：
下面是多头注意力的代码实现。我们在实现AttentionRNN类时将注意力计算封装在成员函数里面， 因此实现多头注意力时可以直接继承AttentionRNN 类， 只需改写构造函数和 attention() 成员方法。
python
# 多头注意力循环神经网络
class MultiHeadAttentionRNN(AttentionRNN):
def __init__(self, input_size, hidden_size, num_heads=4):
super().__init__(input_size, hidden_size)
# 简单起见，一般要求hidden_size能够被num_heads整除
assert hidden_size % num_heads == 0
self.num_heads = num_heads
# 多头注意力参数，用于将查询、键、值映射到子空间
self.W_aq = nn.Parameter(normal((hidden_size, hidden_size)))
self.b_aq = nn.Parameter(torch.zeros(hidden_size))
self.W_ak = nn.Parameter(normal((hidden_size, hidden_size)))
self.b_ak = nn.Parameter(torch.zeros(hidden_size))
self.W_av = nn.Parameter(normal((hidden_size, hidden_size)))
self.b_av = nn.Parameter(torch.zeros(hidden_size))
self.W_ac = nn.Parameter(normal((hidden_size, hidden_size)))
self.b_ac = nn.Parameter(torch.zeros(hidden_size))
# 多头缩放点乘注意力
def attention(self, query, keys, values):
"""
query: batch_size * hidden_size
keys/values: batch_size * prev_len * hidden_size
"""
query = torch.mm(query, self.W_aq) + self.b_aq
ori_shape = keys.size()
keys = torch.reshape(torch.mm(torch.flatten(keys,
start_dim=0, end_dim=1), self.W_ak) +
self.b_ak, ori_shape)
values = torch.reshape(torch.mm(torch.flatten(values,
start_dim=0, end_dim=1), self.W_av) +
self.b_av, ori_shape)
# batch_size * 1 * hidden_size
query = torch.unsqueeze(query, 1)
# batch_size * hidden_size * prev_len
keys = torch.permute(keys, (0, 2, 1))
head_size = self.hidden_size // self.num_heads
query = torch.split(query, head_size, 2)
keys = torch.split(keys, head_size, 1)
values = torch.split(values, head_size, 2)
heads = []
for i in range(self.num_heads):
# batch_size * 1 * prev_len
head_scores = torch.bmm(query[i], keys[i]) / np.sqrt(
self.hidden_size // self.num_heads)
# batch_size * 1 * prev_len
head_weights = F.softmax(head_scores, dim=1)
# batch_size * head_size
head_state = torch.squeeze(torch.bmm(head_weights,
values[i]))
heads.append(head_state)
heads = torch.cat(heads, dim=1)
attention_state = torch.mm(heads, self.W_ac) + self.b_ac
return attention_state
data_loader = DataLoader(torch.tensor(sent_tokens,
dtype=torch.long), batch_size=16, shuffle=True)
mha_rnn = MultiHeadAttentionRNN(128, 128)
train_rnn_lm(data_loader, mha_rnn, vocab_size, hidden_size=128,
epochs=200, learning_rate=1e-3)
python
epoch-199, loss=1.3427: 100%|█| 200/200 [50:47<00:00, 15.24s
5.Transformer模型
上面在循环神经网络的基础上增加了注意力机制，循环神经网络基于循环连接来间接访问历史隐状态，而注意力机制能够直接访问历史隐状态。一个很自然的问题是， 能否去掉循环神经网络，只利用注意力机制来完成语言模型呢? 基于这样的想法， 我们就得到了 Transformer模型。
Transformer将循环神经网络中相邻隐状态之间的连接完全去除， 只保留注意力机制，因此不同位置的隐状态之间不存在计算上的依赖关系， 完全可以并行计算， 如下图所示。
并行计算是 Transformer相比于循环神经网络的一个显著优势。但是这样一来也引入了一个新的问题，即模型完全没有考虑词的顺序信息，而把输入文字序列看作词的集合，这对于建模自然语言而言显然是不妥的。为了解决这个问题，可以在模型中引入位置编码。一种做法是绝对位置编码，即给输入序列中的每个位置指定或者学习一个位置嵌入， 将其加到对应位置的词嵌入上作为模型的输入。另一种做法是相对位置编码，即在计算注意力时编码词之间的相对位置, 具体做法如下。首先计算相对位置索引：
rᵢⱼ=clip(i-j,-s,s)
其中， clip()是截断函数， s为预先设定的相对位置的截断上界。然后计算考虑了相对位置编码的注意力分数：
其中， 为与
维度相同的向量， 注意每个相对位置索引
都对应一个不同的向量
。最后在计算注意力输出向量时也加入另一组相对位置编码：
其中，类似于
,
为注意力分数
归一化后得到的权重。
仅使用注意力机制的另一个问题是，注意力输出只是对值向量进行了线性组合， 而以往的工作表明非线性变换可以增加模型的表达能力。因此，我们在注意力的输出上增加一个使用非线性激活函数的两层前馈神经网络( feed- forward neural network, FNN)。前馈神经网络有时也被称为多层感知机( multi- layer perceptron, MLP)。
其中，、
通常分别对其输入进行升维和降维操作(例如在经典的 Transformer 结构中, a=4b)。
Transformer模型会将上述注意力机制和前馈神经网络堆叠若干层，以增加模型的表达能力。这种方式类似于6.3.3节介绍的多层双向循环神经网络。为了增加这样的多层模型的训练稳定性, 降低训练难度, 我们进一步引入两个技巧。一是引入残差连接( residual connection).
xˡ=F(xˡ⁻¹)+xˡ⁻¹
其中，xˡ⁻¹为残差连接的输入，xˡ 为残差连接的输出, F()为一层注意力机制或前馈神经网络,二是层归一化 ( layer normalization) . 将每一层的输出归一化到均值为0、方差为1，再进行可学习的仿射变换：
其中，xˡ为层归一化的输入，μˡ与σ分别为输入.xˡ的均值和标准差，xˡ'是层归一化的输出， ∈是一个用于维持数值稳定性的很小的常数，gˡ和bˡ为可学习的仿射变换的参数。
如同上文所讨论的那样, 当 Transformer用于语言模型时， 还需要加上注意力掩码， 以保证每一步查询不会和该步之后的键、值做计算。最后，将模型顶层所输出的每个位置的隐状态输入一个线性分类器中， 得到下一个词的预测分布。整个 Transformer模型的架构如下图所示。
transformer结构图
下面来实现 Transformer模型， 包括加入了位置编码的嵌入层，缩放点注意力，多头注意意力、层归一化等具体实现。
python
# 实现Transformer模型
class EmbeddingLayer(nn.Module):
def __init__(self, vocab_size, max_len, embed_size):
super().__init__()
self.vocab_size = vocab_size
self.max_len = max_len
self.embed_size = embed_size
self.word_embedding = nn.Embedding(vocab_size, embed_size)
self.pos_embedding = nn.Embedding(max_len, embed_size)
def forward(self, input_ids, pos_ids):
"""
input_ids/pos_ids: batch_size * seq_len
return: batch_size * seq_len * embed_size
"""
word_embed = self.word_embedding(input_ids)
pos_embed = self.pos_embedding(pos_ids)
# 将词嵌入和位置嵌入相加得到嵌入层输出
return word_embed + pos_embed
# 缩放点乘注意力
class ScaledDotProductAttention(nn.Module):
def __init__(self, dropout):
super().__init__()
self.dropout = nn.Dropout(dropout)
def forward(self, queries, keys, values, attention_mask):
"""
queries/keys/values: batch_size * seq_len * hidden_size
attention_mask: batch_size * seq_len * seq_len
return: batch_size * seq_len * hidden_size
"""
d = queries.size(-1)
# 根据点乘注意力的矩阵形式计算注意力分数，除以查询向量或键向量
# 维度的平方根，即为缩放点乘注意力
scores = torch.bmm(queries, torch.transpose(keys, 1, 2)) / np.sqrt(d)
# 将掩码为0的位置的注意力分数设为一个大负数，根据softmax函数
# 的性质，这些注意力分数归一化后接近0
scores[attention_mask == 0] = -1e6
self.attention_weights = F.softmax(scores, dim=-1)
return torch.bmm(self.dropout(self.attention_weights), values)
class MultiHeadSelfAttention(nn.Module):
def __init__(self, hidden_size, num_heads, dropout):
super().__init__()
assert hidden_size % num_heads == 0
self.hidden_size = hidden_size
self.num_heads = num_heads
self.W_q = nn.Linear(hidden_size, hidden_size)
self.W_k = nn.Linear(hidden_size, hidden_size)
self.W_v = nn.Linear(hidden_size, hidden_size)
self.W_o = nn.Linear(hidden_size, hidden_size)
self.attention = ScaledDotProductAttention(dropout)
def transpose_qkv(self, states):
# 将长度为hidden_size的向量分成num_heads个长度相等的向量
states = states.reshape(states.shape[0], states.shape[1],\
self.num_heads, self.hidden_size // self.num_heads)
states = torch.permute(states, (0, 2, 1, 3))
return states.reshape(-1, states.shape[2], states.shape[3])
# 与transpose_qkv的变换相反
def transpose_output(self, states):
states = states.reshape(-1, self.num_heads, states.shape[1],\
states.shape[2])
states = torch.permute(states, (0, 2, 1, 3))
return states.reshape(states.shape[0], states.shape[1], -1)
def forward(self, queries, keys, values, attention_mask):
"""
querys/keys/values: batch * seq_len * hidden_size
attention_mask: batch * seq_len * seq_len
return:
"""
# (batch_size * num_heads) * seq_len * (hidden_size / num_heads)
queries = self.transpose_qkv(self.W_q(queries))
keys = self.transpose_qkv(self.W_k(keys))
values = self.transpose_qkv(self.W_v(values))
# 重复张量的元素，用以支持多个注意力头的运算
# (batch_size * num_heads) * seq_len * seq_len
attention_mask = torch.repeat_interleave(attention_mask,\
repeats=self.num_heads, dim=0)
# (batch_size * num_heads) * seq_len * (hidden_size / num_heads)
output = self.attention(queries, keys, values, attention_mask)
# batch * seq_len * hidden_size
output_concat = self.transpose_output(output)
return self.W_o(output_concat)
# 两层前馈神经网络
class PositionWiseFNN(nn.Module):
def __init__(self, hidden_size, intermediate_size):
super().__init__()
self.dense1 = nn.Linear(hidden_size, intermediate_size)
self.relu = nn.ReLU()
self.dense2 = nn.Linear(intermediate_size, hidden_size)
def forward(self, X):
return self.dense2(self.relu(self.dense1(X)))
# 层归一化
class LayerNorm(nn.Module):
def __init__(self, normalized_shape, eps=1e-6):
super().__init__()
self.gamma = nn.Parameter(torch.ones(normalized_shape))
self.beta = nn.Parameter(torch.zeros(normalized_shape))
# 一个小量用于数值稳定（防止除0）
self.eps = eps
def forward(self, hidden_states):
mean = torch.mean(hidden_states, -1, keepdim=True)
std = torch.std(hidden_states, -1, keepdim=True)
return self.gamma * (hidden_states - mean) / (std +\
self.eps) + self.beta
# 将两个输入相加并归一化
class AddNorm(nn.Module):
def __init__(self, hidden_size, dropout):
super().__init__()
self.dropout = nn.Dropout(dropout)
self.layer_norm = LayerNorm(hidden_size)
def forward(self, X, Y):
return self.layer_norm(self.dropout(Y) + X)
# 一个完整的Transformer层
class TransformerLayer(nn.Module):
def __init__(self, hidden_size, num_heads, dropout, intermediate_size):
super().__init__()
self.self_attention = MultiHeadSelfAttention(hidden_size,\
num_heads, dropout)
self.add_norm1 = AddNorm(hidden_size, dropout)
self.fnn = PositionWiseFNN(hidden_size, intermediate_size)
self.add_norm2 = AddNorm(hidden_size, dropout)
def forward(self, X, attention_mask):
Y = self.add_norm1(X, self.self_attention(X, X, X, attention_mask))
return self.add_norm2(Y, self.fnn(Y))
python
# 在Transformer模型基础上加上语言模型需要的输入输出、损失计算等
class TransformerLM(nn.Module):
def __init__(self, vocab_size, max_len, hidden_size, num_layers,\
num_heads, dropout, intermediate_size):
super().__init__()
self.embedding_layer = EmbeddingLayer(vocab_size, max_len,\
hidden_size)
self.num_layers = num_layers
# 使用ModuleList保存多个Transformer层，注意不能使用Python列表，
# Python列表保存的PyTorch变量无法自动求导
self.layers = nn.ModuleList([TransformerLayer(hidden_size,\
num_heads, dropout, intermediate_size)\
for _ in range(num_layers)])
self.output_layer = nn.Linear(hidden_size, vocab_size)
def forward(self, input_ids):
# 这里实现的forward()函数一次只能处理一句话，
# 如果想要支持批次运算，实现起来会更复杂，也会引入冗余操作
seq_len = input_ids.size(0)
assert input_ids.ndim == 1 and seq_len <= \
self.embedding_layer.max_len
# 1 * seq_len
input_ids = torch.unsqueeze(input_ids, dim=0)
pos_ids = torch.unsqueeze(torch.arange(seq_len), dim=0)
# 定义下三角掩码，用于语言模型训练
# 1 * seq_len * seq_len
attention_mask = torch.unsqueeze(torch.tril(torch.ones((seq_len,\
seq_len), dtype=torch.int32)), dim=0)
# 1 * seq_len * hidden_size
hidden_states = self.embedding_layer(input_ids, pos_ids)
for layer in self.layers:
hidden_states = layer(hidden_states, attention_mask)
outputs = self.output_layer(hidden_states)
loss_fct = nn.CrossEntropyLoss(ignore_index=0)
loss = loss_fct(outputs[:, :-1].squeeze(),\
input_ids[:, 1:].squeeze())
return loss
python
# 在Transformer模型基础上加上语言模型需要的输入输出、损失计算等
class TransformerLM(nn.Module):
def __init__(self, vocab_size, max_len, hidden_size, num_layers,\
num_heads, dropout, intermediate_size):
super().__init__()
self.embedding_layer = EmbeddingLayer(vocab_size, max_len,\
hidden_size)
self.num_layers = num_layers
# 使用ModuleList保存多个Transformer层，注意不能使用Python列表，
# Python列表保存的PyTorch变量无法自动求导
self.layers = nn.ModuleList([TransformerLayer(hidden_size,\
num_heads, dropout, intermediate_size)\
for _ in range(num_layers)])
self.output_layer = nn.Linear(hidden_size, vocab_size)
def forward(self, input_ids):
# 这里实现的forward()函数一次只能处理一句话，
# 如果想要支持批次运算，实现起来会更复杂，也会引入冗余操作
seq_len = input_ids.size(0)
assert input_ids.ndim == 1 and seq_len <= \
self.embedding_layer.max_len
# 1 * seq_len
input_ids = torch.unsqueeze(input_ids, dim=0)
pos_ids = torch.unsqueeze(torch.arange(seq_len), dim=0)
# 定义下三角掩码，用于语言模型训练
# 1 * seq_len * seq_len
attention_mask = torch.unsqueeze(torch.tril(torch.ones((seq_len,\
seq_len), dtype=torch.int32)), dim=0)
# 1 * seq_len * hidden_size
hidden_states = self.embedding_layer(input_ids, pos_ids)
for layer in self.layers:
hidden_states = layer(hidden_states, attention_mask)
outputs = self.output_layer(hidden_states)
loss_fct = nn.CrossEntropyLoss(ignore_index=0)
loss = loss_fct(outputs[:, :-1].squeeze(),\
input_ids[:, 1:].squeeze())
return loss
#%%
# 训练TransformerLM，由于不再采取批次训练，因此不再使用RNNLM和data_loader
def train_transformer_lm(data, model, epochs=50, learning_rate=1e-3):
optimizer = Adam(model.parameters(), lr=learning_rate)
model.zero_grad()
model.train()
epoch_loss = []
with trange(epochs, desc='epoch', ncols=60) as pbar:
for epoch in pbar:
step_loss = []
np.random.shuffle(data)
for step, x in enumerate(data):
loss = model(torch.tensor(x, dtype=torch.long))
pbar.set_description(f'epoch-{epoch},'+\
f' loss={loss.item():.4f}')
loss.backward()
grad_clipping(model)
optimizer.step()
model.zero_grad()
step_loss.append(loss.item())
# 本章前面的模型训练使用batch_size为16，
# TransformerLM出于简便实现只能使用batch_size为1
# 因此TransformerLM每一步的损失方差会更大，
# 为便于对比，取每个epoch最后16个样本的平均损失
epoch_loss.append(np.mean(step_loss[-16:]))
epoch_loss = np.array(epoch_loss)
plt.plot(range(len(epoch_loss)), epoch_loss)
plt.xlabel('training epoch')
plt.ylabel('loss')
plt.show()
sent_tokens = dataset.convert_tokens_to_ids()
max_len=40
for i, tokens in enumerate(sent_tokens):
tokens = tokens[:max_len]
tokens += [0] * (max_len - len(tokens))
sent_tokens[i] = tokens
sent_tokens = np.array(sent_tokens)
model = TransformerLM(vocab_size, max_len=40, hidden_size=128,\
num_layers=1, num_heads=4, dropout=0., intermediate_size=512)
train_transformer_lm(sent_tokens, model)
python
epoch-49, loss=0.3176: 100%|█| 50/50 [13:47<00:00, 16.54s/it