根据张量形状判断线性层与批次、序列维度的矩阵运算合法性

你好，欢迎开始这门面向大模型后训练的基础课程。后训练中的 SFT、偏好优化、PPO/GRPO、奖励模型训练，表面上是不同算法，但都依赖同一项基本能力：看到一个张量，能立刻判断每个轴代表什么、某个计算是否能执行，以及输出会是什么形状。

本模块先补齐这套“形状语言”。今天聚焦最常见的一类操作：线性层。你会学会判断批次维、序列维和特征维在 nn.Linear 与矩阵乘法中如何变化；这也是日后阅读 Transformer、训练框架和报错日志时最先需要用到的能力。

---

先把形状读成“有语义的坐标轴”

一个张量的 shape 是各轴长度组成的元组。例如：

$$X\in {\mathbb{R}}^{B\times T\times H}$$

在语言模型训练中，最常见的解释如下：

符号	常见名称	含义	例子
$B$	batch size	一次送入模型的样本数	8 条对话
$T$	sequence length	每条样本包含的 token 数	512 个 token
$H$	hidden size	每个 token 的隐藏表示长度	4096
$O$	output features	线性层输出特征数	11008
$V$	vocabulary size	词表大小	151936

因此，形状为 [8, 512, 4096] 的隐藏状态张量，通常表示：

• 有 $8$ 条序列；
• 每条序列有 $512$ 个 token 位置；
• 每个 token 位置有一个 $4096$ 维向量。

这里有一个必须建立的习惯：形状合法，不等于语义正确。

例如，[T, B, H] 和 [B, T, H] 都是三维张量，送入 nn.Linear(H, O) 都不会报错，因为最后一维都是 $H$。但若某个模块约定第 0 轴是 batch，而你误把序列轴放在第 0 轴，模型虽然“能跑”，数据含义却已错位。后训练工程里，这类静默错误往往比直接报错更难排查。

为了建立对“矩阵形状”的直觉，可以先看一段简短的可视化说明。

---

线性层：只检查最后一维，并保留前面的轴

PyTorch 的线性层写作：

nn.Linear(in_features, out_features)

若某个 token 的隐藏向量为

$$x\in {\mathbb{R}}^H$$

线性层参数为

$$W\in {\mathbb{R}}^{O\times H},b\in {\mathbb{R}}^O$$

则输出为

$$y=x{W}^{\top }+b$$

其中：

• 输入维度是 $H$，所以权重的第二维必须是 $H$；
• 输出维度是 $O$，所以权重的第一维是 $O$；
• 输出向量的维度是 $O$。

严格地说，含 bias 的 nn.Linear 是仿射变换，而非纯线性变换；但深度学习工程中通常仍称其为“线性层”。

PyTorch 以 (out_features, in_features) 存储权重。因此：

proj = nn.Linear(4096, 11008)

对应的参数形状是：

proj.weight.shape  # [11008, 4096]
proj.bias.shape    # [11008]

最关键的 API 规则是：

$$[\dots ,H]\text{经过}\mathrm{Linear}(H,O)\text{得到}[\dots ,O]$$

也就是说，nn.Linear 把最后一维视为输入特征维；此前的任意多个维度都会被保留。

现在把这条规则用在语言模型隐藏状态上：

import torch
import torch.nn as nn

B, T, H, O = 4, 128, 768, 3072

hidden_states = torch.randn(B, T, H)
up_proj = nn.Linear(H, O)

output = up_proj(hidden_states)

print(hidden_states.shape)  # torch.Size([4, 128, 768])
print(output.shape)         # torch.Size([4, 128, 3072])

这个计算合法，因为：

1. hidden_states 的最后一维是 $768$；
2. up_proj.in_features 也是 $768$；
3. batch 维 $B=4$ 与序列维 $T=128$ 原样保留；
4. 输出特征维由 $768$ 变成 $3072$。

从概念上说，线性层对每一个位置的 $H$ 维向量应用同一组参数。也可以把 [B, T, H] 暂时想成 [B*T, H]，经过线性层后得到 [B*T, O]，再恢复成 [B, T, O]。但实现时不需要手动展平，nn.Linear 已经支持任意数量的前导维度。

相反，下面的计算必然非法：

bad_hidden_states = torch.randn(4, 128, 512)
output = up_proj(bad_hidden_states)

因为输入最后一维是 $512$，而线性层要求 $768$。无论 batch size 和 sequence length 是多少，都无法弥补这个不匹配。

