Transformer

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

---

参考资料

• 代码

• 掩码作用

• 李宏毅老师讲解Transformer

• 李沐老师讲解原论文