19 从矩阵到张量

张量:就是 n 维数组(在这本书里)

先把名词的水分挤掉。物理和微分几何里的「张量」有一套坐标变换律的严格定义(一个对象在换基时分量按特定规律变,才配叫张量——第 13 章「主动被动」之争的高阶版);深度学习里的 tensor 只是「多维数组」的时髦叫法,继承自 NumPy 的 ndarray,对坐标变换律毫无承诺。两个学科为这个词打过不少嘴仗,你只需要知道:读 ML 文献时按「多维数组」理解,读物理文献时按「变换律」理解,别拿一边的定义去纠正另一边。维度数叫「阶」或「轴数」:0 阶是标量(一个数),1 阶是向量,2 阶是矩阵,3 阶以上就统称张量。没有新数学——一个 (32, 128, 4096) 的三阶张量,你完全可以理解成「32 张 128×4096 的矩阵摞成一叠」。

真正的新东西是轴的语义。线性代数里矩阵的行列没有身份;AI 代码里每根轴都有名字,最通用的一套是 (B, L, D):B = batch(同时处理几条样本),L = 序列长度(几个 token),D = 特征维度(每个 token 的向量长度)。视觉任务的方言是 (B, C, H, W)(通道、高、宽),音频是 (B, T, F)(时间、频率)——方言不同,语法一致。「32 句话、每句 128 个词、每词 4096 维」就是 (32, 128, 4096)——念出这句话的能力,比记住任何 API 都值钱。读任何模型代码的第一步,都是在脑子里给每根轴挂牌;写代码的好习惯,是在行尾注释形状——# (B, L, D),大厂代码库里这类注释密如蛛网,它们是形状思维的路标。

形状三板斧:reshape、transpose、切片

shapes.py

import numpy as np

x = np.arange(24)                  # (24,)   0..23
a = x.reshape(2, 3, 4)             # (2,3,4) 重新装箱,数据不动
a = x.reshape(2, -1, 4)            # -1 = 「这维你自己算」

b = a.transpose(1, 0, 2)          # (3,2,4) 换轴:轴 0 和 1 对调
c = a[:, 0, :]                     # (2,4)   切片:抽掉一根轴
d = a[:, :2, :]                    # (2,2,4) 切段:轴还在,变短

# 心算练习:注意力里的经典四连
q = np.zeros((32, 128, 4096))                    # (B, L, D)
q = q.reshape(32, 128, 32, 128)                  # (B, L, H, Dh) 切成 32 个头
q = q.transpose(0, 2, 1, 3)                     # (B, H, L, Dh) 头提到前面
# 此后 q @ k.transpose(0,1,3,2) → (B, H, L, L):每头一张注意力表

三条铁律配三板斧:① reshape 只换包装不动数据,总元素数必须守恒(2×3×4 = 24),而且只在「相邻打包」语义下安全——想交换轴序千万别用 reshape 硬凑,数据会被错位穿插,图像变雪花的经典事故就是 reshape 当 transpose 用;② transpose 真的改变「谁挨着谁」,矩阵转置是它的二维特例;③ 切片抽轴、切段保轴——a[:, 0, :] 和 a[:, 0:1, :] 的形状差一根轴,这个差别马上会在广播里要人命。最后那段「切头四连」是每个 Transformer 实现里都有的真实代码,现在看不懂没关系,第 22 章回来再看会全懂——今天只需要能跟上每一步的形状。

沿轴归约:把一根轴吃掉

第二类日常动作是归约(reduction):沿某根轴做 sum/mean/max/argmax,把那根轴吃掉。语义全靠 axis 参数,读法要练到条件反射:对 (B, L, D) 的张量,mean(axis=0) 是「跨样本平均」(得 (L, D),batch 被吃),mean(axis=1) 是「一句话里所有 token 平均」(得 (B, D)——句向量的最简做法,sentence embedding 的祖传配方),mean(axis=-1) 是「每个 token 自己 4096 维的均值」(得 (B, L),LayerNorm 的第一步)。axis 指向谁,谁消失;搞反 axis 的代码通常能跑、指标全错,是仅次于广播血案的第二大隐形杀手。

