介绍
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()生成上三角矩阵，可根据需求选择。