Zhanbo's Blog
Back to home

CS546 Learning Notes:从 Word Embedding 到 RNN,理解 BPTT 与梯度消失

Learning Notes/Algorithm
2026-09-12
14 min read
0 reads
AIMachine Learning

In brief

上次把 Softmax、Cross-entropy 和一个 ReLU MLP 的前向、反向传播手算了一遍。这次继续问:如果输入是一句话,文字怎么变成数字?句子长度不一样、词序不一样,模型又怎么处理? 今天走通的主线是:token → embedding → RNN hidden state → 下一个词的预测 → loss → 沿时间反向传播。最需要停下来推的地方,是 embedding 梯度里的转置,以及为什么 RNN 的梯度传得越远可能越小。

1. Embedding matrix:每个词对应一行

先假设每个单词就是一个 token。真实 tokenizer 不一定按完整单词切分,这里用单词只是为了方便手算。

词表顺序为:

[I,like,cats,dogs,you][\text{I},\text{like},\text{cats},\text{dogs},\text{you}]

token ID 从 0 开始,所以 cats 的 ID 是 2。词表大小为 V=5V=5,每个 embedding 有 d=3d=3 个数字:

ER5×3E\in\mathbb R^{5\times3}

例如:

E=[0.20.10.70.50.30.20.40.80.10.60.50.90.00.20.3]E=\begin{bmatrix} 0.2&-0.1&0.7\\ 0.5&0.3&-0.2\\ -0.4&0.8&0.1\\ 0.6&-0.5&0.9\\ 0.0&0.2&-0.3 \end{bmatrix}

Embedding lookup 就是根据 token ID 取出对应行:

E[2]=[0.4, 0.8, 0.1]E[2]=[-0.4,\ 0.8,\ 0.1]

这里是索引 2,也就是写出来的第 3 行。

句子 I like cats 有 3 个 token,查表后得到 3×33\times3 的矩阵。两个 3 分别代表 token 数量与 embedding dimension,只是恰好相同。

输入Shape维度含义
整张 embedding matrixV×dV\times d词表大小、向量维度
一个长度为 TT 的句子T×dT\times d序列长度、向量维度
BB 个补齐到长度 TT 的句子B×T×dB\times T\times dbatch、序列长度、向量维度

因此 2 个句子、每句补齐到 4 个 token、每个向量长度为 3,输入 tensor 就是 2×4×32\times4\times3。这是常见的 batch-first 排列约定。

2. One-hot 乘矩阵,就是 embedding lookup

cats 的 one-hot 写成行向量:

x=[0,0,1,0,0]R1×5x=[0,0,1,0,0]\in\mathbb R^{1\times5}

则:

x1×5E5×3=e1×3=E[2]\underbrace{x}_{1\times5}\underbrace{E}_{5\times3} =\underbrace{e}_{1\times3} =E[2]

其他行都乘以 0,只有 cats 那行乘以 1。程序通常直接查表,不需要真的创建 one-hot 再做大量包含零的乘法。

Embedding matrix 通常是可学习参数:后面的预测产生 loss,梯度反向传回来,再更新对应的词向量。

3. 为什么 embedding 的梯度是 xTgex^Tg_e

当时看到这个转置,我又有点不明白:

这里 xx 的转置是怎么求出来的?能给我推一遍吗?

先从一个元素开始。前向计算为 e=xEe=xE,把第 jj 个输出分量展开:

ej=i=15xiEije_j=\sum_{i=1}^{5}x_iE_{ij}

这里公式下标使用从 1 开始的数学记法,因此 cats 对应 i=3i=3;它的程序索引仍是 2。

设后面传回来的梯度为行向量:

ge=Le=[g1,g2,g3]R1×3g_e=\frac{\partial L}{\partial e} =[g_1,g_2,g_3]\in\mathbb R^{1\times3}

其中 gj=L/ejg_j=\partial L/\partial e_j。参数 EijE_{ij} 只直接影响 eje_j,而:

ejEij=xi\frac{\partial e_j}{\partial E_{ij}}=x_i

链式法则给出:

LEij=LejejEij=gjxi\frac{\partial L}{\partial E_{ij}} =\frac{\partial L}{\partial e_j} \frac{\partial e_j}{\partial E_{ij}} =g_jx_i

把所有元素排回矩阵:

LE=[x1g1x1g2x1g3x2g1x2g2x2g3x5g1x5g2x5g3]=xTge\frac{\partial L}{\partial E} =\begin{bmatrix} x_1g_1&x_1g_2&x_1g_3\\ x_2g_1&x_2g_2&x_2g_3\\ \vdots&\vdots&\vdots\\ x_5g_1&x_5g_2&x_5g_3 \end{bmatrix} =x^Tg_e

Shape 也吻合:

(5×1)(1×3)=5×3(5\times1)(1\times3)=5\times3

这叫外积。转置的根据是逐元素求导得到的 xigjx_i g_j;shape 是检查结果的工具。直接写 xgexg_e 时,中间维度 515\neq1,连矩阵乘法都做不了。

例如 ge=[2,1,3]g_e=[2,-1,3],则:

LE=[000000213000000]\frac{\partial L}{\partial E} =\begin{bmatrix} 0&0&0\\ 0&0&0\\ 2&-1&3\\ 0&0&0\\ 0&0&0 \end{bmatrix}

若 learning rate 为 η=0.1\eta=0.1,普通 Gradient Descent 更新为:

E[2]new=[0.4,0.8,0.1]0.1[2,1,3]=[0.6,0.9,0.2]\begin{aligned} E[2]_{\text{new}} &=[-0.4,0.8,0.1]-0.1[2,-1,3]\\ &=[-0.6,0.9,-0.2] \end{aligned}

我第一次把第一个数写成了 0.60.6。实际上 0.40.2=0.6-0.4-0.2=-0.6;第二个数增加,是因为减去负梯度:0.80.1(1)=0.90.8-0.1(-1)=0.9

“只更新被查到的行”在这里有条件:EE 仅用于输入 lookup,使用普通 Gradient Descent,且没有正则化等其他更新来源。若有 weight decay、优化器历史状态,或与输出层共享 embedding 权重,就不能简单断言其他行一定不变。

如果同一个词出现两次,两次 lookup 使用的是同一行参数,因此两处传回来的梯度贡献相加。

4. 下一个词预测的正确答案从哪里来?

我当时的问题是:

以前有确定的标签,这里要怎么给答案?

先区分一下:回归任务的标签通常是数值;one-hot 是之前分类任务里的标签表示。语言模型的下一个 token 预测,也可以看成词表上的分类任务。

原始文本本身就提供了答案。对于:

[I,like,cats,today][\text{I},\text{like},\text{cats},\text{today}]

把序列错开一位:

时间步当前输入 token已看到的上下文Target
1IIlike
2likeI likecats
3catsI like catstoday

所以训练用的 input 是 [I, like, cats],target 是 [like, cats, today]。第二步虽然只新读入 like,但 hidden state 还携带前面的 I 的信息。

这种从数据本身构造监督信号的方式叫 self-supervised learning。Target 可以概念性地写成 one-hot,程序里通常直接保存 token ID。

也可以加入 <BOS><EOS> 表示序列开始、结束:

Input<BOS>Ilikecatstoday
TargetIlikecatstoday<EOS>

特殊 token 的具体使用方式依赖模型与数据处理约定,并不是所有实现都必须使用这两个名字。加入新 token 时,词表大小也要相应计算。

训练时,这个例子每一步喂入原文中的正确 token,称为 teacher forcing。生成时,则把前一步选出的 token 喂回下一步。RNN 的 hidden state 依然需要按时间顺序计算;知道全部标签不等于能并行算出所有时间步的 hidden state。

5. Embedding 相似,不是人工规定每个数字的含义

Embedding 中每个坐标通常没有人工指定的单独语义。向量通过完成训练目标逐步调整,经常出现在相似上下文中的词,可能学到相似表示。