softmax 是「归约的亲戚」:它沿某根轴归一化但不吃轴(输出形状不变,值变成和为 1)。softmax(scores, axis=-1) 在注意力里的含义精确到不能再精确:对每个查询各自的那一行分数做预算分配(第 04 章)。把 axis 写成 -2,注意力就变成了「每个 key 把预算分给查询们」——语义完全不同,形状却一样,错了都没报错。形状不报错,语义才是终审法官,这是本章最想立住的一句话。

拼接与堆叠:把张量装订起来

反方向的动作:concatenate 沿已有轴接龙((B, L, D₁) 拼 (B, L, D₂) 得 (B, L, D₁+D₂)——多模态里「图向量拼文向量」就这一下);stack 则新造一根轴(32 张 (L, D) 堆成 (32, L, D)——把零散样本装订成 batch);split/chunk 是逆操作。多头注意力的收尾「把 32 个头的输出拼回 4096 维再过一个投影」,代码就是一次 transpose + reshape(等价于 concat)——第 07 章分块矩阵的「拼接再投影」,在形状语言里只是两板斧。

广播:让小张量自动长大

第 05 章那句「形状对了,循环就消失了」的引擎,就是广播(broadcasting)。规则只有两条,右对齐着比:

从最右边的轴开始逐轴配对:两轴相等 → 兼容;其中一方是 1 → 兼容(把 1 那方「复制扩张」到对方大小);缺轴 → 当 1 处理。
任何一对轴既不相等又没有 1 → 报错。

broadcast.py

X = np.random.randn(32, 128, 4096)   # (B, L, D)
mu = X.mean(axis=-1, keepdims=True)   # (32, 128, 1)   keepdims 保轴!
X_centered = X - mu                   # (32,128,4096)-(32,128,1) ✓ 广播

w = np.random.randn(4096)             # (4096,)
X * w                                 # ✓:缺的轴当 1,右对齐 → 每个 token 逐维加权

col = np.random.randn(128)            # (128,)
# X * col                             ✗:128 对 4096,不等又无 1,报错
X * col[:, None]                      # (128,1) ✓:None 手动补一根轴

两个惯用法值得划线:keepdims=True 让归约(mean/sum/max)后的轴以 1 的身份留守,方便接着广播——LayerNorm 的实现里必有它;[:, None](或 np.newaxis)手动插轴,是「我知道我要广播到哪根轴」的显式声明。广播的本质是虚拟复制:没有真的拷贝内存,只是索引时循环利用——免费的外积、免费的两两配对(第 05 章那个 100×5000 距离矩阵),都是它请客。

⚠ 坑

广播最阴的不是报错,是不报错。经典血案:(n,) 减 (n,1)。你以为是逐元素相减得 (n,),广播规则却右对齐出 (n,n) 的全配对差值表——n=10000 时白造一亿个数,损失函数悄悄算错还不崩,指标诡异到让人怀疑人生,而代码评审时它看起来完全无辜。药方三帖:① 时刻区分 (n,)、(n,1)、(1,n) 三种形状,它们在广播眼里是三个物种;② 关键运算后 assert x.shape == (...),一行断言值十次 debug;③ sklearn/torch 报 warning 提示「shapes 不一致但可广播」时,别嫌烦,去看。另一个亲戚坑:view 与 copy——切片返回的常是视图,改它会穿透改到原数组;transpose 后内存不连续,某些操作需要 .copy()/.contiguous()。「我改了 A 怎么 B 也变了」,九成是视图在作祟。

广播的几何身世

广播不是工程 hack,它有干净的线代出身:「向量减标量」是减去 c·1(全 1 向量的倍数),「矩阵每行减行均值」是乘一个中心化投影矩阵(I − 𝟙𝟙ᵀ/n,第 10 章的投影兽!)。框架只是把这些「结构极规则的矩阵运算」优化成了不构造矩阵的逐元素循环——第 10 章「识货省钱」原则的系统级落实。所以每当你写下一个广播表达式,可以自问它对应哪个线性代数动作;答不上来的广播,往往就是埋 bug 的那个。

einsum:形状运算的正则表达式

当 @、transpose、sum 叠到三层嵌套时,有一个大杀器能把它们压成一行:einsum(爱因斯坦求和记号)。规则:给每根轴一个字母,写「输入轴 → 输出轴」的签名,没出现在输出里的字母自动求和:

