重点

加性注意力函数中的问题1有助于更好的理解注意力机制

掩蔽softmax操作

#@save
def masked_softmax(X, valid_lens):
    """通过在最后一个轴上掩蔽元素来执行softmax操作"""
    # X:3D张量，valid_lens:1D或2D张量
    if valid_lens is None:
        return nn.functional.softmax(X, dim=-1)
    else:
        shape = X.shape
        if valid_lens.dim() == 1:
            valid_lens = torch.repeat_interleave(valid_lens, shape[1])
        else:
            valid_lens = valid_lens.reshape(-1)
        # 最后一轴上被掩蔽的元素使用一个非常大的负值替换，从而其softmax输出为0
        X = d2l.sequence_mask(X.reshape(-1, shape[-1]), valid_lens,
                              value=-1e6)
        return nn.functional.softmax(X.reshape(shape), dim=-1)

函数1：dim()

valid_lens.dim()			#获得维度个数

函数2：torch.repeat_interleave

torch.repeat_interleave(input, repeats, dim=None) → Tensor
# 不指定维度会输出一个向量

例子：
# 一维张量，不指定dim，每个元素重复2次
x = torch.tensor([1, 2, 3])
torch.repeat_interleave(x, 2)
# 输出: tensor([1, 1, 2, 2, 3, 3])
# 二维张量，不指定dim，每个元素重复2次
y = torch.tensor([[1, 2], [3, 4]])
torch.repeat_interleave(y, 2)
# 输出: tensor([1, 1, 2, 2, 3, 3, 4, 4])
# 二维张量，指定dim=0，每行元素重复3次
torch.repeat_interleave(y, 3, dim=0)
# 输出: tensor([[1, 2], [1, 2], [1, 2], [3, 4], [3, 4], [3, 4]])
# 二维张量，指定dim=1，每列元素重复3次
torch.repeat_interleave(y, 3, dim=1)
# 输出: tensor([[1, 1, 1, 2, 2, 2], [3, 3, 3, 4, 4, 4]])
# 二维张量，指定dim=0，不同行元素重复不同次数
torch.repeat_interleave(y, torch.tensor([1, 2]), dim=0)
# 输出: tensor([[1, 2], [3, 4], [3, 4]])

函数3：sequence_mask

#代码十分优雅
def sequence_mask(X, valid_len, value=0):
    """在序列中屏蔽不相关的项"""
    maxlen = X.size(1)  # 获得数据表示范围的长度
    mask = torch.arange((maxlen), dtype=torch.float32,
                        device=X.device)[None, :] < valid_len[:, None]
    X[~mask] = value  # ~mask取反，True -> False    False -> True
    return X

X = torch.tensor([[1, 2, 3], [4, 5, 6]])
sequence_mask(X, torch.tensor([1, 2]))

结果：

mask = torch.arange((maxlen), dtype=torch.float32,device=X.device)[None, :] < valid_len[:, None]

关键代码分析

步骤1：创建位置索引

python

torch.arange((maxlen), dtype=torch.float32, device=X.device)
# 生成: [0, 1, 2, ..., maxlen-1]
# 例如 maxlen=5: [0, 1, 2, 3, 4]

步骤2：调整形状以便广播

python

[None, :]  # 形状从 [maxlen] 变为 [1, maxlen]
valid_len[:, None]  # 形状从 [batch_size] 变为 [batch_size, 1]

步骤3：创建布尔掩码

python

mask = position_indices < valid_len_expanded

举例：

text

假设:
- maxlen = 5
- valid_len = [3, 4] (批次大小=2)

位置索引: [[0, 1, 2, 3, 4]]        (形状: [1, 5])
有效长度: [[3], [4]]              (形状: [2, 1])

广播比较:
[[0,1,2,3,4] < [3]] → [T,T,T,F,F]   # 第一个序列
[[0,1,2,3,4] < [4]] → [T,T,T,T,F]   # 第二个序列

最终掩码:
[[True, True, True, False, False],
 [True, True, True, True, False]]

加性注意力

