因果注意力

约 3000 字大约 10 分钟

设计思路与核心概念

1. 因果注意力的背景与动机

因果注意力（Causal Attention）是自回归语言模型的核心机制，主要用于GPT系列模型中，解决以下关键问题：

时序约束：在语言生成任务中，当前位置只能看到之前的信息，不能"偷看"未来
自回归生成：确保模型在训练和推理时的一致性，避免信息泄露
因果关系建模：保持序列的时间顺序和因果依赖关系
训练效率：通过掩码机制在训练时并行处理整个序列

2. 核心设计思想

因果注意力的核心思想是通过掩码机制限制注意力只能关注当前位置及之前的位置：

因果掩码：使用上三角掩码屏蔽未来位置的信息
信息流向：确保信息只能从过去流向现在，不能逆向流动
位置约束：第i个位置只能关注位置1到i，不能关注i+1及之后
训练一致性：训练时的并行计算与推理时的逐步生成保持一致

3. 数学公式

因果注意力的完整公式

对于输入序列 $X \in \mathbb{R}^{n \times d_{in}}$ ，因果注意力的计算过程为：

线性变换生成QKV：
$Q = XW_Q, \quad K = XW_K, \quad V = XW_V$
注意力分数计算：
$S = \frac{QK^T}{\sqrt{d_k}}$
因果掩码应用：
$S_{masked}[i,j] = \begin{cases} S[i,j] & \text{if } j \leq i \\ -\infty & \text{if } j > i \end{cases}$
注意力权重计算：
$A = \text{softmax}(S_{masked})$
上下文向量计算：
$\text{CausalAttention}(Q,K,V) = AV$

其中掩码矩阵 $M$ 定义为：

M[i,j] = \begin{cases} 0 & \text{if } j \leq i \\ 1 & \text{if } j > i \end{cases}

执行流程

1. 整体执行流程图

2. 详细计算流程图

计算步骤详解

步骤	操作	PyTorch代码	数学表达式	说明
1	输入序列	`x`	$X \in \mathbb{R}^{b \times n \times d_{in}}$	批次大小b，序列长度n
2	生成Query	`self.W_query(x)`	$Q = XW_Q$	查询矩阵(b, n, d_out)
3	生成Key	`self.W_key(x)`	$K = XW_K$	键矩阵(b, n, d_out)
4	生成Value	`self.W_value(x)`	$V = XW_V$	值矩阵(b, n, d_out)
5	计算分数	`queries @ keys.transpose(1, 2)`	$S = QK^T$	注意力分数(b, n, n)
6	应用掩码	`attn_scores.masked_fill_(mask, -torch.inf)`	$S_{masked}$	屏蔽未来位置
7	缩放归一化	`torch.softmax(scores / sqrt(d_k), dim=-1)`	$A = \text{softmax}(\frac{S_{masked}}{\sqrt{d_k}})$	概率分布
8	Dropout	`self.dropout(attn_weights)`	$A_{drop}$	防止过拟合
9	加权求和	`attn_weights @ values`	$C = A_{drop}V$	最终上下文(b, n, d_out)

完整代码实现

因果注意力.py
# 导入PyTorch库
import torch
import torch.nn as nn


class CausalAttention(nn.Module):
    """实现因果自注意力机制的模块"""
    
    def __init__(self, d_in, d_out, context_length, dropout, qkv_bias=False):
        """
        参数说明：
            d_in: 输入特征维度
            d_out: 输出特征维度
            context_length: 上下文长度（最大序列长度）
            dropout: dropout比率
            qkv_bias: 是否在Q/K/V线性变换中添加偏置项
        """
        super().__init__()
        self.d_out = d_out
        
        # 定义Q/K/V的线性变换层
        self.W_query = nn.Linear(d_in, d_out, bias=qkv_bias)  # 查询（Query）变换
        self.W_key = nn.Linear(d_in, d_out, bias=qkv_bias)    # 键（Key）变换
        self.W_value = nn.Linear(d_in, d_out, bias=qkv_bias)  # 值（Value）变换
        
        self.dropout = nn.Dropout(dropout)  # dropout层
        
        # 注册因果掩码（不可训练的缓冲区）
        # 创建上三角矩阵（主对角线以上为1，其余为0），用于屏蔽未来信息
        self.register_buffer(
            'mask',
            torch.triu(torch.ones(context_length, context_length), diagonal=1)
        )

    def forward(self, x):
        """
        前向传播过程
        输入形状: (batch_size, num_tokens, d_in)
        输出形状: (batch_size, num_tokens, d_out)
        """
        b, num_tokens, d_in = x.shape  # 获取输入张量的形状信息
        
        # 生成键（K）、查询（Q）、值（V）
        keys = self.W_key(x)      # 形状: (b, num_tokens, d_out)
        queries = self.W_query(x) # 形状: (b, num_tokens, d_out)
        values = self.W_value(x)  # 形状: (b, num_tokens, d_out)
        
        # 计算注意力分数（缩放点积注意力）
        attn_scores = queries @ keys.transpose(1, 2)  # (b, num_tokens, num_tokens)
        
        # 应用因果掩码（屏蔽未来位置的注意力）
        attn_scores.masked_fill_(
            self.mask.bool()[:num_tokens, :num_tokens],  # 动态调整掩码尺寸
            -torch.inf   # 用负无穷填充被屏蔽的位置
        )
        
        # 计算注意力权重（应用softmax和dropout）
        attn_weights = torch.softmax(
            attn_scores / keys.shape[-1] ** 0.5,  # 缩放因子sqrt(d_k)
            dim=-1
        )
        attn_weights = self.dropout(attn_weights)
        
        # 计算上下文向量（加权和）
        context_vec = attn_weights @ values  # (b, num_tokens, d_out)
        
        return context_vec