---

Transformer 中的批次维、序列维与特征维

下面这张 Transformer 架构图值得作为今后读代码时的形状参考。

以 decoder-only 大语言模型为例，一条常见的形状链是：

模块或张量	形状	发生了什么
input_ids	[B, T]	每个位置是一个整数 token ID
token embedding 输出	[B, T, H]	每个 token ID 被查表为 $H$ 维向量
Transformer block 输出	[B, T, H]	注意力和 MLP 改变表示内容，但通常保持形状
LM head 输出 logits	[B, T, V]	每个 token 位置得到对整个词表的打分

注意：embedding 层和线性层都可使张量最后一维变化，但机制不同。

• Embedding 将离散 token ID 映射成连续向量，常见输入是 [B, T]，输出是 [B, T, H]。
• 线性层接收已经是浮点表示的张量，例如 [B, T, H]，输出 [B, T, O]。

在后训练中，最常见的最终线性投影是语言模型头：

$$\mathrm{LMHead}:[B,T,H]\mapsto [B,T,V]$$

它为每个序列位置、每个词表 token 产生一个 logit。之后的 softmax、标签移位与交叉熵损失会在后续课程中展开。

还可以用 MLP 层观察特征维如何变化。一个典型 Transformer MLP 中可能有：

$$[B,T,H]$$

经过上投影层得到：

$$[B,T,4H]$$

经过激活函数后形状不变，随后经下投影层回到：

$$[B,T,H]$$

这解释了为何残差连接通常合法：主分支输出和残差分支输入都具有 [B, T, H]，能够逐元素相加。

---

矩阵乘法：中间维必须相同

线性层的核心是矩阵乘法。对于二维矩阵：

$$A\in {\mathbb{R}}^{m\times k},C\in {\mathbb{R}}^{k\times n}$$

乘积合法：

$$AC\in {\mathbb{R}}^{m\times n}$$

判断规则只有一句：

左侧矩阵的列数，必须等于右侧矩阵的行数。

其中 $k$ 是被消去的“中间维”。输出保留左矩阵的行数 $m$ 与右矩阵的列数 $n$。

例如：

$$[128,768]\times [768,3072]=[128,3072]$$

而：

$$[128,768]\times [3072,768]$$

不合法，因为相邻的维度 $768$ 和 $3072$ 不相等。

这恰好解释了一个常见困惑。若权重以 PyTorch 的格式存储：

$$W\in {\mathbb{R}}^{O\times H}$$

而输入以“每行一个样本”的格式表示：

$$X\in {\mathbb{R}}^{N\times H}$$

那么应使用：

$$X{W}^{\top }\in {\mathbb{R}}^{N\times O}$$

因为：

$$[N,H]\times [H,O]=[N,O]$$

所以，下面手写实现与 nn.Linear(H, O) 的权重方向一致：

x = torch.randn(4, 128, 768)
layer = nn.Linear(768, 3072)

y_manual = x @ layer.weight.T + layer.bias
y_module = layer(x)

print(y_manual.shape)  # torch.Size([4, 128, 3072])

实际工程中优先使用 layer(x) 或 torch.nn.functional.linear；手动写矩阵乘法主要用于理解、调试和验证自定义实现。

torch.matmul 或 @ 对高阶张量采用“最后两维是矩阵维”的规则。因而：

$$[B,T,H]\times [H,O]=[B,T,O]$$

对于左侧张量，最后两维 [T, H] 可视为一个 $T\times H$ 矩阵，前面的 $B$ 轴则表示有 $B$ 组这样的矩阵。

这与模型语义中“序列轴”的叫法并不冲突：

• 对 nn.Linear 而言，B 和 T 都只是被保留的前导轴；
• 对 matmul 而言，T 在这次操作中扮演左矩阵的行数；
• 对注意力而言，T 还会表示 query 或 key 的位置数。

同一个轴的角色取决于当前操作，不能只凭轴的位置机械记忆。

例如，注意力分数的形状常见为：

$$Q\in {\mathbb{R}}^{B\times T\times D},K^{\top }\in {\mathbb{R}}^{B\times D\times T}$$

因此：

$$Q{K}^{\top }\in {\mathbb{R}}^{B\times T\times T}$$

这里两个 $T$ 分别对应 query 位置与 key 位置。我们之后会手工计算注意力；现在只需能检查：相乘的中间维都是 $D$，因此这一矩阵乘法合法。

---

