交叉注意力融合

交叉注意力融合

本节摘要:投影层把一个图像向量与一个字幕向量对齐。真视觉语言解码器要每文本 token 能注意每 patch token,使模型能把每个词定位到区域。交叉注意力是定位发生的方式:文本查询、视觉键值回答。本节建交叉注意力块、因果文本自注意力、与使两者合法的掩码形状。把图像 token 与文本 token 拼成一序列是一种融合选项(早融合,Chameleon 与 Emu3 走的路);交叉注意力是另一种(晚融合,Flamingo 引入、之后每个 Flamingo 形解码器复制)。晚融合有两优:文本流干净保文本能力;图像流每图算一次、每解码步复用,使长字幕生成廉价。代价是每块多一注意力子层。

对应原课程:Phase 19 · Lesson 61 · cross-attention-fusion(原英文 phases/19-capstone-projects/61-cross-attention-fusion/docs/en.md)。本节属「多模态/VLM」赛道第四节。

学习目标

阅读完本节,你应当能够:

1. 实现查询流是文本、键值流是视觉的多头交叉注意力。
2. 组合解码器块:因果自注意力 + 交叉注意力 + 前馈。
3. 把掩码形状搞对:自注意力因果掩码、交叉注意力无掩码。
4. 跑批量文本 token + 固定图像 token 池的前向。

一、问题与直觉

把图像与文本 token 拼成一序列是一种融合选项(早融合,Chameleon 与 Emu3 走的路)。交叉注意力是另一种(晚融合,Flamingo 引入、之后每个 Flamingo 形解码器复制)。晚融合里,文本解码器跑在纯文本 token 上,每层经交叉注意力伸手进图像流。晚融合有两优:文本流干净保文本能力;图像流每图算一次、每解码步复用,使长字幕生成廉价。代价是每块多一注意力子层。

掩码形状

解码器块内两注意力需不同掩码:

本节含一形状校验函数,使混错两者的错误露为 ValueError 而非静默坏的损失曲线。

为什么交叉注意力无掩码

图像在任何文本生成前完全观测。字幕 token t 可注意任意图 patch,patch 无时序。一些 Flamingo 变体在交错多图与文本段时加每样本掩码模式,但单图加字幕时交叉注意力看一切。

键值缓存

图像键值在解码开始算一次持缓存。每新文本 token 用缓存不重算。这是字幕推理快的原因:重 ViT 跑一次,交叉注意力每步复用其键值。本节露缓存并测缓存命中路径。

块组合

解码器块跑:pre-LN → 自注意力 → 残差 → pre-LN → 交叉注意力 → 残差 → pre-LN → 前馈 → 残差。三子层,各自 LayerNorm。Flamingo 论文在交叉注意力加学习门使模型训练时稳定性代价下可退出图像路径;规范基线(本节用)无门。

class DecoderBlock:
  def forward(self, text_tokens, image_tokens, text_mask, cross_mask):
      text_tokens = text_tokens + self.self_attn(self.ln1(text_tokens),
                                                 mask=text_mask)
      text_tokens = text_tokens + self.cross_attn(self.ln2(text_tokens),
                                                  image_tokens,
                                                  mask=cross_mask)
      text_tokens = text_tokens + self.ffn(self.ln3(text_tokens))
      return text_tokens
​

二、从零实现

code/main.py 实现:

• CrossAttention(hidden, heads):多头交叉注意力,分立 q 与 kv 投影。
• CausalSelfAttention(hidden, heads):标准解码器的掩码自注意力。
• DecoderBlock:组三子层带 pre-LN 残差。
• VisionLanguageDecoder:四层解码器,喂 mock 视觉编码器输出与小文本嵌入表。
• causal_mask(length):返 (length, length) 下三角布尔张量。
• demo:喂两长 10 文本序列批 + 长 197 图像记忆,打输出形状、自注意力掩码形状、交叉注意力每位输出范数。

class CrossAttention(nn.Module):
    def __init__(self, hidden=512, heads=8):
        self.q_proj = nn.Linear(hidden, hidden)             # 文本 -> Q
        self.k_proj = nn.Linear(hidden, hidden)             # 图像 -> K
        self.v_proj = nn.Linear(hidden, hidden)             # 图像 -> V
        self.out_proj = nn.Linear(hidden, hidden)
        self.heads = heads; self.d = hidden // heads
    def forward(self, text, image, mask=None):
        B, Nt, _ = text.shape; Nv = image.shape[1]
        q = self.q_proj(text).view(B, Nt, self.heads, self.d).transpose(1, 2)
        k = self.k_proj(image).view(B, Nv, self.heads, self.d).transpose(1, 2)
        v = self.v_proj(image).view(B, Nv, self.heads, self.d).transpose(1, 2)
        scores = (q @ k.transpose(-2, -1)) / math.sqrt(self.d)   # 沿 Nv softmax
        attn = F.softmax(scores, dim=-1)
        out = (attn @ v).transpose(1, 2).reshape(B, Nt, -1)
        return self.out_proj(out)                            # 无掩码, 全图可见