einsum.py

np.einsum('ij,jk->ik', A, B)        # 矩阵乘:j 消失 = 沿 j 求和
np.einsum('i,i->', u, v)            # 点积:全消
np.einsum('i,j->ij', u, v)          # 外积:啥也不消
np.einsum('bld,bmd->blm', Q, K)     # 注意力分数:批量 QKᵀ,一行
np.einsum('blm,bmd->bld', W, V)     # 注意力加权求和,一行

五个例子从上到下正好是本书的进度复习:矩阵乘(07)、点积(04)、外积(07)、注意力分数(22 章预告)、加权求和(03 章配方)。它的可读性哲学与本书一脉相承:把「对哪根轴做什么」写在明面上,而不是让读者去脑补 transpose 链。einsum 还常常更快——它把多步操作交给一个优化过的内核,省掉中间张量的物化,又是一笔激活账。很多论文的开源实现(以及 einops 库的 rearrange)大量使用这种记号;会读它,等于拿到了一批高质量代码的入场券。顺便,你也真正理解了矩阵乘法公式里那个 Σⱼ——「重复下标求和」,爱因斯坦本人懒得写求和号,发明了这个记法,物理学与深度学习共用至今。

从公式到代码:一张翻译表

论文公式与张量代码之间有一张稳定的对照表,值得贴在显示器边上:

论文里写	代码里是	形状备注
Wx + b	`x @ W.T + b`	数据行优先约定,W 存转置(PyTorch Linear)
QKᵀ/√d	`q @ k.transpose(-2,-1) / sqrt(d)`	(B,H,L,Dh) → (B,H,L,L)
Σᵢ αᵢvᵢ	`attn @ v`	加权求和 = 矩阵乘,第 03 章配方
‖x‖²	`(x**2).sum(-1)`	沿特征轴归约
x/‖x‖	`x / norm(x, axis=-1, keepdims=True)`	keepdims 保广播
diag(γ)·x	`gamma * x`	对角乘 = 逐元素乘,第 10 章特权

注意一条贯穿性的差异:数学公式默认「单个列向量」,代码默认「一批行向量」——所以公式里的左乘 W,到代码里常变成右乘 W 的转置。第 06 章埋的这颗雷,现在有了完整的排雷手册:以 shape 咬合为准,公式只看语义。顺带安心丸一颗:NumPy、PyTorch、JAX、TensorFlow 的形状语法有九成同源(都是 NumPy 方言),本章的功夫在四个框架里通用,差异只在函数名(transpose vs permute)和自动微分的外壳。

两本内存账:参数与激活

形状思维的终极用途是算账,而且是面试和生产都常考的两本账。参数账(模型本身占多大):把每层的形状乘出来加总——一层 Transformer 约 12D²(四个注意力投影 4D² + 两个前馈 8D²),乘层数,再加词表的 V×D;7B 模型 × fp16(每参数 2 字节)= 14 GB,「我的显卡放不放得下」一乘便知;LoRA 只训 (m+n)r 个参数,同一笔算术。训练时还要给优化器翻倍再翻倍(Adam 的动量与方差各存一份 fp32),这就是「训练显存 ≈ 推理四倍起步」的出处。激活账(数据流经时的中间结果):一个 (32, 4096, 4096) 的注意力分数张量 × 32 头 × fp16 ≈ 34 GB——比模型还大!这就是为什么长上下文吃显存、为什么 FlashAttention 的核心贡献是「不物化那张 L×L 大表」。读论文时凡遇到「memory-efficient」,先问一句:它省的是参数账还是激活账?两本账,一套形状乘法,你现在都会算了。

激活账还教你一个反直觉事实:batch 越大不一定越快「per 样本」,因为激活内存随 B 线性涨,爆显存前的最后一步是各种「梯度累积」「激活重计算」(checkpoint:前向丢中间结果,反向时重算——拿算力换内存)。这些工程名词背后没有新数学,全是「哪些形状的张量、在哪个阶段、要不要留在内存里」的排班表。会算形状账的人,读得懂一切显存优化技术的动机。

▸ Android 类比