广播：它能帮助逐元素运算，不能修复矩阵乘法

线性层公式中的 bias：

$$Y=X{W}^{\top }+b$$

看起来像把 [B, T, O] 与 [O] 相加。它之所以合法，依赖的是 broadcasting（广播）。

广播用于逐元素加、减、乘、除等操作。比较两个形状时，从最右侧向左对齐；每一对对应维度必须：

1. 数值相等；或
2. 其中一方为 $1$。

例如：

$$[B,T,O]+[O]$$

会把 [O] 看作 [1,1,O]，然后在 batch 和 sequence 方向复制，因此结果是 [B,T,O]。

下面的视频用二维张量加一维向量展示了同一规则。

一个很常见的后训练代码问题是：你想为 batch 中每条样本乘一个系数，但写出了错误形状。

rewards = torch.randn(B, T)
sample_weights = torch.randn(B)

直接计算 rewards * sample_weights 通常不符合预期，因为 [B] 会优先对齐 rewards 的最后一维 $T$，而不是第 0 维的 $B$。

应显式把权重写成：

weighted_rewards = rewards * sample_weights[:, None]

这时形状是：

[B, T] * [B, 1]

广播发生在长度为 $1$ 的序列维上，每条样本使用自己的权重。

对三维隐藏状态 [B, T, H] 做按样本缩放，则可使用：

scales = torch.randn(B, 1, 1)
scaled_hidden_states = hidden_states * scales

广播很强大，但务必记住它的边界：

• 广播可以使逐元素运算合法；
• 广播不改变矩阵乘法的中间维匹配要求；
• H 与 O 不相等时，不能指望广播使 X @ W 自动正确；
• 显式使用 unsqueeze、None 或 reshape，通常比依赖“碰巧对齐”更安全。

---

一套可用于读代码和排错的形状检查法

以后看到任意训练代码中的一行张量操作，可以按下面五步检查。

1. 写出每个张量的完整 shape。
   不满足于“这是 logits”或“这是 reward”；直接标记为 [B, T, V]、[B]、[B, T] 等。

2. 为每个轴补上语义。
   确认它是 batch、sequence、head、hidden、vocab，还是一个临时矩阵维度。

3. 识别操作类别。
   逐元素操作检查广播；矩阵乘法检查相邻中间维；线性层检查输入最后一维；拼接则检查除拼接轴外的所有维度。

4. 在执行前预测输出 shape。
   例如看到 nn.Linear(768, 3072) 作用于 [2, 1024, 768]，应先得出 [2, 1024, 3072]。

5. 用断言把推理写进代码。

assert hidden_states.ndim == 3
assert hidden_states.shape[-1] == up_proj.in_features

projected = up_proj(hidden_states)

assert projected.shape[:2] == hidden_states.shape[:2]
assert projected.shape[-1] == up_proj.out_features

在需要排查框架调用链或自定义 loss 时，形状日志也非常有效：

print(f"input_ids: {input_ids.shape}")
print(f"hidden_states: {hidden_states.shape}")
print(f"logits: {logits.shape}")
print(f"labels: {labels.shape}")

尤其要警惕两类错误：

• 显式报错：如 mat1 and mat2 shapes cannot be multiplied，通常是线性层输入特征维或矩阵乘法中间维不匹配。
• 静默错误：如 batch 与 sequence 轴被交换、广播对齐到了错误轴、mask 的 shape 虽可广播但语义错位。它们不会停止训练，却会让结果异常。

---

小结

今天建立了几条后训练工程中高频使用的形状规则：

• [B, T, H] 通常表示一批 token 的隐藏状态，$B$、$T$、$H$ 分别承担不同语义。
• nn.Linear(H, O) 只要求输入最后一维等于 $H$，并将最后一维替换为 $O$。
• PyTorch 线性层权重的存储形状是 [O, H]；以行向量形式计算时使用 $X{W}^{\top }$。
• 二维矩阵乘法要求左矩阵列数等于右矩阵行数；高阶 matmul 则在最后两维应用这一规则。
• 广播服务于逐元素运算，例如线性层 bias 的 [O] 加到 [B, T, O]；它不能修复错误的矩阵乘法。
• 合法的 shape 不保证语义正确，因此要同时检查轴长度与轴含义。

下一课会把这些形状规则落实到具体数字上：手工计算向量点积和小型矩阵乘法。那时你会更清楚地看到，线性层输出的每一个元素究竟是怎样由输入特征与权重共同计算出来的。