Update post of Assignment3 in cs224n which has finished the code in nmt_model.py

Author: ottercoconut

content/zh/post/cs224n-a3/index.md

@@ -214,6 +214,8 @@ $$
 
 ### Implementation and written questions
 
+#### `__init__()`
+
 在问题 (c) 的`nmt_model.py`的`__inti__()`中，关于神经网络各层的定义，其中`self.att_projection`的定义需注意，虽然它对应的是 $W_{attProj}$ ，但是在编写代码时，要拆解结合律，决定计算顺序。
 
 应该先算$\mathbf{W}_{\text{attProj}} $ 和 $\mathbf{h}_i^{\text{enc}}$ 相乘，再将其乘积与 $(\mathbf{h}_t^{dec})^T$ 相乘求最终结果，这是因为所有的编码器状态 $\mathbf{h}_i^{\text{enc}}$ 在 Decoder 开始工作前就已经全部计算好了。你可以**一次性**把整句话的 $\mathbf{h}^{\text{enc}}$ 丢进 `Linear` 层进行投影（这叫 Pre-computation 预计算）。但是在某一时刻只有一个 $\mathbf{h}_t^{dec}$ 而没有未来的量，导致不能一次性投影完成，而只能把线性层放到了循环中，引发严重的性能问题（虽然数学原理上没错）。下面是一个例子(generated by Gemini) :
@@ -243,6 +245,8 @@ $$
 > - **计算量**：我们直接拿解码器原生的 $1 \times 512$ 向量，去乘以准备好的“投影后字典”（$512 \times 50$）。只需要进行 $512 \times 50 = \mathbf{25,600}$ 次乘加运算！
 > - **没有额外的层开销**：这里不需要调用 `nn.Linear`，我们在代码里只要写一个极其轻量的纯矩阵乘法指令（比如 `torch.matmul` 或 `@` 符号）就搞定了。不需要加载任何权重矩阵，因为权重矩阵的任务在循环外已经完成了！
 
