【AI Infra小白入门⑥】AI Infra数学基础速成：线代、概率与微积分（工程向，够用就行）
本文是「AI Infra 小白入门」系列的第六篇。很多同学一听到"数学"就头大——别担心！AI Infra不需要你推导复杂的论文公式，只要掌握"工程够用"的数学就行。本文用最通俗的语言，带你快速过一遍必备的数学知识。
一、写在前面：AI Infra需要多少数学？
先给大家吃颗定心丸：AI Infra对数学的要求，比算法研究低得多。
你不需要会：
- 复杂的定理证明
- 凸优化推导
- 高深的概率论
你只需要会：
- 心算矩阵维度和FLOPs（性能分析的基础）
- 理解几个核心概念（softmax、交叉熵、链式法则）
- 建立量化直觉（显存多大、带宽够不够、算力利用率多少）
一句话：AI Infra的数学是"工程数学"——够用就行，会算就好。
先看一张总览图：
二、线性代数：最最重要的部分
线性代数是AI Infra数学基础的"重中之重"。为什么？因为深度学习的核心计算就是矩阵乘法（GEMM），而所有性能分析都围绕着它展开。
2.1 矩阵乘法：维度与FLOPs
先记住这个公式：
(M, K) × (K, N) = (M, N)
两个矩阵相乘，第一个矩阵的列数必须等于第二个矩阵的行数。结果的行数是第一个矩阵的行数，列数是第二个矩阵的列数。
FLOPs怎么算？
FLOPs（Floating Point Operations，浮点运算次数）=2 × M × N × K
为什么有个"2"？因为每个乘加操作算两次运算（一次乘法 + 一次加法）。
举个例子：
- 一个(1024, 512)的矩阵乘以(512, 2048)的矩阵
- FLOPs = 2 × 1024 × 2048 × 512 ≈ 2.1 billion = 2.1 GFLOPs
💡为什么重要：GEMM是一切性能分析的基础单位。算力利用率 = 实际FLOPS / 峰值FLOPS。你连FLOPs都算不对，怎么评估性能？
2.2 转置与访存模式
矩阵乘法中，A矩阵是"按行读"，B矩阵是"按列读"。但在内存中，矩阵通常是**行主序（row-major）**存储的——也就是一行一行连续存的。
这就导致了一个问题：
- 读A的行：连续访问，效率高 ✅
- 读B的列：跳跃访问，效率低 ❌
这就是GEMM优化（tiling + shared memory）的数学根源——把B矩阵的一块先转到共享内存里，变成按行访问。
2.3 分块矩阵运算
矩阵可以"切"成小块来算：
C = [C11 C12] A = [A11 A12] B = [B11 B12] [C21 C22] [A21 A22] [B21 B22] C11 = A11·B11 + A12·B21分块之后，每个子块可以独立计算。
AI Infra中的应用：
- GPU的GEMM优化用了tiling分块思想
- 张量并行（TP）就是把权重矩阵按行或按列切分到多张卡上
- 理解了分块矩阵，就理解了TP为什么可行
三、概率论与统计：两个核心概念
概率论在AI Infra中用得不多，但有两个概念必须掌握：Softmax和交叉熵。
3.1 Softmax：把数值变成概率
Softmax的作用是把一组任意实数，变成"加起来等于1的正数"——也就是概率分布。
公式：
pi=ezi∑jezjp_i = \frac{e^{z_i}}{\sum_j e^{z_j}}pi=∑jezjezi
用人话讲：
- 每个数取指数（保证都是正数）
- 加起来求和
- 每个数除以总和（归一化）
工程问题：数值溢出
如果z_i很大，e^z_i会超级大，可能超出数值范围，导致溢出。
解决方案：每个数都减去最大值——softmax(z - max(z))
💡 这个技巧非常重要！FlashAttention中的"online softmax"就是它的流式版本。
推理中的采样
LLM推理时，得到logits后过softmax得到概率分布，然后采样下一个token：
- Greedy（贪心）：直接选概率最大的（argmax）
- Temperature：调整概率分布的"尖锐程度"
- Top-p / Top-k：只在概率最高的一部分token中采样
这些都是LLM Serving中必须理解的概念。
3.2 交叉熵损失：衡量预测有多准
交叉熵用来衡量"预测的概率分布"和"真实分布"之间的差距。
公式：
L=−∑iyilogpiL = -\sum_i y_i \log p_iL=−i∑yilogpi
在分类任务中，y_i通常是one-hot的（只有正确类别是1，其他都是0），所以公式简化为：L = -log(p_correct)
为什么要关注交叉熵？
因为Softmax和交叉熵联合求导有非常简洁的形式：dL/dz = p - y
正因为求导简洁，深度学习框架通常把Softmax和交叉熵**融合（fuse）**成一个算子——既减少了计算，又避免了数值不稳定。
💡 这就是为什么PyTorch里有
CrossEntropyLoss——它内部已经包含了Softmax，不用自己再写一遍。
四、微积分基础：链式法则就够了
4.1 链式法则 = 反向传播的全部
反向传播的本质就是链式法则在计算图上的系统执行。
公式：
∂L∂x=∂L∂y⋅∂y∂x\frac{\partial L}{\partial x} = \frac{\partial L}{\partial y} \cdot \frac{\partial y}{\partial x}∂x∂L=∂y∂L⋅∂x∂y
用人话讲：L对x的导数 = L对y的导数 × y对x的导数。一层一层"链"下去。
PyTorch的autograd做的就是这件事：把每个算子的局部梯度算好，然后沿计算图反向连乘，最终得到损失对每个参数的梯度。
4.2 梯度的工程含义
- 梯度指向损失"上升最快"的方向，所以优化器沿负梯度更新参数（梯度下降）
- 梯度也是张量，形状和参数完全一致
💡为什么重要：训练显存中，梯度占了很大一块。比如FP32训练，梯度就是4P字节（P是参数量）——这就是训练显存估算的基础之一。
五、工程速查表：这些公式你得会背
以下是AI Infra工程师的"常用公式"，建议记牢：
| 场景 | 公式 | 说明 |
|---|---|---|
| GEMM FLOPs | 2 × M × N × K | 矩阵乘法的计算量 |
| Attention FLOPs（每头） | 4 × N² × d | QK^T和softmax×V两部分 |
| 训练显存（混合精度Adam） | ≈ 16 × 参数字节 | 权重2B + 梯度2B + 优化器状态8B×2 + 激活… |
| Softmax稳定化 | softmax(z - max(z)) | 避免数值溢出 |
| AllReduce通信量（Ring） | ≈ 2S | S为总数据量，与rank数无关 |
🎯检验标准：看到一个模型参数、一个batch size，能大致心算出：
- 推理需要多少显存？
- 训练需要多少显存？
- 一层Attention的计算量大概是多少？
- 多卡训练通信量大概是多少？
能快速心算这些数字，是AI Infra工程师的"基本功"。
六、小白学习建议
建议一：不要死记公式，要理解"物理意义"
比如矩阵乘法，不要只记"行乘列"，要理解它的几何意义和工程含义——它是深度学习的基本运算单元，所有性能优化都围绕它。
建议二：多做"心算练习"
看到一个模型，就估算一下：
- 参数大概多少？
- 推理要多少显存？
- 训练要多少显存？
- 一层Attention的计算量是多少？
算多了，你就建立了"量化直觉"——看到一个数字，大概知道是大是小、是多是少。
建议三：不用学太深入，够用就行
AI Infra的数学是"工具"，不是"目的"。会用、会算、能解决工程问题就行。
如果你后面想深入某个方向（比如编译器的数学优化、分布式的理论分析），再补对应的数学知识也不迟。
七、总结一下
今天这篇文章，我们用最精简的方式过了一遍AI Infra必备的数学基础：
- 线性代数：矩阵乘法（维度+FLOPs）、转置与访存、分块思想——是性能优化和并行切分的根基
- 概率论：Softmax（数值稳定化）、交叉熵（融合优化）——理解推理采样和图优化
- 微积分：链式法则 = 反向传播 = autograd的全部
- 工程速查表：几个必须会背的公式，建立量化直觉
记住：AI Infra的数学是工程数学，够用就行，会算就好。
下一篇预告：我们来拆解Transformer架构——从Attention到FFN，从MHA到MLA，从工程视角看每个组件的计算量和优化点。
📚本系列导航：
- 第1篇：一文搞懂什么是AI基础设施：五大方向与不可能三角
- 第2篇：大厂AI Infra组织架构大揭秘
- 第3篇：AI Infra岗位路线图：级别、薪资与成长路径
- 第4篇：AI Infra学习路线图：从入门到精通的完整指南
- 第5篇：AI Infra必备编程基础：Python、C++与Linux速通指南
- 第6篇：AI Infra数学基础速成：线代、概率与微积分（工程向）（本文）
- 第7篇：Transformer架构工程拆解：从Attention到MLA的优化密码
- 更多内容持续更新中……
