✦星尘知识库

ChatGPT

NLP📅 2026-06-30👁 3 阅读♥ 0 点赞⏱ 约 10 分钟

SingleHeadAttention

class SingleHeadAttention(nn.Module):
    # 单头注意力机制
    def __init__(self, config):
        super().__init__()
        self.key = nn.Linear(config.n_embd, config.head_size)
        self.value = nn.Linear(config.n_embd, config.head_size)
        self.query = nn.Linear(config.n_embd, config.head_size)
        self.head_size = config.head_size

        # 尝试学习新的写法，attention_mask 通过 register_buffer 注册
        # 因为不用计算 梯度，所以节约内存和显存，速度也更快
        self.register_buffer(
            'attention_mask', 
            torch.tril(
                torch.ones(config.block_size, config.block_size)
            ))
        self.dropout = nn.Dropout(config.dropout)

    def forward(self, x):
        batch_size, seq_len, hidden_size = x.size()
        k = self.key(x)
        v = self.value(x)
        q = self.query(x)
        weight = q @ k.transpose(-2, -1)   # @ 就是 torch.matmul 的简化写法
        # 一定要在 softmax 前除以 sqrt(head_size)
        weight = weight.masked_fill(
            self.attention_mask[:seq_len, :seq_len] == 0, 
            float('-inf')
        ) / math.sqrt(self.head_size)  # 这里的 hidden_size 其实是 head_size，因为是单头
        weight = F.softmax(weight, dim=-1)
        weight = self.dropout(weight)
        out = weight @ v
        return out

细节理解

问题一

如果把attention_mask定义成一个普通的下三角张量，和这种方式有什么区别？

在 PyTorch 中，如果你将 attention_mask 定义为一个普通的下三角张量，而不是使用 register_buffer，会有以下几点区别：

模型状态（Model State）

  • **使用 **register_buffer：
    • attention_mask 会被视为模型的一部分，自动包含在模型的状态字典（state_dict）中。
    • 在保存和加载模型时，attention_mask 会随模型参数一同保存和恢复，确保模型在不同环境或训练阶段中保持一致。
  • 定义为普通张量：
    • attention_mask 不会被自动包含在模型的状态字典中。
    • 在保存和加载模型时，你需要手动处理 attention_mask 的保存和加载，否则在加载模型后，attention_mask 可能不会被正确恢复，导致模型行为不一致。

内存和计算效率

  • **使用 **register_buffer：
    • PyTorch 不会为 attention_mask 计算梯度，因为缓冲区（buffer）默认不需要梯度。
    • 这可以节省内存和显存，尤其是在训练大型模型时，因为不需要为 attention_mask 分配梯度存储空间。
    • 训练速度也会略有提升，因为跳过了不必要的梯度计算步骤。
  • 定义为普通张量：
    • 如果你将 attention_mask 定义为普通张量并在模型中使用，PyTorch 默认会为其分配梯度（除非你手动设置 requires_grad=False）。
    • 如果不设置 requires_grad=False，PyTorch 会为 attention_mask 计算梯度，这会浪费内存和计算资源，因为 attention_mask 是一个固定的掩码，不需要更新。
    • 即使你设置 requires_grad=False，attention_mask 仍然不会被自动包含在模型的状态字典中，你仍需手动处理其保存和加载。

灵活性

  • **使用 **register_buffer：
    • attention_mask 作为模型的一部分，管理更加简单和统一。
    • 在模型迁移、部署或多 GPU 训练时，PyTorch 会自动处理缓冲区的分配和同步。
  • 定义为普通张量：
    • 你需要手动管理 attention_mask，例如在模型的 __init__ 方法中创建，并在 forward 方法中使用。
    • 在多 GPU 训练或模型迁移时，你可能需要额外处理 attention_mask 的设备分配和同步。

总结

  • **使用 **register_buffer：
    • 更高效：避免不必要的梯度计算，节省内存和显存。
    • 更一致：attention_mask 自动包含在模型的状态中，确保模型在保存和加载时完整迁移。
    • 更简洁：PyTorch 自动管理缓冲区的设备和同步，减少手动操作。
  • 定义为普通张量：
    • 需要手动处理保存和加载，可能导致模型状态不一致。
    • 如果不设置 requires_grad=False，会浪费资源计算不必要的梯度。
    • 即使设置 requires_grad=False，仍然需要手动管理 attention_mask 的状态。

因此，对于像 attention_mask 这样固定的、不需要更新的张量，使用 register_buffer 是 PyTorch 推荐的做法，能够优化资源使用并确保模型的一致性。

问题二

weight = weight.masked_fill(
            self.attention_mask[:seq_len, :seq_len] == 0, 
            float('-inf')
        ) / math.sqrt(self.head_size)

为什么这段代码需要[:seq_len, :seq_len]?

要搞清楚这个问题需要知道weight 的形状，weight 的形状是 (batch_size, seq_len, seq_len)，因此在进行掩码时要切片。

DEEPSEEK:

在注意力机制中，使用self.attention_mask[:seq_len, :seq_len]进行切片是处理变长序列的关键技术，原因如下：

预分配掩码矩阵

代码中初始化时创建了一个固定大小的掩码矩阵：

torch.ones(config.block_size, config.block_size)

这里config.block_size是模型支持的最大序列长度（如512或1024）。这是一个预分配的方形矩阵，大小始终为[block_size, block_size]。

实际序列长度可变

在训练/推理时：

batch_size, seq_len, hidden_size = x.size()

seq_len（实际序列长度）可能小于block_size（最大长度）。例如：

  • block_size = 512
  • 当前输入seq_len = 128

切片的作用

[:seq_len, :seq_len]从预分配的大矩阵中动态提取子矩阵：

self.attention_mask[:128, :128]

这保证了：

  • 形状匹配：生成的掩码矩阵大小[128, 128]，与注意力权重weight的[batch, 128, 128]维度匹配
  • 资源优化：避免为每个不同长度的序列重新创建新掩码

实际示例

假设：

  • block_size = 4（最大长度4）
  • 预分配掩码：
[[1, 0, 0, 0],
 [1, 1, 0, 0],
 [1, 1, 1, 0],
 [1, 1, 1, 1]]

情况1：序列长度seq_len=2
切片后：

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

情况2：序列长度seq_len=3
切片后：

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

为什么必须切片？

如果不切片：

  • 当seq_len < block_size时：尝试用[512,512]的掩码处理[128,128]的权重 → 形状不匹配错误
  • 当seq_len > block_size时：超出预分配范围 → 索引越界错误

设计优势

  1. 计算高效：预分配一次，重复使用（避免每个batch重新计算tril）
  2. 内存友好：大矩阵常驻内存，动态切片无额外开销
  3. 灵活性：适应不同长度的输入序列
  4. 与位置编码兼容：配合可学习的位置编码处理变长序列

这种设计模式在Transformer实现中很常见，如HuggingFace的BertModel也使用类似机制处理attention_mask。

评论（0）

评论加载中…