【模型讲解】理解与编码LLM中的自注意力、多头注意力、因果注意力和交叉注意力

### 自注意力机制在LLM中的工作原理 自注意力机制（Self-Attention Mechanism）是一种用于捕获序列中不同位置之间的关系的技术，在大语言模型（LLM）中起着至关重要的作用。其核心思想是让模型能够关注输入序列的不同部分，从而更好地理解上下文信息。 #### 1. **自注意力机制的核心概念** 自注意力机制通过计算输入序列中每个单词与其他单词的相关性来决定当前单词的重要性权重。这种机制允许模型动态调整对不同词语的关注程度，而不是固定地对待每一个词[^3]。 具体来说，自注意力机制涉及三个主要组件：Query（查询）、Key（键）和Value（值）。这些向量由输入嵌入经过线性变换得到。对于给定的位置 \(i\)，模型会计算该位置上的 Query 向量与所有其他位置上 Key 向量的相似度得分，并将其转换为概率分布作为加权系数，最终用来聚合 Value 向量形成新的表示。 #### 2. **数学表达形式** 假设我们有一个长度为 \(n\) 的输入序列 \(\{x_1, x_2, ..., x_n\}\)，其中每个 \(x_i\) 都是一个 d 维的向量，则可以通过以下公式定义自注意力： \[ \text{Attention}(Q, K, V) = \text{softmax} \left( \frac{QK^\top}{\sqrt{d_k}} \right)V \] 这里： - \( Q = XW_Q \), \( K = XW_K \), \( V = XW_V \) - \( W_Q, W_K, W_V \) 是可学习参数矩阵 - \( d_k \) 表示 Key 的维度大小 此操作可以被高效并行化实施于现代硬件之上如 GPU 或 TPU 上面[^2]。 #### 3. **多头注意力机制扩展** 为了进一步增强模型的表现力，通常采用多头注意力机制(Multi-head Attention)。这种方法独立运行多个平行的注意子层并将它们的结果连接起来再做一次投影: ```python import torch import torch.nn as nn class MultiHeadAttention(nn.Module): def __init__(self, embed_size, heads): super(MultiHeadAttention, self).__init__() self.embed_size = embed_size self.heads = heads self.head_dim = embed_size // heads assert ( self.head_dim * heads == embed_size ), "Embedding size needs to be divisible by heads" self.values = nn.Linear(self.head_dim, self.head_dim, bias=False) self.keys = nn.Linear(self.head_dim, self.head_dim, bias=False) self.queries = nn.Linear(self.head_dim, self.head_dim, bias=False) self.fc_out = nn.Linear(heads * self.head_dim, embed_size) def forward(self, values, keys, query, mask): N = query.shape[0] value_len, key_len, query_len = values.shape[1], keys.shape[1], query.shape[1] # Split embedding into self.heads pieces values = values.reshape(N, value_len, self.heads, self.head_dim) keys = keys.reshape(N, key_len, self.heads, self.head_dim) queries = query.reshape(N, query_len, self.heads, self.head_dim) energy = torch.einsum("nqhd,nkhd->nhqk", [queries, keys]) if mask is not None: energy = energy.masked_fill(mask == 0, float("-1e20")) attention = torch.softmax(energy / (self.embed_size ** (1/2)), dim=3) out = torch.einsum("nhql,nlhd->nqhd", [attention, values]).reshape( N, query_len, self.heads*self.head_dim ) return self.fc_out(out) ``` 上述代码展示了如何构建一个多头注意力模块，利用 PyTorch 实现了基于张量运算的功能。 #### 4. **优化方法——分组查询注意力(GQA)** 除了标准的多头注意力外，还有更高效的变体比如 GQA(Grouped Query Attention)。它试图平衡速度与质量间的差异，即尝试接近 MHA 质量水平下达到类似于 MQA 的效率增益效果。 ---