#@save
class AdditiveAttention(nn.Module):
    """加性注意力"""
    def __init__(self, key_size, query_size, num_hiddens, dropout, **kwargs):
        super(AdditiveAttention, self).__init__(**kwargs)
        self.W_k = nn.Linear(key_size, num_hiddens, bias=False)
        self.W_q = nn.Linear(query_size, num_hiddens, bias=False)
        self.w_v = nn.Linear(num_hiddens, 1, bias=False)
        self.dropout = nn.Dropout(dropout)

    def forward(self, queries, keys, values, valid_lens):
        queries, keys = self.W_q(queries), self.W_k(keys)
        # 在维度扩展后，
        # queries的形状：(batch_size，查询的个数，1，num_hidden)
        # key的形状：(batch_size，1，“键－值”对的个数，num_hiddens)
        # 使用广播方式进行求和
        features = queries.unsqueeze(2) + keys.unsqueeze(1)
        features = torch.tanh(features)
        # self.w_v仅有一个输出，因此从形状中移除最后那个维度。
        # scores的形状：(batch_size，查询的个数，“键-值”对的个数)
        scores = self.w_v(features).squeeze(-1)
        self.attention_weights = masked_softmax(scores, valid_lens)
        # values的形状：(batch_size，“键－值”对的个数，值的维度)
        return torch.bmm(self.dropout(self.attention_weights), values)

整体分析

输入输出分析

输入张量

• queries: (batch_size, num_queries, query_size)
• keys: (batch_size, num_kv_pairs, key_size)
• values: (batch_size, num_kv_pairs, value_dim)
• valid_lens: (batch_size,) 或 (batch_size, num_queries)

输出张量

• output: (batch_size, num_queries, value_dim)

逐步骤详细分析

步骤1: 线性投影

queries, keys = self.W_q(queries), self.W_k(keys)

作用: 将查询和键映射到相同的隐藏维度

• W_q: 将查询从query_size投影到num_hiddens
• W_k: 将键从key_size投影到num_hiddens
• 输出形状:
  • queries: (batch_size, num_queries, num_hiddens)
  • keys: (batch_size, num_kv_pairs, num_hiddens)

---

步骤2: 维度扩展与广播相加

python

features = queries.unsqueeze(2) + keys.unsqueeze(1)

张量形状变换:

• queries.unsqueeze(2): (batch_size, num_queries, 1, num_hiddens)
• keys.unsqueeze(1): (batch_size, 1, num_kv_pairs, num_hiddens)
• 广播机制: 两个张量在维度1和2上分别扩展，形成所有可能的查询-键组合
• 输出形状: (batch_size, num_queries, num_kv_pairs, num_hiddens)

数学意义:

• 为每个(查询i, 键j)对创建一个组合特征
• 相当于: features[i,j,k] = queries[i,j] + keys[i,k]

---

步骤3: 非线性激活

python

features = torch.tanh(features)

作用:

• 引入非线性变换，增强模型表达能力
• 将特征值限制在[-1, 1]范围内，提高数值稳定性
• 形状保持不变: (batch_size, num_queries, num_kv_pairs, num_hiddens)

---

步骤4: 计算注意力分数

python

scores = self.w_v(features).squeeze(-1)

操作分解:

1. self.w_v(features): 将num_hiddens维特征映射到标量分数
   • 输入: (batch_size, num_queries, num_kv_pairs, num_hiddens)
   • 输出: (batch_size, num_queries, num_kv_pairs, 1)
2. squeeze(-1): 移除最后一个维度
   • 输出: (batch_size, num_queries, num_kv_pairs)

数学表达: scores[i,j,k] = w_v^T · tanh(queries[i,j] + keys[i,k])

---

步骤5: 计算注意力权重

python

self.attention_weights = masked_softmax(scores, valid_lens)

关键功能:

• 应用softmax将分数转换为概率分布
• 使用masked_softmax处理变长序列，确保padding位置权重为0
• 保存注意力权重便于可视化和分析
• 输出形状: (batch_size, num_queries, num_kv_pairs)

---

步骤6: 输出计算

python

return torch.bmm(self.dropout(self.attention_weights), values)

操作流程:

1. self.dropout(attention_weights): 防止过拟合
2. torch.bmm(): 批量矩阵乘法
   • 注意力权重: (batch_size, num_queries, num_kv_pairs)
   • 值矩阵: (batch_size, num_kv_pairs, value_dim)
   • 输出: (batch_size, num_queries, value_dim)

数学表达: output[i,j] = Σ_k attention_weights[i,j,k] · values[i,k]

函数1：传播机制

features = queries.unsqueeze(2) + keys.unsqueeze(1)

目的：queries 和 keys 加在一起

张量形状分析

假设：

• queries 形状: (batch_size, num_queries, num_hiddens)
• keys 形状: (batch_size, num_kv_pairs, num_hiddens)

unsqueeze操作

queries.unsqueeze(2)

• 原始形状: (batch_size, num_queries, num_hiddens)
• unsqueeze(2)后: (batch_size, num_queries, 1, num_hiddens)
• 在第2个维度（0-based索引）插入一个长度为1的维度

keys.unsqueeze(1)

• 原始形状: (batch_size, num_kv_pairs, num_hiddens)
• unsqueeze(1)后: (batch_size, 1, num_kv_pairs, num_hiddens)
• 在第1个维度插入一个长度为1的维度

