大模型入门 ch 03：注意力机制

本文是github上的大模型教程LLMs-from-scratch的学习笔记，教程地址：教程链接

Chapter 3： Attention Mechanism

本文首先从固定参数的注意力机制说起，然后拓展到可以训练的注意力机制，然后加入掩码mask，最后拓展到多头注意力机制。

1. 注意力机制

一个句子中的每一个token，都会受到其他token的影响（这里先不考虑忽略未来的单词，掩码的问题后面再说），注意力机制可以让一个token收到其他token的影响，生成一个最终我们想要的embedding。即每个token有一个原始的embedding，通过注意力机制后，得到了一个新的embedding，这个embedding是结合了上下文语义得到的。

举个简单的例子，我们直接使用tokens的embedding之间两两点乘，得到互相之间的点乘结果，然后将点乘结果归一化，得到embeddings之间的注意力得分。

归一化一般使用softmax函数，通过取指数，除以求和得到归一化结果
torch.exp(x) / torch.exp(x).sum(dim=0)

得到token之间的相关权重后，我们就可以加权求和，得到每一个token的最终embedding。

2. 可以训练的注意力头

在上面的例子中，我们直接使用token对应的embedding来计算相关系数，以及最终的加权求和，这显然是不合理的，如果这样的话，那么我们只能训练token对应的词嵌入来学习模型，或者是一些全连接层，因此我们需要引入新的矩阵，来学习到更多的参数，这就是transformer的QKV矩阵。
QKV都是对原始的embedding做线性变换，得到新的向量，但是模型就可以通过训练QKV，学习海量知识。

在这里插入图片描述

QKV的维度不固定，可以与原始嵌入相同，也可以不同。总之，通过QKV三个矩阵，我们将原始token的embedding转换成了3个新的向量。

Query vector: $q^{(i)} = W_q \,x^{(i)}$
Key vector: $k^{(i)} = W_k \,x^{(i)}$
Value vector: $v^{(i)} = W_v \,x^{(i)}$

可以使用矩阵乘法实现：

keys = inputs @ W_key 
values = inputs @ W_value

然后我们计算K和Q之间的点积，作为两两token之间的关联度。为什么要用两个不一样的矩阵，我的猜测是，如果使用的是一个矩阵计算相似度，那么关于对角线对称的元素就会完全相同，但是使用两个不同的矩阵计算，就不会存在这样的情况，可以学习到的内容更多。

我们使用K和Q的点积得到了两两之间的注意力得分，同样使用softmax进行归一化，得到最终的注意力权重。

注意到没有直接对注意力得分softmax，而是除以维度的方根后再softmax，这是因为在计算注意力权重时，如果直接将Query和Key的点积结果用于softmax函数，当Key的维度较高时，点积的结果会变得非常大。这可能导致softmax函数在梯度下降过程中学习困难，因为大的数值会使softmax的梯度变得非常小（接近于0），这在数值稳定性上是一个问题，称为“梯度消失”。

最后一步，不再使用原始的embedding加权，我们使用V矩阵变换后的向量进行加权求和，得到结果向量。

代码如下

class SelfAttention_v2(nn.Module):

    def __init__(self, d_in, d_out, qkv_bias=False):
        super().__init__()
        self.W_query = nn.Linear(d_in, d_out, bias=qkv_bias)
        self.W_key   = nn.Linear(d_in, d_out, bias=qkv_bias)
        self.W_value = nn.Linear(d_in, d_out, bias=qkv_bias)

    def forward(self, x):
        keys = self.W_key(x)
        queries = self.W_query(x)
        values = self.W_value(x)
        
        attn_scores = queries @ keys.T
        attn_weights = torch.softmax(attn_scores / keys.shape[-1]**0.5, dim=-1)

        context_vec = attn_weights @ values
        return context_vec

3. 隐藏未来的单词

对语言任务来说，在训练模型的时候不能使用未来的文本来预测之前的文本。因此我们需要屏蔽未见文本对先前文本的影响。在我们计算得到注意力权重后，我们人为地将上三角矩阵的权重置为0。
有一种naive的方法，就是将上注意力权重都置为0后，重新对剩下的元素归一化。但是我们要介绍的是一般使用的方法：

我们在计算出注意力得分后，对右上角元素都赋值为负无穷大，负无穷大在经过softmax后就变为了0。

mask = torch.triu(torch.ones(context_length, context_length), diagonal=1)
masked = attn_scores.masked_fill(mask.bool(), -torch.inf)
attn_weights = torch.softmax(masked / keys.shape[-1]**0.5, dim=-1)

最后为了防止过拟合，一般会使用dropout，对注意力权重矩阵进行随机丢弃，加强模型泛化性能。

总结以上的所有内容，我们现在就可以写出一个单头的注意力机制了，并且加入了对batch输入的处理：