+#### `encode`
+
 问题 (d) 是encoder部分，以及对decoder部分的初始化：
 
 ```python
@@ -289,5 +293,81 @@ init_decoder_hidden = self.h_projection(torch.cat((last_hidden[0],last_hidden[1]
 
 `last_cell`同理，不再赘述。
 
+#### `decode`
+
+问题 (e) 是decode部分：
+`enc_hiddens_proj = self.att_projection(enc_hiddens)` 对应公式 (6) 中 $\mathbf{W}_{\text{attProj}} \mathbf{h}_i^{\text{enc}}$ 的部分。正如在 `__init__()` 部分讨论的，编码器的全部隐藏状态在进入 `decode` 时已经固定，所以可以**在循环外一次性**完成投影预计算，将 `enc_hiddens` 从 `(b, src_len, 2h)` 投影为 `(b, src_len, h)`。
+
+`Y = self.model_embeddings.target(target_padded)` 和 `encode` 中对源语言做查表的操作对称，这里调用**目标语言**的嵌入矩阵，将 `target_padded` 中的词索引映射成词向量，得到形状为 `(tgt_len, b, e)` 的张量 `Y`。
+
+`torch.split(Y, 1)` 沿第 0 维（时间维度）将 `Y` 切成一系列 `(1, b, e)` 的张量。`squeeze(Y_t, 0)` 去掉多余的维度变为 `(b, e)`。注意**必须指定 `dim=0`**，否则当 `batch_size = 1` 时会误删 batch 维度。
+
+`torch.cat((Y_t, o_prev), dim=1)` 对应前文的拼接操作：将 $\mathbf{y}_t$（维度 $e$）与 $\mathbf{o}_{t-1}$（维度 $h$）拼接为 $\overline{\mathbf{y}_t} \in \mathbb{R}^{(e+h)}$。
+
+`self.step(...)` 内部一次性完成公式 (5)-(11) 的计算：Decoder LSTM 前向传播 → 注意力评分与分布 → 加权求和得到上下文向量 → 线性投影、tanh、dropout，最终产出联合输出向量 $\mathbf{o}_t$。每步将 $\mathbf{o}_t$ 存入列表并更新 `o_prev`。
+
+`torch.stack(combined_outputs, dim=0)` 将列表中所有 `(b, h)` 的张量堆叠为 `(tgt_len, b, h)`，随后送入公式 (12) 的词汇投影层生成概率分布。
+
+#### `step`
+
+问题 (f) 是step部分，即解码器单步计算，内部完成公式 (5)-(11)。核心代码分两段：
+
+**第一段（~3行）：Decoder LSTM 前向 + 注意力评分**
+
+```python
+dec_state = self.decoder(Ybar_t, dec_state)
+dec_hidden, dec_cell = dec_state
+e_t = torch.bmm(enc_hiddens_proj, torch.unsqueeze(dec_hidden, 2)).squeeze(2)
+```
+
+前两行直接对应公式 (5)：将 $\overline{\mathbf{y}_t}$ 和上一步状态送入 Decoder LSTM，得到新的 `dec_state`，再拆分为 `dec_hidden`（$\mathbf{h}_t^{\text{dec}}$）和 `dec_cell`（$\mathbf{c}_t^{\text{dec}}$），形状均为 `(b, h)`。
+
+第三行对应公式 (6) 中的注意力评分。这里的关键在于 `torch.bmm` 的形状要求——它执行**批量矩阵乘法**，要求输入严格为三维张量 `(b, n, m)` × `(b, m, p)` → `(b, n, p)`。而我们手里的张量：
+
+- `enc_hiddens_proj`：形状 `(b, src_len, h)` — 已经是三维，无需处理
+- `dec_hidden`：形状 `(b, h)` — 只有二维，不满足 `bmm` 的要求
+
+所以需要 `torch.unsqueeze(dec_hidden, 2)` 在第 2 维（最后）插入一个维度，将 `(b, h)` 变为 `(b, h, 1)`。这样 `bmm` 的乘法就变成了：
+
+$$\underbrace{(b, \text{src\_len}, h)}_{\text{enc\_hiddens\_proj}} \times \underbrace{(b, h, 1)}_{\text{dec\_hidden}} = \underbrace{(b, \text{src\_len}, 1)}_{e_t}$$
+
+这实质上就是对 batch 内的每一条样本，让 `enc_hiddens_proj` 的每一行（某个源词的投影）与 `dec_hidden` 做点积，得到该源词的注意力分数——正是公式 (6) 的 $(\mathbf{h}_t^{\text{dec}})^T \mathbf{W}_{\text{attProj}} \mathbf{h}_i^{\text{enc}}$。
+
+最后 `.squeeze(2)` 去掉末尾多余的维度 `1`，从 `(b, src_len, 1)` 变回 `(b, src_len)`，得到注意力得分向量 $\mathbf{e}_t$。
+
+**第二段（~6行）：注意力加权 → 联合输出**
+
+```python
+alpha_t = F.softmax(e_t, dim=1)
+a_t = torch.bmm(torch.unsqueeze(alpha_t, 1), enc_hiddens).squeeze(1)
+U_t = torch.cat((a_t, dec_hidden), dim=1)
+V_t = self.combined_output_projection(U_t)
+O_t = self.dropout(torch.tanh(V_t))
+```
+
+`F.softmax(e_t, dim=1)` 对应公式 (7)，沿 `dim=1`（即 `src_len` 维度）做 softmax，将分数归一化为注意力分布 $\alpha_t$，形状仍为 `(b, src_len)`。
+
+计算上下文向量 $\mathbf{a}_t$（公式 (8)）时又遇到了 `bmm` 的三维要求。`alpha_t` 是 `(b, src_len)`，需要 `unsqueeze(alpha_t, 1)` 在第 1 维插入，变为 `(b, 1, src_len)`：
+
+$$\underbrace{(b, 1, \text{src_len})}_{\alpha_t} \times \underbrace{(b, \text{src_len}, 2h)}_{\text{enc_hiddens}} = \underbrace{(b, 1, 2h)}_{a_t}$$
+
+这就是用注意力权重对编码器隐藏状态做加权求和。`.squeeze(1)` 去掉中间的 `1`，得到 `(b, 2h)` 的上下文向量 $\mathbf{a}_t$。
+
+后三行依次对应公式 (9)(10)(11)：`torch.cat` 拼接 $\mathbf{a}_t$ 与 $\mathbf{h}_t^{\text{dec}}$ 得到 `(b, 3h)` 的 $\mathbf{u}_t$；线性层投影回 `(b, h)` 的 $\mathbf{v}_t$；最后 tanh + dropout 得到联合输出 $\mathbf{o}_t$。
+
+#### Attention Masking
+
+问题 (g) 是关于 `step()` 中注意力掩码的作用。在 `step()` 的两段代码之间，有这样一段：
+
+```python
+if enc_masks is not None:
+    e_t.data.masked_fill_(enc_masks.bool(), -float('inf'))
+```
+
+`enc_masks` 由 `generate_sent_masks()` 生成：它创建一个 `(b, src_len)` 的零矩阵，然后对 batch 中每条句子，在其**实际长度之后**的位置全部填 `1`。换言之，`1` 标记的是 `<pad>` token 的位置，`0` 标记的是真实词的位置。
+
+**掩码对注意力计算的影响：** `masked_fill_` 将 $\mathbf{e}_t$ 中所有对应 `<pad>` 位置的注意力分数替换为 $-\infty$。当这些 $-\infty$ 的值随后经过 `softmax` 时，$e^{-\infty} = 0$，因此 `<pad>` 位置的注意力权重 $\alpha_{t,i}$ 会变为 $0$，而所有真实词的权重之和仍归一化为 $1$。这意味着在公式 (8) 的加权求和中，`<pad>` 位置的编码器隐藏状态对上下文向量 $\mathbf{a}_t$ 完全没有贡献。
+
+**为什么必须这样做：** 由于 batch 内各句子长度不一，短句会被 `<pad>` 填充到统一长度。如果不加掩码，`<pad>` 对应位置的编码器隐藏状态（本质上是无意义的噪声）会分走一部分注意力权重，从而污染上下文向量，导致翻译质量下降。