Attention Is All You Need

Transformer

Self-Attention

Self-Attention机制是一种在序列处理中非常重要的技术，它允许模型在处理一个序列的特定位置时，能够考虑到序列中的其他位置，从而捕捉序列内部的复杂依赖关系。

自注意力机制的核心思想是通过计算序列中每个元素对其他元素的影响权重来实现内部的关联，这些权重反映了元素之间的相似性或者关联性。具体来说，自注意力机制包含以下几个步骤：

1. 自注意力机制首先将输入序列的每个元素转换成三个不同的向量，分别是Query（查询）、Key（键）和Value（值）。这通常是通过将输入向量与三个不同的权重矩阵相乘来实现的。这样，对于序列中的每个元素，我们都有对应的三个向量。

2. 接下来，自注意力机制计算每个Query与其他所有元素的Key之间的点积，得到一个分数，这个分数反映了Query与每个元素的匹配程度。然后，通过应用一个softmax函数，将这些分数转换成概率，即注意力权重。

3. 一旦我们有了注意力权重，就可以用它们来加权对应的Value向量。具体来说，就是将每个Value与其对应的注意力权重相乘，然后将所有加权后的Value向量求和，得到最终的输出向量。

4. 为了稳定学习过程并提高模型性能，自注意力机制通常会并行地执行多次（称为“头”），每个头都有自己的权重矩阵。最后，这些头的输出会被拼接起来，并通过一个线性变换，得到最终的输出。

Masked Self-Attention

• Masked self-attention是对self-attention的一种改进，它在计算注意力时引入了masking机制。
• 遮蔽机制通过将某些位置的注意力权重设置为一个很小的值（通常是负无穷大），从而阻止模型在这些位置上进行关注。
• 这种机制通常用于处理那些模型在训练时不应该关注的信息，例如在语言模型中，当模型尝试预测下一个词时，它不应该关注到未来的词汇。在序列生成任务中，如机器翻译或文本生成，masked self-attention确保模型在生成当前位置的输出时，只能使用之前的元素，而不能使用之后未来的元素。

Scaled Dot-Product Attention

为什么需要除以？
作者在原文中的解释：

While for small values of the two mechanisms perform similarly, additive attention outperforms dot product attention without scaling for larger values of . We suspect that for large values of , the dot products grow large in magnitude, pushing the softmax function into regions where it has extremely small gradients. To counteract this effect, we scale the dot products by .

Multi-Head Attention

import math
import torch

def scaled_dot_product_attention(query, key, value, query_mask=None, key_mask=None, mask=None):
    dim_k = key.size(-1)
    scores = torch.bmm(query, key.transpose(1, 2)) / math.sqrt(dim_k)
    if query_mask is not None and key_mask is not None:
        mask = torch.bmm(query_mask.unsqueeze(-1), key_mask.unsqueeze(1))
    if mask is not None:
        scores = scores.masked_fill(mask == 0, -float('inf'))
    weights = torch.nn.functional.softmax(scores, dim=-1)
    return torch.bmm(weights, value)

class AttentionHead(nn.Module):
    def __init__(self, embed_dim, head_dim):
        super().__init__()
        self.q = nn.Linear(embed_dim, head_dim)
        self.k = nn.Linear(embed_dim, head_dim)
        self.v = nn.Linear(embed_dim, head_dim)
        
    def forward(self, query, key, value, query_mask=None, key_mask=None, mask=None):
        attn_outputs = scaled_dot_product_attention(
            self.q(query), self.k(key), self.v(value), query_mask, key_mask, mask
        )
        return attn_outputs

class MultiHeadAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        embed_dim = config.hidden_size
        heads_num = config.num_attention_heads
        head_dim = embed_dim // heads_num
        self.heads = nn.ModuleList([AttentionHead(embed_dim, head_dim) for _ in range(heads_num)])
        self.output_linear = nn.Linear(embed_dim, embed_dim)
        
    def forward(self, query, key, value, query_mask=None, key_mask=None, mask=None):
        x = torch.cat(
            [head(query, key, value, query_mask=None, key_mask=None, mask=None) for head in self.heads],
            dim=-1
        )
        x = self.output_linear(x)
        return x

from torch import nn
from transformers import AutoTokenizer
from transformers import AutoConfig

model_name = "bert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(model_name)

text = "to the sound of stirring iced tea"
tokens = tokenizer(text, return_tensors="pt", add_special_tokens=False)
print(tokens.input_ids)

config = AutoConfig.from_pretrained(model_name)
token_embeds = nn.Embedding(config.vocab_size, config.hidden_size)
input_embeds = token_embeds(tokens.input_ids)
print(input_embeds.shape)

tensor([[ 2000,  1996,  2614,  1997, 18385, 28248,  5572]])
torch.Size([1, 7, 768])

multihead_attn = MultiHeadAttention(config)
query = key = value = input_embeds
outputs = multihead_attn(query, key, value)
print(outputs.shape)

torch.Size([1, 7, 768])

Encoder

Positional Encoding

• No position information in self-attention.
• Each position has a unique positional vector .

Decoder

Seq2seq Model

AT VS NAT

Efficient Transformers: A Survey

https://arxiv.org/abs/2009.06732

平方复杂度问题：标准Transformer模型中的self-attention机制具有的时间和空间复杂度。因为自注意力机制需要计算序列中所有元素对之间的注意力分数，导致计算量和内存需求随着序列长度的增加而呈平方级增长。这种复杂度在长序列处理任务中尤其成问题，因为它限制了模型可处理的序列长度，并增加了计算资源的需求。

Swin Transformer