例如 catsdogs 的用法经常相近,但这是一种训练结果的趋势,并非由一两个句子保证。

尤其要区分输入 embedding 与输出层权重:在 I like → cats 这条样本中,如果两者没有共享参数,输入 lookup 用到的是 Ilike,不会因为 target 是 cats 就直接查找并更新输入端的 cats 那一行。输出层则会通过词表上的分类梯度学习预测 cats

因此,“相似上下文对应相似 embedding”不能直接从这一个样本推出,还要看完整数据、训练目标与参数共享方式。比较词向量时常用 cosine similarity,它主要衡量方向是否相近。

6. 为什么不能把一句话的 embedding 直接相加?

比较:

  • dog bites man
  • man bites dog

我最开始觉得它们相加后的表示会不同,但向量加法满足交换律:

edog+ebites+eman=eman+ebites+edoge_{\text{dog}}+e_{\text{bites}}+e_{\text{man}} =e_{\text{man}}+e_{\text{bites}}+e_{\text{dog}}

对于这种固定的 token embedding,直接求和会丢掉词序。按词序把向量排成矩阵,才能在输入表示中保留这种差异。

如果每个 embedding 长度为 3,三个词展平后长度为 9,五个词展平后长度为 15。同一个普通 MLP 的第一层权重 shape 固定,不能直接接收这两种长度。

Padding 到固定长度可以配合 MLP 使用,也能通过固定位置保留词序,但需要预先处理长度范围。RNN 提供了另一种方式:每次读一个 token,把之前的信息传给下一步。

7. RNN:当前输入加上之前的 hidden state

基础 RNN 的计算是:

ht=tanh(Wxxt+Whht1+bh)h_t=\tanh(W_xx_t+W_hh_{t-1}+b_h)

从这一节开始,xtx_t 表示当前 token 的 embedding 列向量,不是前面 lookup 推导中的 one-hot 行向量。若 EE 每行存一个词向量,则 xt=E[IDt]Tx_t=E[\mathrm{ID}_t]^T

设 embedding dimension 为 d=3d=3,hidden dimension 为 H=4H=4

符号Shape含义
xtx_t3×13\times1当前 token 的 embedding
ht1h_{t-1}hth_t4×14\times1上一步、当前的 hidden state
WxW_x4×34\times3将输入映射到 hidden space
WhW_h4×44\times4将历史状态映射到当前计算
bhb_h4×14\times1偏置

三项都得到 4×14\times1,因此可以相加;tanh\tanh 逐元素作用,不改变 shape。

每个时间步共享同一组 Wx,Wh,bhW_x,W_h,b_h。序列变长会增加计算步数,但不会因此增加这组参数的数量。Hidden state 是历史信息的学习表示,不保证完整记住所有前文。

8. 手算验证:RNN 为什么能感知顺序?

暂时去掉 tanh\tanh,把输入和 hidden state 都简化成标量:

ht=2xt+0.5ht1,h0=0h_t=2x_t+0.5h_{t-1},\qquad h_0=0

输入顺序是 [1,3][1,3]

h1=2(1)+0.5(0)=2h_1=2(1)+0.5(0)=2 h2=2(3)+0.5(2)=7h_2=2(3)+0.5(2)=7

交换顺序,输入 [3,1][3,1]

h1=2(3)+0.5(0)=6h_1=2(3)+0.5(0)=6 h2=2(1)+0.5(6)=5h_2=2(1)+0.5(6)=5

相同的两个输入,最终状态分别为 7 和 5。较早的输入会经过 recurrent connection,和当前输入对最终状态的作用不同,因此顺序可以影响结果。

这说明 RNN 能表达词序差异,不代表任何参数设置下都一定能区分所有不同序列。

9. 已经能区分顺序,为什么还要 tanh\tanh

我想到之前的 ReLU:全是线性层的 MLP,不管叠多少层,都能合并成一个线性变换。RNN 也有类似的问题。

