从零实现大语言模型：深入理解因果注意力机制

从零实现大语言模型：深入理解因果注意力机制引言在自然语言处理领域，Transformer架构已成为构建大语言模型的基础。其中，注意力机制是Transformer的核心组件。本文将重点探讨一种特殊的注意力机制——因果注意力（Causal Attention），这是构建GPT等自回归语言模型的关键技术。什么是因果注意力？因果注意力，也称为遮蔽注意力（Masked Attention），是自注意...

乔媚倩June

721人浏览 · 2025-06-19 09:08:51

乔媚倩June · 2025-06-19 09:08:51 发布

从零实现大语言模型：深入理解因果注意力机制

引言

在自然语言处理领域，Transformer架构已成为构建大语言模型的基础。其中，注意力机制是Transformer的核心组件。本文将重点探讨一种特殊的注意力机制——因果注意力（Causal Attention），这是构建GPT等自回归语言模型的关键技术。

什么是因果注意力？

因果注意力，也称为遮蔽注意力（Masked Attention），是自注意力机制的一种变体。与标准的自注意力不同，因果注意力限制模型在处理任何给定Token时，只能"看到"序列中当前及之前的Token，而不能访问未来的Token信息。

这种限制对于语言模型至关重要，因为它模拟了人类阅读和生成文本的方式——我们总是从左到右依次处理文本，而无法预先知道后面的内容。

因果注意力的实现原理

1. 标准注意力机制回顾

在标准自注意力中，模型计算注意力权重时考虑序列中的所有Token，包括当前Token之前和之后的Token。这适用于机器翻译等任务，但在语言建模中会导致信息泄露问题。

2. 引入因果遮蔽

因果注意力的核心思想是通过遮蔽矩阵（mask）来阻止模型访问未来的Token信息。具体来说：

首先计算标准的注意力得分矩阵
创建一个下三角矩阵作为遮蔽，将对角线以上的元素设为零
将遮蔽应用于注意力得分矩阵
重新标准化注意力权重，确保每行和为1

# 示例代码：创建因果遮蔽
context_length = attn_scores.shape[0]
mask = torch.tril(torch.ones(context_length, context_length))
masked_attn = attn_weights * mask

3. 高效实现技巧

在实践中，我们可以利用softmax函数的数学特性，采用更高效的方式实现因果遮蔽：

使用上三角矩阵创建遮蔽（对角线以上为1）
将这些1替换为负无穷大（-inf）
直接应用softmax函数

这种方法利用了softmax对-inf输入会输出0的特性，避免了额外的标准化步骤。

# 高效实现因果注意力
mask = torch.triu(torch.ones(context_length, context_length), diagonal=1)
masked = attn_scores.masked_fill(mask.bool(), -torch.inf)
attn_weights = torch.softmax(masked / keys.shape[-1]**0.5, dim=1)

防止过拟合：注意力Dropout

在训练大型语言模型时，过拟合是一个常见问题。为了增强模型的泛化能力，我们可以在注意力机制中引入Dropout技术。

Dropout原理

Dropout在训练期间随机"关闭"一部分神经元（将其输出设为零），迫使模型不依赖于任何特定的神经元组合。在注意力机制中，我们通常：

在计算注意力权重后应用Dropout
对保留的注意力权重进行缩放（乘以1/(1-dropout_rate)）

# 注意力Dropout示例
dropout = torch.nn.Dropout(0.5)  # 50%的Dropout率
dropped_attn = dropout(attn_weights)

为什么需要缩放？

Dropout在训练时随机丢弃部分权重，但在推理时使用全部权重。为了保持训练和推理时注意力权重的期望值一致，我们需要在训练时对保留的权重进行放大。

完整实现：CausalAttention类

结合上述技术，我们可以实现一个完整的因果注意力类，支持批量处理和Dropout：

class CausalAttention(nn.Module):
    def __init__(self, d_in, d_out, dropout=0.1):
        super().__init__()
        self.W_query = nn.Linear(d_in, d_out)
        self.W_key = nn.Linear(d_in, d_out)
        self.W_value = nn.Linear(d_in, d_out)
        self.dropout = nn.Dropout(dropout)
        
    def forward(self, x):
        queries = self.W_query(x)
        keys = self.W_key(x)
        values = self.W_value(x)
        
        attn_scores = queries @ keys.transpose(-2, -1)
        mask = torch.triu(torch.ones(attn_scores.shape[-2:]), diagonal=1)
        masked = attn_scores.masked_fill(mask.bool(), -torch.inf)
        attn_weights = torch.softmax(masked / keys.shape[-1]**0.5, dim=-1)
        attn_weights = self.dropout(attn_weights)
        
        return attn_weights @ values

这个类的主要特点包括：