Self-Attention的架构
时间: 2025-03-31 11:14:05 浏览: 34
### Self-Attention机制的详细架构及工作原理
#### 架构组成
Self-Attention是一种特殊的注意力机制,其核心在于通过线性变换将输入序列映射到三个不同的向量空间:Query(查询)、Key(键)和Value(值)。具体来说,对于给定的输入序列 \( X \in \mathbb{R}^{n \times d_{\text{model}}} \),其中 \( n \) 是序列长度,\( d_{\text{model}} \) 是嵌入维度,会分别经过三个权重矩阵 \( W_Q, W_K, W_V \) 的乘法操作来获得 Query 向量、Key 向量以及 Value 向量:
\[ Q = XW_Q,\ K = XW_K,\ V = XW_V \]
这里 \( W_Q, W_K, W_V \) 都是可学习参数矩阵[^2]。
#### 工作流程
1. **计算注意力分数**
自注意力机制的核心部分是对齐不同位置的信息。这一步骤通过对 Query 和 Key 进行点积运算完成,并除以缩放因子 \( \sqrt{d_k} \) 来稳定梯度传播,其中 \( d_k \) 表示 Key 的维度。公式如下所示:
\[
\text{Attention}(Q,K,V) = \text{softmax}\left(\frac{QK^\top}{\sqrt{d_k}}\right)V
\]
上述公式的含义是从输入序列的不同位置抽取特征并重新组合它们,从而捕捉全局依赖关系[^3]。
2. **Softmax归一化处理**
得到的原始注意力建议分值会被送入 Softmax 函数中进行概率分布转换,使得每一项都落在区间 [0, 1] 内且总和等于 1。此步骤确保后续加权平均操作具有统计意义。
3. **加权求和生成输出**
经过 softmax 处理后的权重应用于对应的 Value 向量上,最终形成新的表示形式作为该层的输出。这一过程可以视为一种基于内容的选择性聚合方式,突出重要信息的同时抑制无关干扰[^4]。
#### 特殊优化——Flash Attention
为了进一步提升效率,在大规模预训练模型中提出了 Flash Attention 技术。它采用单次遍历(one-pass)策略代替传统多步方法,减少了中间状态保存需求,进而降低了显存消耗与延迟时间。此外,还引入了一些近似算法加速计算速度而不显著影响精度。
```python
import torch
import math
def scaled_dot_product_attention(query, key, value):
dk = query.size(-1)
scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(dk) # 缩放点积
attn_weights = torch.softmax(scores, dim=-1) # 转换为概率分布
output = torch.matmul(attn_weights, value) # 加权求和
return output, attn_weights
```
上述代码片段展示了如何实现标准版本的 Scaled Dot-Product Attention 功能模块。
阅读全文
相关推荐


















