简化的自注意力

约 2643 字大约 9 分钟

设计思路与核心概念

1. 自注意力机制的背景与动机

自注意力（Self-Attention）机制是Transformer架构的核心组件，由Vaswani等人在2017年的论文《Attention Is All You Need》中提出，主要解决以下问题：

序列建模的并行化：传统RNN/LSTM无法并行处理，自注意力允许同时计算所有位置的关系
长距离依赖捕获：直接建模任意两个位置之间的关系，无需通过中间状态传递
动态权重分配：根据输入内容动态决定每个位置的重要性，而非固定的权重
上下文理解：每个词都能"看到"整个序列，获得丰富的上下文信息

2. 核心设计思想

自注意力的核心思想是让序列中的每个元素都能关注到序列中的所有元素：

自我关注：序列内部元素之间的相互关注，不依赖外部信息
全局视野：每个位置都能直接访问所有其他位置的信息
权重动态性：注意力权重根据输入内容动态计算，体现内容的相关性
信息聚合：通过加权求和的方式聚合全局信息

3. 数学公式

简化版自注意力公式

对于输入序列 $X = [x_1, x_2, ..., x_n]$ ，简化的自注意力计算过程为：

注意力分数计算：
$S_{ij} = x_i \cdot x_j^T$
注意力权重归一化：
$A_{ij} = \frac{\exp(S_{ij})}{\sum_{k=1}^{n} \exp(S_{ik})} = \text{softmax}(S_i)$
上下文向量计算：
$c_i = \sum_{j=1}^{n} A_{ij} \cdot x_j$

其中：

$S_{ij}$ 表示位置i对位置j的注意力分数
$A_{ij}$ 表示归一化后的注意力权重
$c_i$ 表示位置i的上下文向量

执行流程

1. 整体执行流程图

2. 详细计算流程图

计算步骤详解

步骤	操作	PyTorch代码	数学表达式	说明
1	输入序列	`inputs`	$X \in \mathbb{R}^{n \times d}$	n个词，每个词d维向量
2	选择查询	`query = inputs[1]`	$q = x_2$	选择第2个词作为查询
3	计算分数	`torch.dot(x_i, query)`	$s_i = x_i \cdot q$	点积计算相似度
4	Softmax归一化	`torch.softmax(scores, dim=0)`	$a_i = \frac{\exp(s_i)}{\sum_j \exp(s_j)}$	转换为概率分布
5	矩阵化计算	`inputs @ inputs.T`	$S = XX^T$	并行计算所有分数
6	批量Softmax	`torch.softmax(scores, dim=-1)`	$A = \text{softmax}(S)$	每行独立归一化
7	上下文聚合	`attn_weights @ inputs`	$C = AX$	加权求和得到上下文

完整代码实现

简化的自注意力.py
# 导入PyTorch深度学习框架
import torch

# 定义输入矩阵（包含6个词的3维向量表示）
# 每个向量代表一个词的嵌入表示，例如第一个向量[0.43,0.15,0.89]对应"Your"的编码
inputs = torch.tensor(
    [[0.43, 0.15, 0.89], # Your (x^1)
    [0.55, 0.87, 0.66], # journey (x^2)
    [0.57, 0.85, 0.64], # starts (x^3)
    [0.22, 0.58, 0.33], # with (x^4)
    [0.77, 0.25, 0.10], # one (x^5)
    [0.05, 0.80, 0.55]] # step (x^6)
)

# 选取第2个词向量作为查询向量（索引从0开始）
query = inputs[1]

# 初始化注意力分数容器（存储每个词与查询向量的相关度）
attn_scores_2 = torch.empty(inputs.shape[0])

# 遍历所有词向量计算注意力分数
for i, x_i in enumerate(inputs):
    # 通过点积计算相似度（查询向量与每个词向量的关联强度）
    attn_scores_2[i] = torch.dot(x_i, query)
print(attn_scores_2)

# 应用Softmax函数将分数转换为概率分布（总和为1）
attn_weights_2 = torch.softmax(attn_scores_2, dim=0)
print("Attention weights:", attn_weights_2)
print("Sum:", attn_weights_2.sum())  # 验证概率分布性质

# 通过矩阵乘法并行计算所有词的注意力分数（输入矩阵与其转置相乘）
attn_scores = inputs @ inputs.T  # 等价于torch.matmul(inputs, inputs.T)
print(attn_scores)

# 对每行应用Softmax得到注意力权重矩阵（每个词对所有词的关注度）
attn_weights = torch.softmax(attn_scores, dim=-1)
print(attn_weights)
print("All row sums:", attn_weights.sum(dim=-1))  # 每行概率和应为1

# 计算最终上下文向量（每个词的加权组合表示）
all_context_vecs = attn_weights @ inputs  # 权重矩阵与输入矩阵相乘
print(all_context_vecs)

代码详细解析

1. 数据结构设计

输入表示

inputs = torch.tensor([
    [0.43, 0.15, 0.89], # Your (x^1)
    [0.55, 0.87, 0.66], # journey (x^2)
    [0.57, 0.85, 0.64], # starts (x^3)
    [0.22, 0.58, 0.33], # with (x^4)
    [0.77, 0.25, 0.10], # one (x^5)
    [0.05, 0.80, 0.55]  # step (x^6)
])

设计要点：

形状: (6, 3) - 6个词，每个词3维向量
含义: 每行代表一个词的嵌入向量
数值: 模拟真实的词嵌入，范围在[0,1]之间

