CS336 学习笔记

发表于 2026-07-23 分类于笔记阅读次数： Valine：

课程网站：https://cs336.stanford.edu/

Tokenization

文本的本质： Unicode字符串。

语言模型会对token序列建立概率分布。Tokenizer是完成encode和decode的类。

Token麻烦的点：

有时单词和空格+单词是不同的token
开头的单词和句中的单词token是不一样的
如何tokenize数字？

压缩比 = 字符串的字节数 / token数量

提升压缩比会增加效率，过度增加压缩比会造成sparsity的问题。词表数量增加-单个词出现的频率下降->数据稀疏->embedding层学习能力不充分->泛化差

如何构建分词器？

按照Unicode编码？词表巨大；有的字符利用率低，压缩比不理想。
首先将String转成字节。同上。
按照空格或正则来切：词表巨大，可能是无限的。

Byte Pair Encoding (BPE)

让罕见的序列拆分成多个token，同时让常见的序列拆分成一个token

基本思路：使用原始文本训练分词器

初始：每个字节都是一个token。

循环：将出现频率最高的相邻token对合并起来（原来的不会删除）。使用新的token更新整个训练语料库。重新统计pair

边界：直到词表达到设定的大小

BPE 训练算法

算法：BPE训练

输入：
    训练语料 C
    目标词表大小 V_target

输出：
    词表 V
    合并规则 M

1.  将训练语料 C 拆分为基本符号序列
2.  使用所有基本符号初始化词表 V
3.  初始化合并规则 M ← 空

4.  当 |V| < V_target 时，重复：

5.      统计 C 中所有相邻 Token 对的出现频率

6.      找到出现频率最高的 Token 对 (a, b)

7.      创建新 Token：
            new_token ← 合并(a, b)

8.      将 C 中所有相邻的 (a, b)
        替换为 new_token

9.      将 new_token 加入词表 V

10.     将合并规则 (a, b) 加入 M

11.  结束循环

12. 返回 V, M

BPE 编码算法

算法：BPE编码

输入：
    待编码文本 x
    合并规则 M

输出：
    Token序列 T

1.  将文本 x 拆分为基本符号，得到 Token 序列 T

2.  重复：

3.      找出 T 中所有相邻的 Token 对

4.      从这些 Token 对中，
        找出存在于合并规则 M 中的 Token 对

5.      如果不存在可以合并的 Token 对：
            结束循环

6.      根据合并规则 M 的优先级，
        选择优先级最高的 Token 对 (a, b)

7.      将 T 中相邻的 (a, b)
        合并成新的 Token ab

8.  结束循环

9.  返回 T

扩展方法：

只遍历重要的合并规则
处理特殊token
将文本分解为块，对每个块应用分词器

展望

未来可能不需要Tokenization，而是实现从字节开始的端到端。

但是仍然需要满足：

transformer依赖某种序列的chunks运行。

Chunks应该是大小可变的。并非所有字节都同等重要。

Resource Accounting

要回答的问题：在给定的资源（计算量与内存）下，可以训练最好的模型是什么。

Memory Accounting

Tensor Memory

tensor是储存数据的基本单元。

1
2
3

a = torch.zeros(4) # vector
b = torch.zeros(4, 8) # matrix
c = torch.zeros(4, 8, 2) # rank-3 tensor

通常储存浮点数

float32: fp32 单精度，用于传统科学计算，可训练小模型
float64: 双精度
float16: 半精度，动态范围小，训练时不稳定
bfloat16：动态范围和float32一样。

混合精度训练：

参数、激活值、梯度：bf16
优化器状态：fp32
Pytorch的AMP库可以实现

甚至有fp8，fp4

Compute Counting

Tensor Einops

Einops 是用来操作张量的库，命名了每个维度。

Ensum

Ensum 爱因斯坦求和约定：一套统一的方法完成很多tensor运算。

约定：

字母表示Tensor的维
输入中存在、但输出中消失的维度，会被求和。因此："i,i->"$= \sum_i a_ib_i$
输出中字母顺序决定输出Tensor的维度顺序

可以用...前面所有剩余的维度

1
2
3

