Attention
什么是注意力机制?
对于人类来说,注意力机制是在注意力有限的情况下,只关注接受信息的一部分,而忽略其他部分。
对于Transformer来说,以NLP为例,注意力机制就是对于当前token来说,为其所在序列中
对任务而言更重要的元素赋予更高权重(注意力)。
感知机可以认为是对不同选项赋予不同优先级(权重),以做出最终决策:
而注意力机制是对一个句子/一张图片内部元素赋予不同注意力(权重)的过程.
关于\(Q,K,V\)的Intuition
\(key-value\)可以理解给定一个键值对(map),对于一个询问\(query\),我们需要计算\(query\)作为键,
在给定键值对\(key-value\)对中其值是多少。例如我们有乒乓球选手各项能力以及其胜率的键值对:
如果查询是一个新选手的各项能力,我们想知道其对应胜率。
为回答这个问题,我们大概率不能在\(key-value\)中找到与\(query\)等值的\(key\),进而直接得到\(value\).
一个解决方案是求解\(query\)与所有\(key\)的相似度,在用相似度与\(value\)相乘,用最终的和作为\(qeury\)对应的值。
例如:
假设这位选手和键值对表中选手的相似度为\(0.7, 0.2, 0.1\), 那么我们可以预测这位选手的胜率为
\(0.7\times 0.7 + 0.2\times 0.5 + 0.1\times 0.8 = 0.67\).
对于Transformer来说,相似度是使用向量点积量化的。
什么是self-attention
以NLP问题为例,当计算注意力数值时,\(Q, K, V\)均来自某个句子序列自身(self).
Transformer
self-attention
An attention function can be described as mapping a query and a set of key-value pairs to an output, where the query, keys, values, and outputs are all vectors. The output is computed as a weighted sum of values, where the weight assigned to each value is computed by a compatibility function of the query with the corresponding key.
输入: 向量序列; 输出: 相同个数的向量,每个向量均考虑了所有输入向量.
具体的计算过程如下,以\(b_1\)为例:根据\(a_1\),考虑序列所有元素与\(a_1\)的相关性(权重越大相关性越大,将注意力集中
在相关性大的元素)
接下来的问题是,两个向量之间的相关程度\(a\)如何量化呢?一个方法是用两个向量的点积:
注意Transformer中在计算点积之前,通过线性投影矩阵将输入向量投影至\(query\)和\(value\)空间.
通过计算投影至\(query\)和\(value\)空间的向量点积,我们得到输入序列各元素与\(a_1\)的相关程度.
接着将相关程度通过\(Softmax\),得到归一化后的注意力权重. (非必须,使用\(ReLU\)或其他激活函数也可以,
只是Transformer是这样设计的)
根据得到的attention scores,从序列中提取与\(a_1\)有关的信息:
self-attention矩阵计算
Transformer中的self attention相比RNN的一个优势是计算可以并行化.
从\(a_i\)到\(Q, K, V\)的计算过程: \(Q = W^q I, K = W^k I, V = W^v I\).
通过\(Q, K\)计算attention score: \(A = Q K^T, A' = softmax(A)\).
通过注意力分数\(A'\)和\(V\)计算self attention的输出: \(O = V A'\):
综上,self attention的计算过程的矩阵乘法表示:
代码过程如下:
Self-attention
import torch import torch.nn as nn import math
NEG_INF = float('-inf')
class ScaledDotProductAttention(nn.Module):
def init(self,
d_model: int,
d_query_key: int,
d_value: int):
"""
d_model: 输入的词嵌入向量的长度
d_query_key: 单个query, key的向量长度
d_value: 单个value向量的长度
"""
super(ScaledDotProductAttention, self).init()
# key, query, value对应的线性投影矩阵
self.wk = nn.Parameter(torch.zeros(d_model, d_query_key))
self.wq = nn.Parameter(torch.zeros(d_model, d_query_key))
self.wv = nn.Parameter(torch.zeros(d_model, d_value))
self.div = math.sqrt(d_query_key) # 长的Q, K的点积数值过大会让softmax饱和
self.softmax = nn.Softmax(dim=2)def forward(self, mask: torch.Tensor, x_key_value: torch.Tensor, x_query: torch.Tensor): """ mask: bool Tensor, 为true的地方表示允许分配注意力 """ # 对于Encoder, Q, K, V的输入相同 # 对于Decoder, K, V来自Encoder最后一层输出, Q来自翻译后的token if x_query is None: x_query = x_key_value # linear transform k = x_key_value @ self.wk v = x_key_value @ self.wv q = x_query @ self.wq # batch matrix multiply: Q K.transpose(1, 2), 得到注意力分数 attention = torch.einsum('nik, njk -> nij', q, k) / self.div attention = torch.where(mask, attention, NEG_INF) attention = self.softmax(attention) output = torch.einsum('nij, njk -> nik', attention, v) return output
Multi-head self-attention
当我们获得\(Q, K, V\)时,我们可以用线性投影矩阵将其投影至多个空间,接着并行计算各个空间的
attention score, output. 最后将所有output链接作为最终输出.
这样做的目的是一方面模拟了卷积操作多个输出channel; 每个\(Q K^T\)可以认为是从一个视角的注意力
分布,计算不同形式的相关性可以获得不同维度的信息。
对\(Q, K, V\)做进一步投影,计算各自子空间中的attention score, output:
最后将输出concatenate, 并用一个线性变换\(W^O\)将信息投影至输出空间.
代码如下:
Multi-head Self-attention
import torch import torch.nn as nn
class MultiHeadAttention(nn.Module):
def init(self,
num_heads: int,
d_model: int,
p_drop: float=0.1):
super(MultiHeadAttention, self).init()# num_heads个注意力模块并行计算, 考虑残差连接以及计算效率每个头的输出维度为d_model(输入维度) / num_heads self.multi_head_attention = nn.ModuleList( ScaledDotProducAttention(d_model, d_model // num_heads, d_model // num_heads) for _ in range(num_heads) ) self.wo = nn.Parameter(torch.zeros(d_model, d_model)) # value空间线性投影至output空间 self.dropout = nn.Dropout(p_drop) def forward(self, mask: torch.Tensor, x_key_value: torch.Tensor, x_query: torch.Tensor=None): output = torch.concatenate([ attention(mask, x_key_value, x_query) for attention in self.multi_head_attention ], dim=2) output = output @ self.wo output = self.dropout(output) # residual connection output += x_key_value if x_query is None else x_query # layer normalization return torch.layer_norm(output, normalized_shape=output.size()[1:])
Positional Encoding
上述计算过程均可以用矩阵乘积实现,其中并不设计每个输入元素的位置信息(交换两个元素的位置对输出没有影响).
而语句或图片等信息是具有位置的结构信息,Positional Encoding通过对每个输入向量加入一个
唯一的等长向量,嵌入作为其位置特征。
Encoder
输入一串向量,输出相等数目向量。
代码实现:
EncoderLayer & Encoder
class EncoderLayer(nn.Module): """ x --> Multi-head Attention --> residual & layer norm --> Feed Forward --> residual & layer norm """ def __init__(self, num_heads: int, d_model: int, d_ffn: int, p_drop: float=0.1): """ d_model: 输入特征向量(embedding)长度 d_ffn: FFN隐藏层维度 """ super(EncoderLayer, self).__init__() self.muti_head_attention = MutiHeadAttention(num_heads, d_model, p_drop) self.feed_forward = PositionWiseFeedForward(d_model, d_ffn, p_drop)
def forward(self, padding_mask: torch.Tensor, x: torch.Tensor): """ x: embedding / 上一层EncoderLayer """ x = self.multi_head_attention(padding_mask, x) x = self.feed_forward(x) return x
class Encoder(nn.Module):
"""
N x EncoderLayer
"""
def init(self,
num_layers: int,
num_heads: int,
d_model: int,
d_ffn: int,
p_drop: float=0.1):
super(Encoder, self).init()
self.layers = nn.ModuleList(
EncoderLayer(num_heads, d_model, d_ffn, p_drop)
for _ in range(num_layers)
)def forward(self, padding_mask: torch.Tensor, src: torch.Tensor): for layer in self.layers: src = layer(padding_mask, src) return src
Decoder
Transformer中的Decoder是自回归的(Autoregressive), 接受来自Encoder的输入的同时,输出预测句子的单词token, 并且token作为
下一个单词预测的输入. 也就是说AT Decoder的输出逐个生成,且本次输出作为下一次预测的输入.
代码实现:
Masked Self-attention
对于输出\(b_i\)来说, 计算时只考虑\(a_{1\sim i}\).
why? :
masked代码实现: 在\(i\)之后的位置用一个很小的数覆盖,之后的Softmax处理会近似忽略这些位置.
Cross-attention
Decoder中的Self-attention的\(K, V\)来自Encoder输出,\(Q\)来自自身.
具体计算过程:
Decoder代码:
DecoderLayer & Decoder
class DecoderLayer(nn.Module): """ x - Masked Multi-head Attention - q (k, v), q - Multi-head Attention - redidual & layer norm - FFN - residual & layer norm """ def __init__(self, num_heads: int, d_model: int, d_ffn: int, p_drop: float=0.1): """ d_ffn: FFN隐藏层维度 """ super(DecoderLayer, self).__init__() self.masked_multi_head_attention = MultiHeadAttention(num_heads, d_model, p_drop) self.muli_head_attention = MultiHeadAttention(num_heads, d_model, p_drop) self.feed_forward = PositionWiseFeedForward(d_model, d_ffn, p_drop)
def forward(self, padding_mask: torch.Tensor, mask: torch.Tensor, x_decoder: torch.Tensor, x_encoder: torch.Tensor): x_decoder = self.masked_multi_head_attention(mask, x_decoder) output = self.muli_head_attention(padding_mask, x_encoder, x_decoder) output = self.feed_forward(output) return output
class Decoder(nn.Module):
"""
N x Decoder Layer
"""
def init(self,
num_layers: int,
num_heads: int,
d_model: int,
d_ffn: int,
p_drop: float=0.1):
super(Decoder, self).init()
self.layers = nn.ModuleList(
DecoderLayer(num_heads, d_model, p_drop)
for _ in range(num_layers)
)def forward(self, padding_mask: torch.Tensor, mask: torch.Tensor, tgt: torch.Tensor, encoder_out: torch.Tensor): for layer in self.layers: tgt = layer(padding_mask, mask, tgt, encoder_out) return tgt
Transformer完整架构
代码实现: https://github.com/CodesChangeHair/LearningDL/tree/main/Transformer