去掉激活函数和偏置:

h1=Wxx1+Whh0h_1=W_xx_1+W_hh_0 h2=Wxx2+Whh1=Wxx2+WhWxx1+Wh2h0\begin{aligned} h_2 &=W_xx_2+W_hh_1\\ &=W_xx_2+W_hW_xx_1+W_h^2h_0 \end{aligned}

对于固定步数,它仍是输入和初始状态的线性组合;保留偏置时则是仿射变换。虽然有顺序信息,但这段状态计算无法表达非线性关系。

加入 tanh\tanh 后会出现嵌套:

h2=tanh(Wxx2+Whtanh(Wxx1+Whh0+bh)+bh)h_2=\tanh\left(W_xx_2+W_h\tanh(W_xx_1+W_hh_0+b_h)+b_h\right)

这通常不能再合并成一个线性层。

tanh\tanh 还把每个 hidden 分量限制在 (1,1)(-1,1) 内。在线性例子 ht=2ht1h_t=2h_{t-1}h0=1h_0=1 中,状态会变成 2,4,8,2,4,8,\ldots;经过 tanh\tanh 的 hidden state 不会这样无限增长。但前向状态有界,不意味着反向梯度一定稳定。

10. 从 hidden state 输出下一个词

输出层接回了上次学过的 Softmax 分类:

zt=Woht+bo,pt=softmax(zt)z_t=W_oh_t+b_o,\qquad p_t=\operatorname{softmax}(z_t)

H=4H=4、词表大小 V=5V=5

Wo5×4ht4×1+bo5×1=zt5×1\underbrace{W_o}_{5\times4}\underbrace{h_t}_{4\times1} +\underbrace{b_o}_{5\times1} =\underbrace{z_t}_{5\times1}

ptp_t 同样是 5×15\times1,每个元素是对应词的概率,所有元素非负且总和为 1。

沿用五词词表,如果 target 是 cats

pt=[0.1,0.1,0.4,0.3,0.1]T,yt=[0,0,1,0,0]Tp_t=[0.1,0.1,0.4,0.3,0.1]^T,\qquad y_t=[0,0,1,0,0]^T

单个位置的 Cross-entropy 是:

Lt=logpt(cats)=log0.4L_t=-\log p_t(\text{cats})=-\log0.4

梯度为:

Ltzt=ptyt=[0.1,0.1,0.6,0.3,0.1]T\frac{\partial L_t}{\partial z_t}=p_t-y_t =[0.1,0.1,-0.6,0.3,0.1]^T

正确词的概率越高,loss 越小。把 logits 看作独立变量时,负梯度方向会提高正确词的 logit、降低其他词的 logit。但实际训练更新的是共享参数,不能保证整个模型或一个 batch 更新后,每个样本的 logits 都严格按这个方向变化。

多个位置的目标可以相加:

L=t=1TLtL=\sum_{t=1}^{T}L_t

也可以取平均;取平均时,对应总梯度也要除以 TT。有 padding 时,通常只对有效 target 计算 loss。

11. BPTT:沿时间展开后,仍然是链式法则

假设只看三个时间步,总 loss 为:

L=L1+L2+L3L=L_1+L_2+L_3

h1h_1 会影响 L1L_1,还会通过 h2h_2 影响 L2L_2、通过 h2,h3h_2,h_3 影响 L3L_3。因此:

Lh1=L1h1+L2h1+L3h1\frac{\partial L}{\partial h_1} =\frac{\partial L_1}{\partial h_1} +\frac{\partial L_2}{\partial h_1} +\frac{\partial L_3}{\partial h_1}

h2h_2 只收到 L2,L3L_2,L_3 的贡献。对基础单向 RNN,前向图中没有从 h2h_2 返回去影响 L1L_1 的路径。

这就是 Backpropagation Through Time(BPTT):把 RNN 沿时间展开,再使用普通反向传播。

在标量情况下,来自 L3L_3 的一条路径为:

L3h1=L3h3h3h2h2h1\frac{\partial L_3}{\partial h_1} =\frac{\partial L_3}{\partial h_3} \frac{\partial h_3}{\partial h_2} \frac{\partial h_2}{\partial h_1}

仍然是熟悉的规则:沿路径相乘,多路径相加。

向量情况下不能随意交换矩阵顺序。定义 Jt=ht/ht1RH×HJ_t=\partial h_t/\partial h_{t-1}\in\mathbb R^{H\times H},并把 loss 对状态的梯度写成列向量,则从 L3L_3 传回 h1h_1 的贡献是:

h1L3=J2TJ3Th3L3\nabla_{h_1}L_3=J_2^TJ_3^T\nabla_{h_3}L_3

12. 共享参数为什么要把梯度相加?

同一个 WhW_h 在不同时间步被重复使用。反向传播时,各次使用对总 loss 的贡献需要累加,最终更新同一个参数。

为了把“每步的贡献”写清楚,定义激活前的量:

at=Wxxt+Whht1+bha_t=W_xx_t+W_hh_{t-1}+b_h

δt=L/atRH×1\delta_t=\partial L/\partial a_t\in\mathbb R^{H\times1},它已经包含当前及后续 loss 传来的贡献,则:

LWh=t=1Tδtht1T\frac{\partial L}{\partial W_h} =\sum_{t=1}^{T}\delta_t h_{t-1}^T

每项 shape 为 (H×1)(1×H)=H×H(H\times1)(1\times H)=H\times H。这里与 embedding 一样,又出现了外积。

在这次完整序列反向传播结束后,普通 Gradient Descent 更新:

WhWhηLWhW_h\leftarrow W_h-\eta\frac{\partial L}{\partial W_h}

不会因为展开了三个时间步,就变成三个独立学习的 WhW_h。同样,也不要把已经包含后续影响的梯度再重复计入。

13. 梯度消失与爆炸:连续相乘的后果

先用最简单的标量 RNN:

ht=wht1h_t=wh_{t-1}

每一步的局部导数都是 ww。假设末端 loss 对 h3h_3 的梯度为 1,那么:

L3h0=1×w3\frac{\partial L_3}{\partial h_0}=1\times w^3

w=0.5w=0.5

0.53=0.125,0.5100.000980.5^3=0.125,\qquad 0.5^{10}\approx0.00098

w=2w=2

210=10242^{10}=1024

前者展示 vanishing gradient,后者展示 exploding gradient。这里用 h0h_0 观察信号跨越时间后的变化,并不表示固定为零的初始状态必须作为参数更新。

梯度消失意味着,后面 loss 对早期计算的学习信号可能非常弱,因此长期依赖难学。早期位置如果还有自己的 loss,仍可能收到较强的局部梯度;消失的是我们关注的长距离贡献,不是所有梯度都必然为零。

真实的 tanh RNN 中:

Jt=htht1=diag(1ht2)WhJ_t=\frac{\partial h_t}{\partial h_{t-1}} =\operatorname{diag}(1-h_t^2)W_h

这里 ht2h_t^2 表示逐元素平方,diag\operatorname{diag} 把这些激活导数排到对角线上。整个 JtJ_t 的 shape 是 H×HH\times H

因为:

tanh(a)=1tanh2(a)\tanh'(a)=1-\tanh^2(a)

hth_t 接近 111-1,导数接近 0,这叫进入 saturation(饱和区)。梯度经过这些方向时容易被压小。

最后的手算题是:末端梯度为 1,连续五步的局部导数都是 0.2:

1×0.25=0.000321\times0.2^5=0.00032

这次算对了。一个看起来不算特别小的局部导数,重复相乘后也会让远处的学习信号迅速变弱。

矩阵情况取决于 Jacobian 连乘,不能只看 WhW_h 某个元素是否大于 1。虽然每一步共享 WhW_h,但激活导数随状态变化,因此各时间步的完整 Jacobian 通常不同。

ZB

Zhanbo Chen

Java Backend & AI Agent Developer

Back to home
Comments