广播机制

现在两个张量的形状：

• queries: (batch_size, num_queries, 1, num_hiddens)
• keys: (batch_size, 1, num_kv_pairs, num_hiddens)

PyTorch的广播规则：

1. 从尾部维度开始比较
2. 当维度大小相等或其中一个为1时，可以进行广播
3. 大小为1的维度会被扩展以匹配另一个张量

广播后的实际计算：

text

queries 被广播为: (batch_size, num_queries, num_kv_pairs, num_hiddens)
keys    被广播为: (batch_size, num_queries, num_kv_pairs, num_hiddens)

函数2：torch.bmm

在 PyTorch 中，torch.bmm 函数用于执行批量矩阵乘法。具体来说，它计算两个 3D 张量的矩阵乘法。假设 input 是一个形状为 (b, n, m) 的张量，mat2 是一个形状为 (b, m, p) 的张量，那么 torch.bmm(input, mat2) 将返回一个形状为 (b, n, p) 的张量。

注意：

如果 input 是一个 (b×n×m) 张量，mat2 是一个 (b×m×p) 张量，则 out 将是一个 (b×n×p) 张量。

示例

import torch
# 创建两个 3D 张量
input = torch.randn(10, 3, 4)
mat2 = torch.randn(10, 4, 5)
# 执行批量矩阵乘法
result = torch.bmm(input, mat2)
# 打印结果的形状
print(result.size())

输出：

torch.Size([10, 3, 5])

问题1 ：self.attention_weights为什么是这个形状的？

self.attention_weights形状分析

self.attention_weights的形状是(batch_size, num_queries, num_kv_pairs)，这个形状的设计有深刻的数学和语义含义。

形状的数学含义

三维张量的语义

• 第0维 batch_size: 批量中的样本数量
• 第1维 num_queries: 每个样本中的查询数量
• 第2维 num_kv_pairs: 每个样本中的键值对数量

具体解释

对于批量中的每个样本i，每个查询j，每个键k：

text

attention_weights[i, j, k] = 第i个样本中，第j个查询对第k个键的注意力权重

为什么是这个形状？

1. 注意力机制的本质

注意力机制的核心思想是：每个查询应该关注所有键值对中的哪些部分

python

# 伪代码表示
for i in range(batch_size):      # 遍历每个样本
    for j in range(num_queries): # 遍历每个查询
        for k in range(num_kv_pairs): # 遍历每个键
            weight = 计算注意力(查询[j], 键[k])
            attention_weights[i, j, k] = weight

2. 实际应用场景示例

机器翻译场景

text

输入句子: "I love machine learning" (4个词)
输出句子: "我 喜欢 机器 学习" (4个词)

batch_size = 1
num_queries = 4 (输出句子的4个词)
num_kv_pairs = 4 (输入句子的4个词)

attention_weights[0, :, :] 的形状是 (4, 4)，表示：
- 行：输出词的查询（我，喜欢，机器，学习）
- 列：输入词的键（I, love, machine, learning）
- 每个元素：输出词对输入词的关注程度

文本分类场景

text

输入文档: 包含多个句子，每个句子有多个词
查询: [CLS] token (用于分类)

batch_size = N
num_queries = 1 ([CLS] token)
num_kv_pairs = 文档长度

attention_weights[:, 0, :] 表示每个文档中[CLS] token对所有词的关注程度

形状的合理性验证

与最终输出的匹配

python

# 注意力权重形状: (batch_size, num_queries, num_kv_pairs)
# 值矩阵形状: (batch_size, num_kv_pairs, value_dim)
# 输出形状: (batch_size, num_queries, value_dim)

output = torch.bmm(attention_weights, values)

矩阵乘法维度验证：

text

(batch, nq, nk) @ (batch, nk, v_dim) = (batch, nq, v_dim) ✓

与掩码机制的兼容

python

self.attention_weights = masked_softmax(scores, valid_lens)

• scores形状: (batch_size, num_queries, num_kv_pairs)
• valid_lens形状: (batch_size,) 或 (batch_size, num_queries)
• 掩码操作在最后一个维度num_kv_pairs上应用，符合逻辑

为什么不是其他形状？

可能的替代形状及其问题

1. (batch_size, num_kv_pairs, num_queries)
   • 问题：与值矩阵的矩阵乘法无法直接进行
   • 需要转置，增加计算复杂度
2. (batch_size, num_queries \* num_kv_pairs)
   • 问题：丢失了查询-键的对应关系结构
   • 难以解释和可视化
3. (batch_size, 1, num_queries, num_kv_pairs)
   • 问题：不必要的维度，增加内存占用
   • 对实际计算没有帮助