你和张量内存布局早就交过手:相机的 NV21/YUV_420_888 帧、Bitmap.copyPixelsToBuffer 的像素排布,本质都是「多维数据拍平进一维内存」的约定——和张量的行主序(C order)、以及深度学习框架的 NCHW vs NHWC 之争是同一件事:谁在内存里相邻,谁的批量操作就快(缓存友好,第 07 章的搬运账)。TFLite 在移动端偏爱 NHWC,正是因为逐像素的通道数据挨在一起,适配移动 GPU 的访存模式。你调 Bitmap 像素时踩过的 stride 坑,就是 .contiguous() 的前世。

◆ AI 连接

把「形状思维」升格成方法论:读任何一篇模型论文,最高效的姿势不是逐行啃公式,而是追踪张量形状的旅程——输入 (B, L, D) 进来,每经过一个模块问一次「形状变成了什么、哪根轴被混合了」:Linear 混 D 轴,注意力混 L 轴(token 之间交流),卷积混空间轴,batch 轴永远隔离(样本互不串门)。用这条线索读架构图,Transformer、CNN、Mamba 的设计差异会一目了然:它们的分歧,归根到底是「选择在哪根轴上、用什么方式让信息流动」。这句话请带进下一章。

最后一件小事,关于「读代码的姿势」。形状 bug 的调试成本与发现时间成指数关系:写下当场发现是十秒,跑完训练靠指标发现是十小时。所以老手的肌肉记忆是三件套常开:print(x.shape) 当呼吸、关键节点 assert 形状、复杂变换先用 (2, 3, 4) 这种迷你张量在 REPL 里排练一遍再上真数据。这些习惯朴素得不像技术,却是「会 NumPy」和「能维护千行模型代码」之间真正的分水岭。

✋ 动手自测

① (8, 1, 128) 和 (16, 128) 能广播吗?结果形状?(右对齐慢慢比);

② x.mean(axis=1) 和 x.mean(axis=1, keepdims=True) 对 (B, L, D) 各给出什么形状?哪个能直接拿去 x - mean?

③ 用 einsum 写「batch 内每条样本的向量两两点积」:输入 (B, N, D),输出 (B, N, N);

④ 13B 参数的模型,fp16 与 int4 各需多少存储?

⑤ (4096,) 的向量被误存成 (4096, 1) 后与 (4096,) 相减——结果形状是什么,为什么这是本章最贵的一道题?

三个五分钟小游戏,练出形状肌肉

形状思维是肌肉,不是知识,给三个日常可练的小游戏。游戏一「盲算」:读任何模型代码时遮住注释,只看操作链心算每一步的 shape,再对答案——正确率上 90% 之前别写自己的模型。游戏二「翻译」:随手拿一条论文公式(比如 LayerNorm 的定义),徒手写出带 axis 和 keepdims 的 NumPy——写完跑一下,用随机张量验证形状与数值。游戏三「找轴」:看到任何 AI 新名词,先问它动了哪根轴——「梯度累积」动 B 轴(拆小攒大),「滑动窗口注意力」动 L 轴(限制交流半径),「MoE」动 D 轴之后的专家轴(稀疏路由)。三个游戏共同的训练目标,是让「形状」从需要思考的东西,变成像缩进一样扫一眼就自动进脑的东西。

小结

深度学习的张量 = 有语义轴的多维数组,(B, L, D) 是最通用的世界观;reshape 换包装、transpose 换邻居、切片抽轴切段,reshape 冒充 transpose 是雪花图事故的祖师爷;归约吃轴、softmax 保轴,axis 写错语义全翻;广播右对齐、1 可扩、缺轴当 1,keepdims 与 None 插轴是两个日用惯用法,(n,) vs (n,1) 是最贵的血案;einsum 把轴操作写在明面上,是论文代码的通用方言;参数账和激活账两本账,一套形状乘法通吃;读架构 = 追踪「信息在哪根轴上流动」——Linear 混 D、注意力混 L、卷积混空间、batch 永远隔离;三个小游戏把这一切练成本能。装备齐了。下一章正式开拆:神经网络的前向传播——你会看到它从头到尾只有两种动作,而其中一种,你已经在这本书里足足练了十八章。