介绍
张量的 .masked_fill() 方法，是实现注意力掩码的核心操作，专门用于根据掩码张量对目标张量的指定位置进行填充，是大模型注意力机制中不可或缺的工具。
语法
根据布尔类型的 mask 张量，将目标张量中与 mask 里 True 位置对应的元素替换为指定的 value，False 位置的元素保持不变。
提示：mask 中 True = 需要填充的位置，False = 保留原值的位置。
Tensor.masked_fill(mask, value) → Tensor
参数：
mask：布尔型张量（必须通过.bool()转换得到），形状需与目标张量广播兼容（通常是完全相同，或仅少 batch 等维度，可通过unsqueeze补维）；value：要填充的数值（如-1e9、0等），需与目标张量数据类型兼容；- 返回值：与原张量形状、设备一致的新张量（原张量不会被修改）。
例子
简单例子
import torch
# 1. 创建原始张量
x = torch.tensor([[1, 2, 3],
[4, 5, 6],
[7, 8, 9]
])
print("原始张量 x：")
print(x)
# 2. 创建布尔掩码（形状和 x 相同）
# 掩码中 True 的位置会被填充，False 位置保持不变
mask = torch.tensor([[False, True, True],
[False, False, True],
[False, False, False]])
# 3. 使用 masked_fill 填充（将 mask 为 True 的位置填充为 0）
x_filled = x.masked_fill(mask, 0)
print("\n填充后的张量 x_filled：")
print(x_filled)
执行结果：
原始张量 x：
tensor([[1, 2, 3],
[4, 5, 6],
[7, 8, 9]])
填充后的张量 x_filled：
tensor([[1, 0, 0],
[4, 5, 0],
[7, 8, 9]])
掩码例子
import torch
# 维度，3行4列
dim = 3, 4
# 创建值全为1的矩阵
mask = torch.ones(dim)
print(mask)
# 返回对角线以下为0的矩阵
mask2 = torch.triu(mask)
print(mask2)
# 返回0为false，非0为true的矩阵
mask3 = mask2.bool()
print("最终掩码：\n", mask3)
# 数据
data = torch.randn(dim)
print("原始数据：", data)
# mask3 中为 True 的位置，将data中同位置的修改为 0，其余不变
data2 = data.masked_fill(mask3, 0)
print("masked_fill()执行的数据：\n", data2)
执行结果：
tensor([[1., 1., 1., 1.],
[1., 1., 1., 1.],
[1., 1., 1., 1.]])
tensor([[1., 1., 1., 1.],
[0., 1., 1., 1.],
[0., 0., 1., 1.]])
最终掩码：
tensor([[ True, True, True, True],
[False, True, True, True],
[False, False, True, True]])
原始数据： tensor([[-0.0093, -0.3703, 2.3280, 0.1476],
[-0.0731, 2.7202, 1.5814, 0.2370],
[ 0.2546, -1.1607, 0.1380, 0.1145]])
masked_fill()执行的数据：
tensor([[ 0.0000, 0.0000, 0.0000, 0.0000],
[-0.0731, 0.0000, 0.0000, 0.0000],
[ 0.2546, -1.1607, 0.0000, 0.0000]])
结合注意力掩码场景
下面结合自回归掩码案例，完整演示 .masked_fill() 的使用，让你直观理解其作用：
import torch
# 1. 模拟场景：batch_size=1, seq_len=4 的注意力分数
batch_size, seq_len = 1, 4
# 随机生成注意力分数（Q@K^T/√d_k 的结果）
attn_scores = torch.randn(batch_size, seq_len, seq_len)
print("原始注意力分数：")
print(attn_scores)
# 输出示例：
# tensor([[[ 0.56, 1.23, -0.45, 0.89],
# [-0.12, 0.78, 1.01, -0.33],
# [ 0.99, -0.21, 0.67, 1.12],
# [-0.55, 0.44, -0.88, 0.77]]])
# 2. 生成自回归掩码（上三角为1，下三角为0）
mask_raw = torch.triu(torch.ones(seq_len, seq_len), diagonal=1)
# 转为布尔掩码（True=未来位置，需要屏蔽）
mask_bool = mask_raw.bool()
# 补batch维度（让mask形状与attn_scores一致：[1,4,4]）
mask_bool = mask_bool.unsqueeze(0)
print("\n布尔掩码（True=需要屏蔽的位置）：")
print(mask_bool)
# 输出：
# tensor([[[False, True, True, True],
# [False, False, True, True],
# [False, False, False, True],
# [False, False, False, False]]])
# 3. 核心操作：用masked_fill屏蔽未来位置
# 填充值选-1e9（极小值），Softmax后这些位置权重趋近于0
attn_scores_masked = attn_scores.masked_fill(mask_bool, -1e9)
print("\n应用masked_fill后的注意力分数：")
print(attn_scores_masked)
# 输出示例：
# tensor([[[ 0.56, -1e9, -1e9, -1e9],
# [-0.12, 0.78, -1e9, -1e9],
# [ 0.99, -0.21, 0.67, -1e9],
# [-0.55, 0.44, -0.88, 0.77]]])
# 4. 验证：Softmax后屏蔽位置权重趋近于0
attn_weights = torch.softmax(attn_scores_masked, dim=-1)
print("\nSoftmax后的注意力权重（屏蔽位置≈0）：")
print(attn_weights)
# 输出示例（数值近似）：
# tensor([[[1.00, 0.00, 0.00, 0.00],
# [0.29, 0.71, 0.00, 0.00],
# [0.54, 0.15, 0.31, 0.00],
# [0.18, 0.35, 0.11, 0.36]]])
注意
（1）形状匹配要求
mask 的形状必须与目标张量广播兼容，常见场景：
- 目标张量形状：
[batch, seq_len, seq_len]（注意力分数）； - 原始掩码形状：
[seq_len, seq_len]（无batch维度）； - 需通过
mask.unsqueeze(0)补batch维度，使形状变为[1, seq_len, seq_len]，PyTorch会自动广播到[batch, seq_len, seq_len]。
（2）填充值的选择
在注意力掩码中，填充值通常选 -1e9（或 -inf），原因：
- Softmax 公式：$Softmax(x_i) = \frac{e^{x_i}}{\sum_j e^{x_j}}$；
- 当 $x_i = -1e9$ 时，$e^{-1e9} ≈ 0$，Softmax后该位置权重几乎为0，实现“屏蔽”效果。
与 masked_fill_() 的区别
.masked_fill()：返回新张量，原张量不变（推荐使用，避免原地修改导致梯度问题）；.masked_fill_()：原地修改原张量（带下划线的方法是原地操作）。
其他常见使用场景
除了注意力掩码，.masked_fill() 还常用于：
- 填充掩码（Pad Mask）：将 padding token 对应的注意力分数填充为
-1e9，避免模型关注无效填充； - 数据清洗：将张量中无效值（如NaN、超出范围的值）替换为指定值；
- 条件替换：根据布尔条件批量替换张量元素。
总结
.masked_fill(mask, value)核心作用是将张量中mask为True的位置替换为value，False位置保留原值；- 在注意力机制中，需配合布尔掩码使用，填充
-1e9实现“屏蔽指定位置”的效果； - 使用时需保证
mask与目标张量形状广播兼容，优先使用非原地的.masked_fill()而非.masked_fill_()。