# ===================== 测试代码 =====================
# 创建测试输入数据（模拟6个词的嵌入向量）
# 每个词用3维向量表示，包含两个相同的样本（batch_size=2）
inputs = torch.tensor(
    [[0.43, 0.15, 0.89],  # 第1个词 "Your" 的嵌入向量
     [0.55, 0.87, 0.66],  # 第2个词 "journey" 的嵌入向量
     [0.57, 0.85, 0.64],  # 第3个词 "starts" 的嵌入向量
     [0.22, 0.58, 0.33],  # 第4个词 "with" 的嵌入向量
     [0.77, 0.25, 0.10],  # 第5个词 "one" 的嵌入向量
     [0.05, 0.80, 0.55]]  # 第6个词 "step" 的嵌入向量
)

# 创建批次数据（将输入数据堆叠两次，模拟batch_size=2的情况）
batch = torch.stack((inputs, inputs), dim=0)
print("输入数据形状:", batch.shape)  # 预期输出: torch.Size([2, 6, 3])

# 初始化模型参数
d_in = inputs.shape[-1]    # 输入维度 = 3
d_out = 2                  # 输出维度 = 2
context_length = batch.shape[1]  # 上下文长度 = 6

# 固定随机种子（保证可重复性）
torch.manual_seed(123)

# 初始化因果自注意力模块（关闭dropout）
ca = CausalAttention(d_in, d_out, context_length, 0.0)

# 前向传播
context_vecs = ca(batch)

# 输出结果形状
print("上下文向量形状:", context_vecs)  # 预期输出: torch.Size([2, 6, 2])

代码详细解析

1. 类设计架构

CausalAttention类结构

class CausalAttention(nn.Module):
    def __init__(self, d_in, d_out, context_length, dropout, qkv_bias=False):
        # 参数初始化和掩码创建
    def forward(self, x):
        # 前向传播逻辑

设计要点：

继承自nn.Module，符合PyTorch模块化设计
支持批处理输入，适用于实际训练场景
内置因果掩码，确保时序约束
集成Dropout正则化，提高泛化能力

2. 核心参数分析

初始化参数详解

参数	类型	作用	默认值
`d_in`	int	输入特征维度	必需
`d_out`	int	输出特征维度	必需
`context_length`	int	最大序列长度	必需
`dropout`	float	Dropout比率	必需
`qkv_bias`	bool	是否使用偏置	False

关键组件设计

# 线性变换层
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)

# 因果掩码（不可训练参数）
self.register_buffer('mask', torch.triu(torch.ones(...), diagonal=1))

3. 因果掩码机制详解

掩码创建原理

# 创建上三角矩阵
mask = torch.triu(torch.ones(context_length, context_length), diagonal=1)

掩码矩阵示例（context_length=4）：

[[0, 1, 1, 1],
 [0, 0, 1, 1],
 [0, 0, 0, 1],
 [0, 0, 0, 0]]

掩码含义：

0: 允许注意力（当前及过去位置）
1: 屏蔽注意力（未来位置）
对角线: diagonal=1 表示主对角线上方开始为1

掩码应用机制

attn_scores.masked_fill_(
    self.mask.bool()[:num_tokens, :num_tokens],
    -torch.inf
)

关键技术点：

动态调整: [:num_tokens, :num_tokens] 适应不同序列长度
负无穷填充: -torch.inf 确保softmax后权重为0
就地操作: masked_fill_ 提高内存效率

4. 前向传播详解

步骤1：QKV生成

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

批处理支持：

输入形状: (batch_size, num_tokens, d_in)
输出形状: (batch_size, num_tokens, d_out)
并行处理: 批次内所有样本同时处理

步骤2：注意力分数计算

attn_scores = queries @ keys.transpose(1, 2)

维度变换：

queries: (b, n, d_out)
keys.transpose(1, 2): (b, d_out, n)
attn_scores: (b, n, n)

步骤3：因果掩码应用

attn_scores.masked_fill_(
    self.mask.bool()[:num_tokens, :num_tokens],
    -torch.inf
)

掩码效果可视化：

原始分数矩阵:        掩码后矩阵:
[s11, s12, s13]  →  [s11, -∞,  -∞ ]
[s21, s22, s23]  →  [s21, s22, -∞ ]
[s31, s32, s33]  →  [s31, s32, s33]

步骤4：权重计算和正则化

