PyTorch 里的 mean(dim)：把第 k 维拍扁 · znzryb
写 RMSNorm 的时候要算 `x.pow(2).mean(dim=-1, keepdim=True)`。`dim` 填几、`keepdim` 要不要，一开始我是靠试的。后来在纸上画了几张图，发现一句话就能记住：
> **`mean(dim=k)` 就是把第 k 维「拍扁」：沿着这一维把数求平均，这一维就消失了。**
先看总图，后面每一节都是在解释它。
![总图：4×2 矩阵沿 dim=0 竖着拍扁得到 [2.5, 3.0]，沿 dim=1 横着拍扁得到 [2.5, 2.5, 2.5, 3.5]](/api/content-images/5607ae8af39cc52eb790a12108c0d65d4ccf5da914d002869542b4ccbaa1a439.svg)
## 二维：竖着拍还是横着拍
拿一个 4 行 2 列的矩阵 `a`，形状是 `(4, 2)`。
**`dim=0`：把第 0 维（行）拍扁。** 4 行压成 1 行，每一列的 4 个数求平均：
- 列 0：`(1 + 2 + 3 + 4) / 4 = 2.5`
- 列 1：`(4 + 3 + 2 + 3) / 4 = 3`
结果 `[2.5, 3]`，形状从 `(4, 2)` 变成 `(2,)`：被拍扁的第 0 维没了。
![手写笔记：对第 0 维取 mean，相当于竖向拍扁，把行给拍扁，第 0 维消失，得到 [2.5, 3]](/api/content-images/57ff510475a55d4161f068f15534586ed05eafa98437862d88e0f7763256aeeb.png)
**`dim=1`：把第 1 维（列）拍扁。** 每一行内部的 2 个数求平均，2 列压成 1 列：
- `(1 + 4) / 2 = 2.5`、`(2 + 3) / 2 = 2.5`、`(3 + 2) / 2 = 2.5`、`(4 + 3) / 2 = 3.5`
结果 `[2.5, 2.5, 2.5, 3.5]`，形状变成 `(4,)`：第 1 维没了。
![手写笔记：对第 1 维取 mean，相当于进行了横向的拍扁，得到 [2.5, 2.5, 2.5, 3.5]](/api/content-images/733b2880f79f465d590802ed13437f0f5abb776d7f9d9301d0cac4c8037e4f74.png)
容易记反的地方：**`dim=0` 不是「对每一行求平均」，而是「把行这一维消灭掉」**，所以是**沿着行的方向**（竖着）求平均，得到的是每一**列**的平均。判断的时候别问「对谁求平均」，问「哪一维要消失」。
## 三维：把长方体拍回一个面
三维看着复杂，其实还是同一件事。
把上面的 `a` 改成 `b = a.reshape(2, 2, 2)`：4 行两两一组分成 2 片，形状 `(2, 2, 2)` 读作「2 片 × 2 行 × 2 列」。把两片前后叠起来，就是一个 2×2×2 的长方体：正面是第 0 片，往里是第 1 片。
三个维度正好是长方体的三个轴：**dim=0 往里（片）、dim=1 往下（行）、dim=2 往右（列）**。`mean(dim=k)` 就是沿第 k 个轴把长方体压扁，压完只剩一个面：
| 写法 | 拍扁哪一维 | 直观上 | 结果 | 形状 |
|---|---|---|---|---|
| `b.mean(dim=0)` | 片 | 两片叠起来、对应位置求平均 | `[[2, 3], [3, 3]]` | `(2, 2)` |
| `b.mean(dim=1)` | 行 | 每片里竖着拍扁 | `[[1.5, 3.5], [3.5, 2.5]]` | `(2, 2)` |
| `b.mean(dim=2)` | 列 | 每片里横着拍扁 | `[[2.5, 2.5], [2.5, 3.5]]` | `(2, 2)` |
不管拍哪一维，长方体都变成一个面，形状都是 `(2, 2)`；区别只在于**哪两维留了下来**。
再看最后一行：`b.mean(dim=2)` 的四个数 `2.5, 2.5, 2.5, 3.5`，和二维时 `a.mean(dim=1)` 的结果一模一样，只是排成了 `(2, 2)`。这很合理：`b` 只是把 `a` 的 4 行分了组，每行内部那 2 个数没变，横着拍扁的结果当然也不变。
## `keepdim=True`：拍成薄片，而不是消失
默认（`keepdim=False`）被拍扁的那一维直接消失。加上 `keepdim=True`，它会留下来、长度变成 1，像是拍成了一张厚度为 1 的薄片：
| 写法 | 形状 |
|---|---|
| `a.mean(dim=1)` | `(4,)` |
| `a.mean(dim=1, keepdim=True)` | `(4, 1)` |
| `b.mean(dim=0, keepdim=True)` | `(1, 2, 2)` |
| `b.mean(dim=-1, keepdim=True)` | `(2, 2, 1)` |
数值完全一样，只差这一维在不在。它的用处在于**拿结果回头和原张量做运算**的时候。
## 落到 RMSNorm
RMSNorm 对每个 token 的 `d_model` 维向量求均方根，再让向量除以它。输入 `x` 的形状是 `(batch, seq_len, d_model)`，比如 `(4, 12, 64)`：
- 要对**每个 token 自己的 64 个数**求平均，所以拍扁的是最后一维：`dim=-1`；
- 拍完要拿去除 `x`，所以用 `keepdim=True` 留下那一维。
把 `x` 画成长方体（往右是 `d_model`），这件事就是：沿 `d_model` 压成一张厚度为 1 的薄片，再把薄片沿 `d_model` 复制回去，和 `x` 逐列相除。不加 `keepdim` 的话，薄片连「厚度」这一维都没了，就对不上：
| 写法 | 结果形状 | `x / 结果` |
|---|---|---|
| `x.pow(2).mean(dim=-1, keepdim=True)` | `(4, 12, 1)` | 可以：最后一维的 1 广播成 64 |
| `x.pow(2).mean(dim=-1)` | `(4, 12)` | 报错：从右往左对齐，`64` 对上了 `12` |
不加 `keepdim` 实际跑出来的报错是：
```text
RuntimeError: The size of tensor a (64) must match the size of tensor b (12) at non-singleton dimension 2
```
写 `dim=-1` 而不是 `dim=2`，是因为前面的维数不固定：输入是 `(seq_len, d_model)` 还是 `(batch, seq_len, d_model)`，`-1` 都指最后一维。
## 自己跑一遍
```python
import torch
a = torch.tensor([[1., 4.],
[2., 3.],
[3., 2.],
[4., 3.]])
print(a.mean(dim=0)) # 竖着拍扁
print(a.mean(dim=1)) # 横着拍扁
print(a.mean(dim=1, keepdim=True).shape) # 拍成薄片，不消失
b = a.reshape(2, 2, 2)
print(b.mean(dim=0))
print(b.mean(dim=-1))
```
输出：
```text
tensor([2.5000, 3.0000])
tensor([2.5000, 2.5000, 2.5000, 3.5000])
torch.Size([4, 1])
tensor([[2., 3.],
[3., 3.]])
tensor([[2.5000, 2.5000],
[2.5000, 3.5000]])
```
`sum`、`max`、`softmax` 这些带 `dim` 参数的函数也是同一个规律：`dim` 指的是**沿着哪一维做**，`sum` / `max` 做完那一维就没了（`softmax` 不拍扁，形状不变，但同样是沿着这一维算）。