class VisionLanguageDecoder(nn.Module):
    def __init__(self, vocab=1000, hidden=512, depth=4, heads=8):
        self.embed = nn.Embedding(vocab, hidden)
        self.blocks = nn.ModuleList([DecoderBlock(hidden, heads) for _ in range(depth)])
        self.lm_head = nn.Linear(hidden, vocab)
    def forward(self, text_ids, image_tokens):
        x = self.embed(text_ids)
        mask = causal_mask(text_ids.size(1))                 # 自注意力因果
        for blk in self.blocks: x = blk(x, image_tokens, mask, None)
        return self.lm_head(x)                               # 交叉无掩码
​

code/test_main.py 覆盖:交叉注意力输出形状;因果掩码下三角;缓存与未缓存路径 logits 等同;掩码混错抛 ValueError。

设计要点:交叉注意力分立 q 与 kv 投影是关键——q 从文本、kv 从图像,使文本「问」图像「答」。因果掩码只在自注意力,交叉无掩码(整图对每文本位可见,无时序)。键值缓存使推理快:重 ViT 跑一次,交叉注意力每步复用其 kv——这是字幕生成长序列廉价的原因。三子层(自注意/交叉/前馈)各 pre-LN 残差,与纯解码器(两子层)区别恰在交叉。

三、框架对比

交叉注意力在两生产族出现:Flamingo 与 IDEFICS(每 K 语言模型块插一交叉注意力子层,冻 LM,视觉语言适配器是交叉注意力块加门);BLIP-2(Q-Former 用从固定 32 查询 token 到图像特征的交叉注意力,再投到 LM 嵌空间)。本节块形状直映两者。掩码纪律(自因果、交叉无)相同。早融合(Chameleon/Emu3)把图像与文本 token 拼一序列、单自注意力处理,更简但文本流不纯、图像每步重算。LLaVA 用「投影后前置图像 token」的折中——图像 token 投影后作前缀拼到文本,单自注意力处理(早融合味),训练简但不如交叉注意力灵活。本节手写让你看清交叉的 q/kv 分离、掩码纪律、缓存复用——这些在用 Flamingo/IDEFICS 时是黑盒。

四、可复用产物

code/main.py + code/test_main.py。demo 在 CPU 上几十秒跑完:解码器产 (2, 10, text_vocab) logits,掩码形状 (10, 10),KV 缓存复用检查确认缓存与未缓存路径 logits 等同。CrossAttention、DecoderBlock、VisionLanguageDecoder 可独立复用——任何「文本查询图像键值」融合场景。第 60 节在对比 + LM 双目标预训里用此解码器。

五、练习

1. 掩码混错:故意把因果掩码传给交叉注意力(或反之),确认 ValueError 而非静默坏损失。
2. 缓存正确性:用缓存与不用缓存跑长 20 序列,断言 logits 逐位等同。
3. 门控交叉:给交叉注意力加学习 sigmoid 门(Flamingo 风格),观察训练稳定性。
4. 早融合对比:把交叉注意力换成「投影后图像 token 作前缀拼文本」早融合,对比参数量与吞吐。
5. 块间隔:每 2 或 4 语言模型块插一交叉注意力(Flamingo K=4),对比精度与算力。

本节要点回顾

1. 晚融合:文本解码器跑纯文本,每层交叉注意力伸手进图像流。
2. 晚融合两优:文本流干净保能力,图像每图算一次复用使长字幕廉价。
3. 交叉分立 q/kv:q 从文本(问)、kv 从图像(答),使词定位到区域。
4. 掩码纪律:自注意力因果(不可前看)、交叉无掩码(整图可见)。
5. 键值缓存:图像 kv 算一次持缓存,每文本 token 复用,推理快。
6. 三子层:自注意/交叉/前馈各 pre-LN 残差,与纯解码器区别在交叉。

下一节,我们做「视觉语言预训练」——InfoNCE 对比损失 + 自回归 LM 损失双目标,同时教模型排名与生成。