attn_weights = torch.softmax(
    attn_scores / keys.shape[-1] ** 0.5,
    dim=-1
)
attn_weights = self.dropout(attn_weights)

处理流程：

缩放: 除以 $\sqrt{d_k}$ 防止梯度消失
Softmax: 转换为概率分布
Dropout: 随机置零部分权重

5. 与标准注意力的对比

功能对比表

特性	标准自注意力	因果注意力
信息访问	全序列可见	仅过去可见
掩码机制	无	上三角掩码
应用场景	编码器	解码器
训练方式	双向上下文	单向上下文
计算复杂度	O(n²)	O(n²)
内存开销	标准	额外掩码存储

注意力模式对比

# 标准注意力权重矩阵（全连接）
[[w11, w12, w13, w14],
 [w21, w22, w23, w24],
 [w31, w32, w33, w34],
 [w41, w42, w43, w44]]

# 因果注意力权重矩阵（下三角）
[[w11,  0,   0,   0 ],
 [w21, w22,  0,   0 ],
 [w31, w32, w33,  0 ],
 [w41, w42, w43, w44]]

6. 实验配置分析

输入数据设计

inputs = torch.tensor([
    [0.43, 0.15, 0.89],  # "Your"
    [0.55, 0.87, 0.66],  # "journey"
    [0.57, 0.85, 0.64],  # "starts"
    [0.22, 0.58, 0.33],  # "with"
    [0.77, 0.25, 0.10],  # "one"
    [0.05, 0.80, 0.55]   # "step"
])

# 创建批次数据
batch = torch.stack((inputs, inputs), dim=0)

数据特点：

序列长度: 6个词
嵌入维度: 3维向量
批次大小: 2个样本
数据复制: 便于观察批处理效果

模型配置

d_in = 3              # 输入维度
d_out = 2             # 输出维度
context_length = 6    # 上下文长度
dropout = 0.0         # 关闭dropout（测试用）

7. 输出结果解读

预期输出特征

形状: (2, 6, 2) - 批次大小×序列长度×输出维度
因果性: 每个位置只包含当前及之前位置的信息
一致性: 相同输入产生相同输出（固定随机种子）

高级特性与扩展

1. 掩码优化技术

内存优化版本

def create_causal_mask(seq_len, device):
    """创建内存优化的因果掩码"""
    mask = torch.triu(torch.ones(seq_len, seq_len, device=device), diagonal=1)
    return mask.bool()

# 使用示例
mask = create_causal_mask(seq_len, device=x.device)
attn_scores.masked_fill_(mask, float('-inf'))

稀疏掩码支持

def create_sparse_causal_mask(seq_len, window_size):
    """创建局部窗口的因果掩码"""
    mask = torch.triu(torch.ones(seq_len, seq_len), diagonal=1)
    # 添加局部窗口限制
    for i in range(seq_len):
        mask[i, :max(0, i-window_size)] = 1
    return mask.bool()

2. 性能优化策略

Flash Attention集成

try:
    from flash_attn import flash_attn_func

    def flash_causal_attention(Q, K, V, causal=True):
        """使用Flash Attention优化的因果注意力"""
        return flash_attn_func(Q, K, V, causal=causal)

except ImportError:
    print("Flash Attention not available, using standard implementation")

梯度检查点

from torch.utils.checkpoint import checkpoint

class CheckpointCausalAttention(CausalAttention):
    def forward(self, x):
        return checkpoint(super().forward, x)

3. 可视化工具

因果掩码可视化

import matplotlib.pyplot as plt
import seaborn as sns

def visualize_causal_mask(context_length):
    """可视化因果掩码模式"""
    mask = torch.triu(torch.ones(context_length, context_length), diagonal=1)

    plt.figure(figsize=(8, 6))
    sns.heatmap(mask.numpy(),
                annot=True,
                cmap='Reds',
                cbar_kws={'label': 'Masked (1) / Visible (0)'})
    plt.title('Causal Attention Mask')
    plt.xlabel('Key Position')
    plt.ylabel('Query Position')
    plt.show()

# 使用示例
visualize_causal_mask(6)

注意力权重可视化

def visualize_causal_attention(model, inputs, words):
    """可视化因果注意力权重"""
    model.eval()
    with torch.no_grad():
        # 获取注意力权重
        x = inputs.unsqueeze(0)  # 添加批次维度
        b, num_tokens, d_in = x.shape

        Q = model.W_query(x)
        K = model.W_key(x)

        attn_scores = Q @ K.transpose(1, 2)
        attn_scores.masked_fill_(
            model.mask.bool()[:num_tokens, :num_tokens],
            float('-inf')
        )
        attn_weights = torch.softmax(attn_scores / (K.shape[-1] ** 0.5), dim=-1)

    plt.figure(figsize=(10, 8))
    sns.heatmap(attn_weights[0].numpy(),
                xticklabels=words,
                yticklabels=words,
                annot=True,
                cmap='Blues',
                fmt='.3f')
    plt.title('Causal Attention Weights')
    plt.xlabel('Key (Attended to)')
    plt.ylabel('Query (Attending from)')
    plt.show()