class CausalAttention(nn.Module):

    def __init__(self, d_in, d_out, context_length,
                 dropout, qkv_bias=False):
        super().__init__()
        self.d_out = d_out
        self.W_query = nn.Linear(d_in, d_out, bias=qkv_bias)
        self.W_key   = nn.Linear(d_in, d_out, bias=qkv_bias)
        self.W_value = nn.Linear(d_in, d_out, bias=qkv_bias)
        self.dropout = nn.Dropout(dropout) # New
        self.register_buffer('mask', torch.triu(torch.ones(context_length, context_length), diagonal=1)) # New

    def forward(self, x):
        b, num_tokens, d_in = x.shape # New batch dimension b
        keys = self.W_key(x)
        queries = self.W_query(x)
        values = self.W_value(x)

        attn_scores = queries @ keys.transpose(1, 2) # Changed transpose
        attn_scores.masked_fill_(  # New, _ ops are in-place
            self.mask.bool()[:num_tokens, :num_tokens], -torch.inf)  # `:num_tokens` to account for cases where the number of tokens in the batch is smaller than the supported context_size
        attn_weights = torch.softmax(
            attn_scores / keys.shape[-1]**0.5, dim=-1
        )
        attn_weights = self.dropout(attn_weights) # New

        context_vec = attn_weights @ values
        return context_vec

4. 多头注意力机制

我们已经实现了单个头的注意力机制，那么要实现多个头，就是使用多个不同的注意力头，各自对输入进行处理，然后将各自得到的输出$z_i$拼接起来，非常显而易见，我们有第一个最直白的写法：

class MultiHeadAttentionWrapper(nn.Module):

    def __init__(self, d_in, d_out, context_length, dropout, num_heads, qkv_bias=False):
        super().__init__()
        self.heads = nn.ModuleList(
            [CausalAttention(d_in, d_out, context_length, dropout, qkv_bias) 
             for _ in range(num_heads)]
        )

    def forward(self, x):
        return torch.cat([head(x) for head in self.heads], dim=-1)

这是一种最简单直白的写法，直接声明num_heads个注意力单元，然后在前向传播的时候，依次调用这num_heads个注意力头，然后将输出拼接起来。(dim=-1代表最后一维拼接）

问题是，这样的话需要循环num_heads次得到结果，并且需要声明num_heads个注意力头，相信熟悉线性代数的朋友已经想到了，可以通过曾广矩阵来拓展注意力头。比如单个的注意力头是(d_in, d_out)，那么有n个头的注意力机制就是(d_in, n*d_out)。
假设输入是(tokens, d_in)，那么(tokens, d_in) @ (d_in, n*d_out) --> (tokens, n*d_out)，输出的结果完美得到了n个头对应的输出，我们只需要按照每d_out列拆开，得到n个(tokens, d_out)的矩阵，就能还原出n个头对应的结果，进行后续的attention score计算。这样写起来虽然麻烦一些，但是效率更高。

class MultiHeadAttention(nn.Module):
    def __init__(self, d_in, d_out, context_length, dropout, num_heads, qkv_bias=False):
        super().__init__()
        assert (d_out % num_heads == 0), \
            "d_out must be divisible by num_heads"

        self.d_out = d_out
        self.num_heads = num_heads
        self.head_dim = d_out // num_heads # Reduce the projection dim to match desired output dim

        self.W_query = nn.Linear(d_in, d_out, bias=qkv_bias)
        self.W_key = nn.Linear(d_in, d_out, bias=qkv_bias)
        self.W_value = nn.Linear(d_in, d_out, bias=qkv_bias)
        self.out_proj = nn.Linear(d_out, d_out)  # Linear layer to combine head outputs
        self.dropout = nn.Dropout(dropout)
        self.register_buffer(
            "mask",
            torch.triu(torch.ones(context_length, context_length),
                       diagonal=1)
        )

    def forward(self, x):
        b, num_tokens, d_in = x.shape

        keys = self.W_key(x) # Shape: (b, num_tokens, d_out)
        queries = self.W_query(x)
        values = self.W_value(x)

        # We implicitly split the matrix by adding a `num_heads` dimension
        # Unroll last dim: (b, num_tokens, d_out) -> (b, num_tokens, num_heads, head_dim)
        keys = keys.view(b, num_tokens, self.num_heads, self.head_dim) 
        values = values.view(b, num_tokens, self.num_heads, self.head_dim)
        queries = queries.view(b, num_tokens, self.num_heads, self.head_dim)

        # Transpose: (b, num_tokens, num_heads, head_dim) -> (b, num_heads, num_tokens, head_dim)
        keys = keys.transpose(1, 2)
        queries = queries.transpose(1, 2)
        values = values.transpose(1, 2)

        # Compute scaled dot-product attention (aka self-attention) with a causal mask
        attn_scores = queries @ keys.transpose(2, 3)  # Dot product for each head

        # Original mask truncated to the number of tokens and converted to boolean
        mask_bool = self.mask.bool()[:num_tokens, :num_tokens]

        # Use the mask to fill attention scores
        attn_scores.masked_fill_(mask_bool, -torch.inf)
        
        attn_weights = torch.softmax(attn_scores / keys.shape[-1]**0.5, dim=-1)
        attn_weights = self.dropout(attn_weights)

        # Shape: (b, num_tokens, num_heads, head_dim)
        context_vec = (attn_weights @ values).transpose(1, 2) 
        
        # Combine heads, where self.d_out = self.num_heads * self.head_dim
        context_vec = context_vec.contiguous().view(b, num_tokens, self.d_out)
        context_vec = self.out_proj(context_vec) # optional projection

        return context_vec