介绍
torch.tril() 函数,是提取矩阵的 下三角部分(Lower Triangular Matrix)的核心工具,在大模型的自注意力机制(如 Decoder 的掩码操作)中是高频用法
作用
返回一个与输入张量形状相同的新张量,其中 下三角部分(包括指定对角线)保留原值,其余位置填充为 0。
下三角矩阵
矩阵中:
- 行号 ≥ 列号 的位置为有效元素(即:矩阵的 左下 是有效元素)
- 行号 < 列号 的位置置 0 (即:矩阵的 右上是0 )
原始矩阵:
[[1., 1., 1., 1.],
[1., 1., 1., 1.],
[1., 1., 1., 1.],
[1., 1., 1., 1.]]
下三角矩阵(diagonal=0):
[[1., 0., 0., 0.],
[1., 1., 0., 0.],
[1., 1., 1., 0.],
[1., 1., 1., 1.]]
语法
torch.tril(input, diagonal=0, *, out=None)
参数
| 参数 | 作用 | 注意事项 |
|---|---|---|
input |
输入张量(支持任意维度,核心作用于最后两个维度) | 若输入是高维张量(如 [batch, seq_len, seq_len]),会对每个二维子矩阵生成下三角 |
diagonal |
有效对角线的偏移量(默认 0) | 正数:保留主对角线上方的对角线;负数:仅保留主对角线下方的对角线 |
out |
输出张量(可选) | 可指定输出张量,避免额外内存分配 |
diagonal 参数:控制“有效对角线”的位置,默认 0(主对角线),正数表示保留主对角线上方的若干对角线,负数表示仅保留主对角线下方的若干对角线。
例子
先通过简单矩阵理解 torch.tril() 的核心效果:
import torch
# 示例1:基础用法(二维矩阵,默认 diagonal=0)
# 定义4×4的全1矩阵
x = torch.ones(4, 4)
# 生成下三角矩阵(保留主对角线及以下)
tril_x = torch.tril(x)
print("原始矩阵:\n", x)
print("下三角矩阵(diagonal=0):\n", tril_x)
# 输出:
# 原始矩阵:
# tensor([[1., 1., 1., 1.],
# [1., 1., 1., 1.],
# [1., 1., 1., 1.],
# [1., 1., 1., 1.]])
# 下三角矩阵(diagonal=0):
# tensor([[1., 0., 0., 0.],
# [1., 1., 0., 0.],
# [1., 1., 1., 0.],
# [1., 1., 1., 1.]])
# 示例2:调整 diagonal 参数(diagonal=1,保留主对角线上方1条对角线)
tril_x_1 = torch.tril(x, diagonal=1)
print("下三角矩阵(diagonal=1):\n", tril_x_1)
# 输出:
# tensor([[1., 1., 0., 0.],
# [1., 1., 1., 0.],
# [1., 1., 1., 1.],
# [1., 1., 1., 1.]])
# 示例3:diagonal=-1(仅保留主对角线下方的元素)
tril_x_neg1 = torch.tril(x, diagonal=-1)
print("下三角矩阵(diagonal=-1):\n", tril_x_neg1)
# 输出:
# tensor([[0., 0., 0., 0.],
# [1., 0., 0., 0.],
# [1., 1., 0., 0.],
# [1., 1., 1., 0.]])
高维张量示例(批量场景,大模型核心用法)
# 示例:3个批量的4×4矩阵(shape [3,4,4])
batch_x = torch.ones(3, 4, 4)
# 对每个批量的矩阵生成下三角
batch_tril = torch.tril(batch_x)
print("批量下三角矩阵形状:", batch_tril.shape) # torch.Size([3, 4, 4])
print("第一个批量的下三角矩阵:\n", batch_tril[0])
# 输出和单个4×4矩阵的下三角结果一致
核心特性与注意事项
仅修改值,不改变张量形状
torch.tril() 是 原地修改值(逻辑上),但不会改变输入张量的形状,输出形状和输入完全一致:
x = torch.randn(4, 4)
tril_x = torch.tril(x)
print("输入形状:", x.shape) # torch.Size([4, 4])
print("输出形状:", tril_x.shape) # torch.Size([4, 4])
数据类型兼容
torch.tril() 兼容所有数值类型(float、int、bool 等),bool 类型常用于注意力掩码:
# bool类型的下三角掩码(注意力机制核心)
bool_x = torch.ones(4, 4, dtype=torch.bool)
bool_tril = torch.tril(bool_x)
print("bool类型下三角掩码:\n", bool_tril)
# 输出:
# tensor([[ True, False, False, False],
# [ True, True, False, False],
# [ True, True, True, False],
# [ True, True, True, True]])
实战核心场景:大模型自注意力掩码(Decoder)
这是 torch.tril() 最核心的应用场景——在 Decoder 的自注意力中,需要生成“未来掩码”(禁止模型看到当前位置之后的token),下三角掩码正好满足这一需求:
# 场景:Decoder自注意力的未来掩码生成
def create_causal_mask(seq_len: int):
# 1. 生成seq_len×seq_len的全1矩阵(bool类型)
mask = torch.ones(seq_len, seq_len, dtype=torch.bool)
# 2. 生成下三角掩码(保留当前位置及之前的token,屏蔽未来token)
causal_mask = torch.tril(mask)
return causal_mask
# 生成序列长度为5的因果掩码
seq_len = 5
causal_mask = create_causal_mask(seq_len)
print("Decoder因果掩码:\n", causal_mask)
# 输出:
# tensor([[ True, False, False, False, False],
# [ True, True, False, False, False],
# [ True, True, True, False, False],
# [ True, True, True, True, False],
# [ True, True, True, True, True]])
# 扩展:批量因果掩码(适配 [batch, seq_len, seq_len])
batch_size = 3
batch_causal_mask = create_causal_mask(seq_len).unsqueeze(0).repeat(batch_size, 1, 1)
print("批量因果掩码形状:", batch_causal_mask.shape) # torch.Size([3, 5, 5])
掩码在注意力计算中的使用
# 模拟注意力分数([batch, seq_len, seq_len])
batch_size = 3
seq_len = 5
attn_scores = torch.randn(batch_size, seq_len, seq_len)
# 应用因果掩码:将未来位置的注意力分数置为 -inf(softmax后为0)
# 先将bool掩码转换为float,False→-inf
mask = batch_causal_mask.float().masked_fill(~batch_causal_mask, -1e9)
# 应用掩码到注意力分数
masked_attn_scores = attn_scores + mask
print("掩码后的注意力分数(第一批量):\n", masked_attn_scores[0])
# 输出中,上三角位置均为 -inf,下三角为原始分数
总结
torch.tril(input, diagonal=0)生成下三角矩阵:保留对角线及以下元素,其余置0,diagonal控制有效对角线的偏移。- 核心特性:支持高维张量(批量处理)、兼容多数据类型(bool 类型用于掩码)、不改变张量形状。
- 核心场景:大模型 Decoder 的自注意力“未来掩码”生成(最常用),也可用于任意需要下三角矩阵的数学计算。
- 互补函数:
torch.triu()生成上三角矩阵,可根据需求选择。