深入解析Transformer架构：从Attention机制到上下文窗口优化

1次阅读

共计 2297 个字符，预计需要花费 6 分钟才能阅读完成。

Transformer 架构自 2017 年由 Google 提出以来，彻底改变了自然语言处理（NLP）领域的格局。与传统 RNN 和 LSTM 相比，Transformer 通过并行计算和全局依赖捕捉能力，显著提升了模型训练效率和性能。如今，从 BERT 到 GPT 系列模型，Transformer 已成为大语言模型的基石架构。

QKV 矩阵生成 ：
输入序列通过三个独立的线性变换层，分别生成 Query、Key 和 Value 矩阵。这三个矩阵的维度通常为（序列长度×头维度），其中头维度 = 嵌入维度 / 头数。
注意力分数计算 ：
通过矩阵乘法计算 Query 和 Key 的点积，然后除以√d_k（键向量维度的平方根）进行缩放，最后应用 softmax 函数得到注意力权重。公式表达为：
```
Attention(Q,K,V)=softmax(QK^T/√d_k)V
```
多头注意力机制 ：
将上述过程并行执行多次（即多个注意力头），每个头学习不同的注意力模式，最后将各头的输出拼接并通过线性变换融合。

Tokenization：
文本被分割为子词或单词级别的 token，常见方案包括 Byte Pair Encoding (BPE) 和 WordPiece。例如 ”unhappiness” 可能被分解为 [“un”, “happiness”]。
位置编码 ：
由于 Transformer 没有循环结构，必须显式注入位置信息。原始论文使用正弦函数生成固定位置编码，公式为：
```
PE(pos,2i)=sin(pos/10000^(2i/d_model))
PE(pos,2i+1)=cos(pos/10000^(2i/d_model))
```
现代实现更常使用可学习的位置嵌入。

固定窗口的限制 ：
传统 Transformer 的注意力计算复杂度为 O(n²)，当序列长度超过预设窗口（如 GPT- 3 的 2048）时，需要截断或分块处理，导致长距离依赖丢失。
优化方案 ：
稀疏注意力（如 Longformer 的滑动窗口）
内存压缩（如 Reformer 的 LSH 注意力）
分块处理配合状态缓存（如 Transformer-XL 的段级递归）

import torch
import torch.nn as nn
import math

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model=512, n_heads=8):
        super().__init__()
        assert d_model % n_heads == 0
        self.d_k = d_model // n_heads
        self.n_heads = n_heads

        # 线性变换层
        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        self.W_o = nn.Linear(d_model, d_model)

    def forward(self, x, mask=None):
        batch_size = x.size(0)

        # 生成 QKV [batch_size, seq_len, d_model]
        Q = self.W_q(x)
        K = self.W_k(x)
        V = self.W_v(x)

        # 分头 [batch_size, n_heads, seq_len, d_k]
        Q = Q.view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
        K = K.view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
        V = V.view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)

        # 计算注意力分数 [batch_size, n_heads, seq_len, seq_len]
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)
        attention = torch.softmax(scores, dim=-1)

        # 加权求和 [batch_size, n_heads, seq_len, d_k]
        context = torch.matmul(attention, V)

        # 合并多头 [batch_size, seq_len, d_model]
        context = context.transpose(1, 2).contiguous().view(batch_size, -1, self.n_heads * self.d_k)

        return self.W_o(context)