x = np.random.rand(2, 3, 4, 5)  # 想把最后两维 4和5 转置为 5和4
y = np.einsum('...ij->...ji', x)
print(y.shape)  # (2, 3, 5, 4)

矩阵乘法：

1	torch.einsum("ik,kj->ij", A, B)

使用的三个步骤：

给输入维度贴标签

Q → b h q d

K → b h k d
看输出

output → b h q k
寻找消失的字母

常用einsum

运算	einsum
向量求和	`"i->"`
向量逐元素乘法	`"i,i->i"`
向量点积	`"i,i->"`
外积	`"i,j->ij"`
矩阵转置	`"ij->ji"`
矩阵所有元素求和	`"ij->"`
每行求和	`"ij->i"`
每列求和	`"ij->j"`
矩阵乘法	`"ik,kj->ij"`
Batch 矩阵乘法	`"bik,bkj->bij"`

Reduce

对消失的维度做：

"sum"     # 求和
"mean"    # 平均
"max"     # 最大值
"min"     # 最小值
"prod"    # 连乘
reduce(x, "b s h -> b h", "mean") # 例子

Rearrange

重新排列、拆分或合并 Tensor 的维度，但不会对数据进行求和等计算。

顺序：括号中从左到右排列，右边的维度变化最快。

1 2	rearrange(x, "batch seq hidden -> (batch seq) hidden") # 合并 rearrange(x, "batch seq (head d) -> batch head seq d", head=8) #拆分

Tensor Operations Flops

FLOP：浮点运算，主要是加法乘法

FLOPs 浮点运算次数

FLOP/s (FLOPS)：浮点运算性能

H100的性能：1979 teraFLOP/s（实际的要除以2）

对于 (B,D) @ (D,K) 的矩阵乘法，$FLOPs = 2BD*K$（乘一次加一次）

另一种理解方式：$FLOPs = 2(tokens)(parameters)$

MFU (Model Flops utilization): 实际FLOPS/承诺FLOPS（除二过的）

FLOPS由硬件和数据类型决定

Arithmetic Intensity

计算时间决定因素：

加速器的速度 (FLOP/s)
内存带宽 (bytes/s)

二者可以并行

bottleneck：

Memory-bound
Compute-bound

arithmetic intensity = flops / bytes

acclelerator intensity = flop/s / bytes / s

Roofline polts：横轴是计算密度，纵轴是加速器峰值计算能力。

训练时的计算量

反向传播的计算量：

$h_1.grad = h_2.grad \times W_2^T$

$W_2.grad = h_1^Th_2.grad$

是前向传播的两倍

总计算量：

