EECS498 · Transformer II · 10

Attention as Dynamic Information Routing

不要从公式背 Q/K/V。先要求 slept 真正取得 animal 的信息，再一步步推到 softmax(QKᵀ)V。

Step 1 · 先定义问题，不写公式

句子 The animal slept 查表并加入位置后，得到三个 vectors：

The      → x₀
animal   → x₁
slept    → x₂

若三个位置各自只通过 per-position MLP，x₂ 仍只来自 embedding(slept)+position(2)；同一套 MLP 不会读取 x₁。它可以把“slept”这个特征变复杂，却没有一条 edge 把本句主语搬过来。

Step 2 · 先接受“加权读取”

我们希望新的 slept 状态可以是其他位置所提供信息的组合：

new_slept
= 0.05 · information(The)
  + 0.80 · information(animal)
  + 0.15 · information(slept)

Step 3 · 从需求推出 Q、K、V

Q = XWQ       [N,K,D]
K = XWK       [N,K,D]
V = XWV       [N,K,Dv]

S = QKᵀ/√D    [N,K,K]     matching scores
A = softmax(S) [N,K,K]     dynamic read addresses
Y = AV         [N,K,Dv]    retrieved information

为什么同一个 X 要投影成三份？

Case A · Q=K=V=X 可以吗？

可以。这会形成一种合法 attention：原始 feature 同时承担匹配与传递。但它强迫“容易被找到的 feature”和“值得发送的 feature”完全相同。

Case B · Q=XWQ, K=XWK, V=XWV

独立 projections 让模型分别学习：如何提出查询、如何被匹配、实际发送什么。例如一个 token 可用“像名词”这一 feature 被匹配，却把数值、语义或位置信息的另一组合发出去。

Case C · 为什么 dot product，不是 Q+K？

加法本身只得到一条合成 vector，还需要另一个 function 才变成 scalar score。dot product 直接提供高维 feature matching，能以矩阵乘法并行计算所有 pairs，且可微分。它是高效 parameterization，不是数学宇宙中唯一可能的 attention；additive attention 也存在。

Attention matrix · 每个元素都说清楚

对 K=3 的句子，一张可能的 attention matrix 是：

reader ↓ / source →   The   animal  slept
The                    .70    .20     .10
animal                 .10    .80     .10
slept                  .05    .80     .15

row i = 谁在读；column j = 从谁读。所以 A[2,1]=0.80 表示输出位置 slept 从 value 位置 animal 读取 80%。它不表示 animal 从 slept 读，也不证明两者 embedding 相同。

为什么“信息交换”不是比喻？

设 values 是二维 vectors：

V₀ = [1,0]   V₁ = [0,2]   V₂ = [1,1]
A₂ = [.05, .80, .15]

Y₂ = .05V₀ + .80V₁ + .15V₂
   = [.05,0] + [0,1.60] + [.15,.15]
   = [.20,1.75]

V₁ 的第二个分量以 0.80×2 真实进入 Y₂。这不是“模型仿佛知道 animal”；它是可在 tensor 中逐项核对的 value mixing。

Your Code · 从两层循环到 batched matrix multiply

scores = torch.bmm(query, key.transpose(1, 2))
scores /= math.sqrt(M)
scores = torch.masked_fill(scores, mask, -1e9)
weights_softmax = F.softmax(scores, dim=-1)
y = torch.bmm(weights_softmax, value)

这段代码不仅是“vectorization trick”。它把 two-loop 中每一个 (i,j) dot product 汇总成 batch matrix multiply。Notebook 中 loop 与 no-loop 的小 relative error，是在验证两种执行方式表达同一机制。

Attention 与 MLP 各做哪种变化？

[N,K,M] → across K → [N,K,M]

M → M_hidden → M, repeated over K

“token mixing vs channel transformation”是好用视角，不是说模块绝对只做一种语义工作。关键结构事实是：标准 position-wise MLP 没有跨 K 轴读取；attention 有。

Multi-Head · 不只是“多个关系空间”

句子 The animal didn't cross the street because it was tired. 中，不同 heads 可能形成不同匹配：代词↔实体、动词↔主语、局部邻接、结构标记。这里是 conceptual intuition，不代表人为指定每个 head 的功能，也不保证训练后出现整洁的可解释分工。

M = H × Dh

per head Q,K,V      [N,K,Dh]
per head output     [N,K,Dh]
concat H heads      [N,K,H·Dh] = [N,K,M]
output projection   [N,K,M]

heads = [head(query, key, value, mask)
         for head in self.heads]
heads = torch.cat(heads, dim=-1)
output = self.linear_mult(heads)

End-of-Lesson Deep Checks