2. 核心算法实现

单个查询的注意力计算

# 方法1：逐步计算（教学目的）
query = inputs[1]  # 选择"journey"作为查询
attn_scores_2 = torch.empty(inputs.shape[0])

for i, x_i in enumerate(inputs):
    attn_scores_2[i] = torch.dot(x_i, query)

关键点分析：

查询选择: 选择第2个词（索引1）作为查询向量
点积计算: torch.dot(x_i, query) 计算相似度
循环遍历: 逐个计算每个词与查询的关系

矩阵化并行计算

# 方法2：矩阵化计算（高效实现）
attn_scores = inputs @ inputs.T

优势分析：

并行化: 一次计算所有词对之间的关系
效率提升: 利用GPU的矩阵运算优化
代码简洁: 一行代码替代嵌套循环

3. 注意力权重归一化

Softmax函数的作用

attn_weights = torch.softmax(attn_scores, dim=-1)

数学原理：

归一化: 确保每行权重和为1
概率解释: 将分数转换为概率分布
平滑性: 保持相对大小关系的同时避免硬选择

维度说明

参数	含义	作用
`dim=-1`	沿最后一个维度	对每行独立进行Softmax
`dim=0`	沿第一个维度	对每列进行Softmax

4. 上下文向量计算

加权求和机制

all_context_vecs = attn_weights @ inputs

计算过程：

权重矩阵: attn_weights 形状为 (6, 6)
输入矩阵: inputs 形状为 (6, 3)
输出矩阵: all_context_vecs 形状为 (6, 3)

物理意义：

每个输出向量是所有输入向量的加权平均
权重反映了词与词之间的相关性
结果包含了全局上下文信息

实验结果分析

1. 注意力分数解读

单个查询的分数

# 示例输出（具体数值会因输入而异）
tensor([1.0179, 1.8524, 1.8111, 0.7625, 0.9106, 0.8792])

分析：

最高分: 位置1（"journey"自己）和位置2（"starts"）
相关性: 分数越高表示与查询词越相关
对称性: 由于使用点积，关系是对称的

注意力权重分布

# 示例输出
Attention weights: tensor([0.1385, 0.3190, 0.3059, 0.1078, 0.1244, 0.1044])
Sum: tensor(1.0000)

特点：

概率性: 所有权重和为1
集中性: 权重主要集中在相关词上
平滑性: Softmax确保了平滑的权重分布

2. 全局注意力矩阵

矩阵结构

# 注意力权重矩阵 (6x6)
# 每行表示一个词对所有词的注意力分布
attn_weights = [
    [w11, w12, w13, w14, w15, w16],  # "Your"的注意力分布
    [w21, w22, w23, w24, w25, w26],  # "journey"的注意力分布
    ...
]

解读规则：

对角线: 通常较高，表示词对自身的关注
相似词: 语义相似的词之间权重较高
行归一化: 每行和为1，满足概率分布要求

3. 上下文向量特性

信息融合

全局信息: 每个上下文向量包含了整个序列的信息
权重调节: 重要的词贡献更多信息
维度保持: 输出维度与输入维度相同

与完整自注意力的对比

简化版 vs 完整版

特性	简化版	完整版
Query/Key/Value	直接使用输入	通过线性变换生成
多头机制	单头	多头并行
缩放因子	无	$\frac{1}{\sqrt{d_k}}$
位置编码	无	有
掩码机制	无	支持因果掩码

简化版的优势

概念清晰: 直观展示自注意力的核心思想
计算简单: 易于理解和实现
教学友好: 适合初学者理解机制

简化版的局限

表达能力: 缺乏Query/Key/Value的灵活性
扩展性: 无法处理复杂的注意力模式
实用性: 在实际应用中效果有限

扩展实现

1. 带缩放的自注意力

import math

def scaled_dot_product_attention(inputs):
    """带缩放因子的自注意力"""
    d_k = inputs.size(-1)  # 特征维度

    # 计算注意力分数并缩放
    attn_scores = inputs @ inputs.T / math.sqrt(d_k)

    # Softmax归一化
    attn_weights = torch.softmax(attn_scores, dim=-1)

    # 计算上下文向量
    context_vecs = attn_weights @ inputs

    return context_vecs, attn_weights

2. 可视化注意力权重

import matplotlib.pyplot as plt
import seaborn as sns

def visualize_attention(attn_weights, words):
    """可视化注意力权重矩阵"""
    plt.figure(figsize=(8, 6))
    sns.heatmap(attn_weights.numpy(),
                xticklabels=words,
                yticklabels=words,
                annot=True,
                cmap='Blues',
                fmt='.3f')
    plt.title('Self-Attention Weights')
    plt.xlabel('Key (Attended to)')
    plt.ylabel('Query (Attending from)')
    plt.show()

# 使用示例
words = ["Your", "journey", "starts", "with", "one", "step"]
visualize_attention(attn_weights, words)

3. 批处理版本

def batch_self_attention(batch_inputs):
    """支持批处理的自注意力"""
    # batch_inputs: (batch_size, seq_len, d_model)

    # 计算注意力分数
    attn_scores = torch.bmm(batch_inputs, batch_inputs.transpose(1, 2))

    # Softmax归一化
    attn_weights = torch.softmax(attn_scores, dim=-1)

    # 计算上下文向量
    context_vecs = torch.bmm(attn_weights, batch_inputs)

    return context_vecs, attn_weights