Forward pass: 2 (# data points) (# parameters) FLOPs

Backward pass: 4 (# data points) (# parameters) FLOPs

Total: 6 (# data points) (# parameters) FLOPs

在上下文不长时对Transformer也使用

结论

当模型一共L层，每层权重$W\in \mathbb{R}^{D\times D}$，一次处理B个token，则有

$flops = 6BN，where N = D^2L$

减少内存积累的方法：

梯度积累 gradient accumulation：
激活检查点 activation checkpoint：将部分没保存的激活重新计算（$\sqrt L$ 个检查点比较好）

Architectures

大模型结构：

输入文字

tokenID

Enbedding词向量（可训练）

很多层 Transformer Block：

RMSNorm (缩放参数可训练)
多头自注意力机制
- 生成KQV（$W_{QKV}$可训练）
- 对QK做RoPE
- 计算注意力分数
- Causal Mask
- 加权汇总V
残差连接
RMSNorm (缩放参数可训练)
MLP （权重可训练，偏置一般不使用）
残差连接

最终归一化

LM Head （参数可训练）

Logits

采样或选最大值

得到下一个Token

新Token加到最后面

重复这个过程

归一化

RMS Norm没有明显的表达能力损失，不需要计算均值中间量更少，运算更快，减少了内存搬运。

$\operatorname{RMS}(x) = \sqrt{\frac{1}{d}\sum_{i=1}^{d}x_i^2+\varepsilon}$ $\mathrm{RMSNorm}(x) = \frac{x}{\operatorname{RMS}(x)} \odot \gamma$

$x$：输入向量（Hidden State）
$d$：Hidden Size
$\varepsilon$：防止除零的小常数（不可训练）
$\gamma$：可训练的缩放参数（与 Hidden Size 等长）
$\odot$：逐元素乘法

现代模型使用Prenorm，保证残差主干干净

$x_{l+1} = x_l + F(\text{Norm}(x_l))$

偏置项

实践证明可以去掉减少内存搬运，避免不稳定

激活函数

现代模型使用：SwiGLU/GeGLU

$\text{GeGLU}(x) = (x W) \otimes \text{GELU}(x V )$ $\text{SwiGLU}(x) = \text{Swish}(x V) \otimes (x W)$ $\text{Swish}_1(x) = x \cdot \text{sigmoid}(x) = \frac{x}{1 + e^{-x}}$ $\text{GELU}(x) = x \cdot P(X \le x) = x \cdot \Phi(x)$

如果要给MLP加上门控结构，需要把原来的W的输出维度缩小到原来的2/3：

原来的参数量：

升维层 $W1$ 参数量：$d \times h{\text{old}}$
降维层 $W2$ 参数量：$h{\text{old}} \times d$
总参数量：$2 \cdot d \cdot h_{\text{old}}$

加门的参数量：

路径 A ($W{\text{up}}$) 参数量：$d \times h{\text{new}}$
路径 B ($W{\text{gate}}$) 参数量：$d \times h{\text{new}}$
降维层 ($W{\text{down}}$) 参数量：$h{\text{new}} \times d$
总参数量：$3 \cdot d \cdot h_{\text{new}}$

保持参数量不变： $h_{\text{new}} = \frac{2}{3} h_{\text{old}}$

RoPE

传统的绝对位置编码的问题：无法自然表示Token之间的相对距离。

旋转位置编码，通过在二维空间中旋转向量来注入位置信息。

注入 token 位置信息，让模型识别token的绝对位置、相对距离、前后顺序。

二维向量旋转：

$R(\theta)=\begin{bmatrix}\cos\theta & -\sin\theta\\\sin\theta & \cos\theta\end{bmatrix}$ $x_1'=x_1\cos\theta-x_2\sin\theta,\quad x_2'=x_1\sin\theta+x_2\cos\theta$

高维向量旋转：将高维向量两两拆分为二维小组，每组独立执行旋转。

位置 m 处，第 i 组的旋转角度：$m\theta_i, \theta_i = \frac{1}{10000^{2i/d}}$ ，其中 d 为单头注意力维度

$\theta_i$较小：旋转角度随位置m变化缓慢，捕捉长距离语义依赖。

$\theta_i$较大：旋转角度岁位置变化速度快，捕捉相邻词之间的局部依赖。

实现：

$f(\mathbf{x}, m) = \mathbf{x} \odot \cos(m\Theta) + \tilde{\mathbf{x}} \odot \sin(m\Theta)$ $\tilde{\mathbf{x}} = [-x_2, x_1, -x_4, x_3, \dots, -x_d, x_{d-1}]$

优势：

作用在单向量上，便于 KV Cache 缓存；表达相对位置。

长文本外推性好，配合NTK-Aware Scaling 和 YaRN，只对低频维度压缩，扩展上下文长度的代价更小。

随着相对距离 $\vert{}m - n\vert{}$ 增加内积的期望值会自动衰减，即词汇关联度降低，符合语言学规律。

Parallel Layers

传统串行：

$y = x + \text{MLP}(\text{LayerNorm}(x + \text{Attention}(\text{LayerNorm}(x))))$

并行：

$y = x + \text{MLP}(\text{LayerNorm}(x)) + \text{Attention}(\text{LayerNorm}(x))$

可能导致残差相加时的方差积累稍快，主流仍然是用串行，保证数值稳定性。

超参数

一些经验法则：

$d{ff}=4d{model}$，算上GLU之后的比率是2.67左右，也可以取3.5来侧重MLP。

$d{head} \approx d{model}/n_{head}$

$d{model}/n{layer} \approx 100$