做多模态模型这几年我经常被问到同一个问题“图像和文本拼成一个序列之后VLM里的注意力到底应该怎么掩码”看似是细节但这个问题直接决定多模态信息能不能正确融合、训练稳不稳定、生成时会不会“剧透”。我觉得掩码注意力在VLM里的地位相当于整个信息流动的路由表。这篇就把它拆开讲清楚从原理到代码再到工程坑一次性说明白。1. 掩码注意力在VLM里的定位它是信息流动的路由表1.1 为什么VLM必须显式设计掩码视觉语言模型VLM的标准做法是把图像切成一堆patch每个patch过视觉编码器变成一个token然后把视觉token和文本token拼成一个序列再丢进Transformer自回归解码器里。这一步看起来简单但问题马上来了Transformer的自注意力机制默认是“全对全”的也就是每个token都能看到序列里的其他所有token。训练时如果你不拦着文本token在预测下一个字的时候就会提前看到“正确答案”这跟单模态语言模型不做因果掩码是一个道理模型直接退化成作弊机器。更麻烦的是图像不是一维序列。文本天然有顺序从左到右、从前到后因果掩码每个位置只能看之前的位置完全契合语言生成。但图像是二维的patch之间没有严格的先后顺序。左上角的patch和右下角的patch谁先谁后强行套一个一维因果掩码会让图像内部的注意力变成一条链早期patch根本看不到远处的视觉信息局部细节交互也被切断图像理解能力大打折扣。所以VLM的掩码不能照搬语言模型的掩码必须把“模态身份”和“序列位置”两个维度同时考虑进去。掩码矩阵里的每一个0和-∞都是在定义一种信息流动规则。这也是为什么不同模型的掩码策略千差万别——你希望图像内部怎么交互、文本怎么参考图像、图像能不能往后看文本全在这个矩阵里体现。1.2 掩码在注意力公式里的准确位置讲清楚原理之前先回到那个所有人都会用到的公式Attention(Q, K, V) softmax(QK^T / sqrt(d) M) V这里M就是掩码矩阵。它是一个和注意力分数矩阵形状一样的矩阵通常由0和-∞组成。0表示允许这个位置的注意力计算-∞表示禁止。softmax之前把-∞加进去经过softmax之后那一个位置的概率就变成0相当于这个query彻底忽略掉对应的key。有一个细节很多人第一次接触时会困惑为什么不用NaN而用-∞因为NaN在反向传播里会传染全部梯度-∞加进softmax里是一个干净的数学操作数值稳定梯度也能正常回传。这里还要区分一个概念掩码和Padding Mask不是一回事。Padding Mask解决的是“序列里的无效填充位”让token不去关注填充内容而掩码注意力解决的是“有效token之间的可见关系”比如因果、双向、模态隔离。工程上两者通常会合成一个最终的attention bias但设计思路要分开。2. 三种基础掩码模式从矩阵出发理解行为差异2.1 全因果掩码全因果掩码就是语言模型里标准的“下三角”掩码。假设序列长度是N掩码矩阵M满足M[i][j] 0允许当且仅当 j iM[i][j] -∞屏蔽当 j i这个矩阵的含义是第i个位置只能看到自己以及前面的位置绝对不能看后面。对于纯文本生成任务这是唯一正确的选择因为训练时的目标就是根据前面的token预测下一个token未来的信息必须隔绝。但在VLM里如果对拼接后的整个序列直接套全因果掩码前面说的图像二维信息破坏问题就会出现。有些早期方案确实这么干过把图像token当成一段“特殊的文本”处理训练起来倒是省事但效果普遍不如显式设计掩码的模型。我自己的实验里纯因果掩码在细粒度图像描述任务上经常出现“看到了某个物体却叫不出颜色/相对位置”的情况——这就是因为图像内部的远距离patch交互被一维顺序切断了。2.2 全双向掩码双向掩码就是没有掩码所有token可以互相看。矩阵里全是0attention计算变成最原始的全局注意力。在VLM中使用全双向掩码通常只出现在对比学习预训练或者Image Encoder内部因为那里不需要生成每个token都可以参考全图的上下文计算出来的表征更全局。但是一旦到了生成阶段全双向掩码就有问题。文本token如果能看到未来token模型学到的就不是“按顺序生成语言”而是“从答案里抄答案”。所以生成式VLM的主体架构里全双向掩码只会出现在特定模块比如视觉编码器不会出现在联合解码器里。2.3 自定义掩码真正干活的部分VLM里真正核心的是自定义掩码根据视觉token区间和文本token区间把矩阵划分成不同的块每个块单独决定“允许”还是“禁止”。假设图像token在前1到N_v文本token在后N_v1到N_vN_t整个掩码矩阵按坐标轴划分成四个象限查询\键图像token区文本token区图像token区图像内部可见性图像是否能看到文本文本token区文本是否能看到图像文本内部可见性每一种VLM的掩码设计本质上就是回答这四个象限的问题。你可以让图像内部双向、文本内部因果也可以让图像和文本互不相通甚至可以细化到让图像内部的某一块patch只能看所属区域的其他patch。这块是真正的设计空间。3. 从零实现VLM里的掩码注意力模块3.1 输入准备与序列拼接直接上可运行的思路。假设我们已经有一个视觉编码器输入一张图得到N_v个视觉token同时有一个文本分词器得到N_t个文本token。在模型内部两者拼成一个序列import torch # 模拟视觉token和文本token num_v 196 # 例如14x14的patch数量 num_t 32 # 文本长度 hidden_dim 768 visual_tokens torch.randn(1, num_v, hidden_dim) text_tokens torch.randn(1, num_t, hidden_dim) sequence torch.cat([visual_tokens, text_tokens], dim1) # [1, num_vnum_t, hidden_dim]拼接后的总长度N num_v num_t。接下来所有掩码构造都以这个总长度为基准。3.2 用向量化方式构造三类掩码构造掩码最忌讳用Python循环一格格填序列一长就慢到离谱。正确做法是用torch.triu这类矩阵操作生成布尔掩码再转成attention bias。下面给出三种模式的构造代码def build_mask_2d(num_v, num_t, modepartial_causal, dtypetorch.float32): N num_v num_t # 默认全保留 mask_bool torch.zeros(N, N, dtypetorch.bool) if mode full_causal: # 整个序列严格下三角屏蔽右上三角 mask_bool torch.triu(torch.ones(N, N, dtypetorch.bool), diagonal1) elif mode partial_causal: # 视觉区内部双向 # 视觉区禁止看文本区图像行、文本列屏蔽 mask_bool[:num_v, num_v:] True # 文本区内部因果文本行、文本列右上三角屏蔽 mask_bool[num_v:, num_v:] torch.triu( torch.ones(num_t, num_t, dtypetorch.bool), diagonal1 ) elif mode modal_isolation: # 模态隔离视觉和文本互相完全不可见 mask_bool[:num_v, num_v:] True mask_bool[num_v:, :num_v] True # 每个模态内部保持双向 # True的位置表示屏蔽转成attention bias里的 -inf attn_bias torch.zeros(N, N, dtypedtype) attn_bias.masked_fill_(mask_bool, float(-inf)) return attn_bias这段代码里mask_bool是布尔掩码True表示“禁止”。转成attn_bias之后直接加到QK^T / sqrt(d)上就能用。之所以用布尔掩码中转是因为调试的时候布尔矩阵一眼就能看出哪块是通的、哪块是断的比看一堆-∞直观得多。3.3 在多头注意力里应用掩码有了attention bias接下来就是标准的注意力计算。为了展示全流程我写一个不依赖任何第三方库的多头注意力实现重点看掩码如何被消费import torch.nn.functional as F class MaskedMultiHeadAttention(torch.nn.Module): def __init__(self, hidden_dim, num_heads): super().__init__() self.num_heads num_heads self.head_dim hidden_dim // num_heads self.wq torch.nn.Linear(hidden_dim, hidden_dim) self.wk torch.nn.Linear(hidden_dim, hidden_dim) self.wv torch.nn.Linear(hidden_dim, hidden_dim) self.wo torch.nn.Linear(hidden_dim, hidden_dim) def forward(self, x, attn_biasNone): batch, seq_len, _ x.shape H, D self.num_heads, self.head_dim q self.wq(x).view(batch, seq_len, H, D).transpose(1, 2) k self.wk(x).view(batch, seq_len, H, D).transpose(1, 2) v self.wv(x).view(batch, seq_len, H, D).transpose(1, 2) scores q k.transpose(-2, -1) / (D ** 0.5) # [batch, H, seq, seq] if attn_bias is not None: # attn_bias 需要扩展成 [batch, H, seq, seq] scores scores attn_bias.unsqueeze(0).unsqueeze(0) attn_weights F.softmax(scores, dim-1) out attn_weights v out out.transpose(1, 2).reshape(batch, seq_len, -1) return self.wo(out), attn_weights有几个细节必须注意。第一attn_bias的形状要广播成[batch, H, seq, seq]我这里直接unsqueeze两次如果batch和head维度不同需要expand。第二掩码必须加在softmax之前加在softmax之后没有任何屏蔽效果。第三如果一行全被屏蔽attention bias整行都是-∞softmax会输出NaN实际场景里要保证每行至少有一个有效位置。这个坑我后面会展开说。3.4 三种模式下注意力权重的可视化验证代码写完了建议你做一件事构造一个简单输入把三种掩码下的注意力权重矩阵拿出来打印一下确认形状和自己预期一致。mha MaskedMultiHeadAttention(hidden_dim768, num_heads8) x torch.randn(1, num_v num_t, 768) for mode in [full_causal, partial_causal, modal_isolation]: bias build_mask_2d(num_v, num_t, mode) _, attn_weights mha(x, bias) # 取第一个batch、第一个head的平均权重 avg_weights attn_weights[0, 0].detach().mean(dim0) print(mode, avg_weights.shape) # 期望都是 [228, 228]full_causal的权重矩阵应该是一个标准下三角形状partial_causal里文本区是下三角视觉区接近全连接但视觉行看不到文本列modal_isolation则是左上和右下两个方块。如果形状不对第一时间检查拼接顺序和索引边界。4. VLM掩码的进阶设计模态隔离与局部视觉掩码4.1 模态隔离在什么时候用有一种常见误解模态隔离掩码只适合做表征学习不适合生成。这话对了一半。在生成式VLM的联合解码阶段模态隔离确实会阻碍文本参考图像效果很差。但在预训练早期模态隔离是一个非常好用的稳定器。为什么因为视觉token和文本token初始时不在同一个语义空间如果一开始就让两者全量交互模型很容易学歪典型的症状是loss下降快但下游任务效果差因为模型学会了“抄近道”——从图像特征里直接猜测文本模式而没有真正建立模态间的对齐。先用模态隔离让各模态内部先形成稳定的表征再逐步放开跨模态交互可以明显降低早期训练的不稳定性。我见过的一个做法是把模态隔离作为warmup阶段开头几千步用隔离掩码之后切回部分因果掩码。4.2 图像内部的双向到底意味着什么回到partial_causal模式里的视觉区。视觉token彼此之间双向可见本质是允许任意patch直接参与其他patch的表征计算。这跟ViT在视觉编码器里的做法一致只不过这里发生在跨模态解码器里。图像内部全连通带来一个好处细粒度关系建模能力更强。比如描述“桌子左边的红色杯子”模型需要同时参考杯子所在patch、红色相关patch、桌子patch才能正确组织语言。代价是计算量比一维因果链大一些但视觉token数量通常远小于文本token数量这个代价完全可控。实践中我通常建议视觉区保持双向不要为了统一形式而强行套因果掩码否则视觉理解能力会明显退化。4.3 区域掩码让VLM学会只看某个局部区域比模态隔离更细一层的是区域掩码。VLM面对的问题不只是“图像和文本怎么交互”还有“图像内部的注意力怎么聚焦”。比如指代表达式任务里用户问“左上角那个蓝色方块是什么”模型应该重点看左上区域的patch而不是全图一视同仁。这时候可以用区域掩码。大致思路是图像被切成N_v个patch后每个patch有一个坐标位置。当你希望模型关注某个区域时给这个区域内的patch之间的注意力偏置设为0区域外的patch设为-∞或一个较大的负值。这样注意力矩阵里视觉token的可见范围被限制在一个矩形区域内。def build_region_mask(num_v_rows14, num_v_cols14, region(0, 7, 0, 7)): # 假设图像是14x14的patch网格 r0, r1, c0, c1 region num_v num_v_rows * num_v_cols mask_bool torch.zeros(num_v, num_v, dtypetorch.bool) def patch_index(r, c): return r * num_v_cols c for i in range(num_v_rows): for j in range(num_v_cols): if not (r0 i r1 and c0 j c1): # 区域外的patch屏蔽所有关注 idx patch_index(i, j) mask_bool[idx, :] True mask_bool[:, idx] True return mask_bool这段代码把区域外patch的所有注意链路都掐断只保留区域内部的交互。实际使用中不一定要这么“硬”给区域外一个小负值比如-0.1做软掩码也能起到聚焦作用同时保留部分全局上下文。软硬掩码的选择我建议先从硬掩码开始调确认区域定位准确后再尝试软掩码看效果有没有进一步提升。5. 工程落地最容易踩的五个坑5.1 Mask与Flash Attention的兼容性问题很多人跑到工程阶段发现自己写的掩码在torch.nn.functional.scaled_dot_product_attention里有效一换FlashAttention就报错或者掩码不生效。原因是FlashAttention这种高效注意力内核默认只支持padding mask这类“按key侧屏蔽”的掩码不支持任意形状的二维自定义注意力偏置。这是内核设计决定的不是你代码问题。解决方案有三条路。第一继续用scaled_dot_product_attention至少顺序模式支持任意掩码但效率比Flash低。第二改用支持attention bias的高效注意力库比如一些性能优化库里的memory_efficient_attention接口能够接受自定义attention bias。第三把自定义掩码转换成可以由上游计算实现的偏置——思路是给不同区域token的向量上乘一个可扰动项让注意力分数本身变小等效于掩码的软版本。我一般按需混用序列长度短用方案一序列长且掩码形态规整用方案二掩码本身需要参与训练时用方案三。5.2 KV Cache下掩码生长的问题推理阶段为了省算力会缓存历史token的K和V新token只需要跟缓存里的K计算注意力。很多人在训练时把掩码写对了一上线推理就蒙了当前query是新tokenkey缓存里都是历史token根本不存在“未来的token”上三角掩码是不是不需要了答案是KV Cache推理阶段通常不需要重新设置因果掩码的上三角部分因为key列表里已经不存在未来的位置。但你仍然需要一个“不完整掩码”作用对象是padding位置。如果batch内序列长度不一致短序列的padding部分仍然需要用-∞屏蔽否则模型会拿padding当有效信息。这个掩码会在每次decode后动态缩短因为它对应当前query长度1和当前有效key长度的关系。我建议把“train mask”和“inference mask”分成两个函数维护不要共用一个逻辑。5.3 全-∞行导致NaN的问题这是我最常见的线上事故。当你把attention mask和padding mask合并时一个query行的所有key位都被屏蔽softmax分母为0输出NaN。特征一旦变NaN后面所有层全部崩溃loss直接变成NaN训练过程彻底报废。排查办法很简单在每次forward后加一行断言检查attention weights是否存在NaN。更根本的解法是构造掩码时保证每一行至少留下一个有效位置。如果你的padding策略是右侧padding文本token右侧的大量-∞加上因果掩码右侧的-∞会让末尾的某些行碰巧交集为空一定要人工检查边界位置。5.4 掩码矩阵索引错位的问题视觉token拼接在前还是后、文本token在前还是后不同模型风格完全不同。有的模型把图像放在前面有的交错放置有的把系统提示放最前面。掩码构造代码里的num_v和num_t一旦和实际序列错位出错的不是崩溃而是静默的性能下降——模型还能跑但注意力完全乱了。我习惯的做法是在构造掩码后立刻用一个固定的小示例打印掩码矩阵人工对照区块结构。视觉区是连续的、文本区的位置和实际索引一致、图像看不到文本、文本能看到图像这四条检查完再进训练循环。别等到训练跑了两天后才怀疑掩码错了。5.5 注意力头“单头垄断”问题不加约束的VLM训练后期常出现某些注意力头熵值特别低几乎只看固定一两个位置其他头退化。这在掩码设计不当的时候会加剧尤其是区域掩码太硬把太多交互通道掐死模型只好依赖极少数头去做所有事。应对手段有两个。一是在损失里加一个注意力熵正则项鼓励注意力分布不要过于集中。二是根据头的编号分配不同掩码——比如4个头用全图双向视觉掩码4个头用区域聚焦掩码让模型既有全局视角又有局部聚焦能力。第二种做法在细粒度视觉定位任务上提升明显但训练时要多观察各头的熵曲线防止掩码分配后某些头几乎没有梯度。6. 掩码模式的实验对比与个人经验6.1 我建议的对比方案如果从零开始设计VLM不要一上来就搞花哨的自定义掩码。先用最基础的三种模式做一组对比实验模式训练任务预期优势预期风险full_causal图像文本整体自回归实现最简单图像内部信息交互受限partial_causal文本自回归图像内部双向图像理解文本生成兼顾多一个掩码分支要维护modal_isolation partial_causal切换分阶段训练训练更稳定切换时机需要调参把顺序跑完记录下游VQA任务准确率、文本生成困惑度、训练总时长三个指标。从我实际做过的几个项目来看partial_causal几乎总是占优full_causal在纯视觉理解任务上会掉点而分阶段切换对最终指标提升的帮助则要看数据量数据量小的时候帮助不明显。6.2 软掩码和注意力偏置的一点心得掩码不一定只能是非0即-∞。把-∞换成一个大负数比如-10000效果上接近硬掩码但保留极少信息通道换成-1就变成了软掩码对注意力分布是一个惩罚而不是彻底关断。软掩码的好处是梯度能够平滑地穿过整个注意力矩阵训练早期更不容易陷入局部最优。我常用的一个技巧是“线性衰减掩码”图像token离当前query的欧氏距离越远attention bias负值越大但不断言为-∞。这个做法让模型既优先看近处patch也不会彻底丧失远处上下文。代价是需要额外维护位置距离矩阵推理时需要预计算。如果只是普通任务直接用硬掩码就够了别为了技巧而技巧。6.3 调试掩码问题的一条完整思路最后分享一个我调试掩码相关问题的常规流程。遇到VLM训练效果异常时不要急着改模型结构按顺序排查第一步打印attention weights的熵值分布看是否出现大范围0熵或NaN第二步检查掩码矩阵的区块结构确认视觉区和文本区的索引边界第三步关闭掩码跑一个小规模过拟合实验看模型是否具备基本记忆能力第四步恢复掩码继续过拟合实验如果loss明显变差拿掩码矩阵可视化逐块对比。第四步通常能直接暴露问题。曾经有个项目我们以为是学习率问题折腾了一周没解决最后发现是视觉区有一块patch索引写错导致部分图像区域的注意力永久关闭。那次之后我养成了一个习惯掩码矩阵的单元测试必须写进项目一开始的代码框架里而不是之后补。掩码注意力机制说穿了不复杂它就是一个控制信息可见性的矩阵。但放在VLM的语境里它把“图像二维结构”“语言因果性”“模态对齐策略”这三个问题拧在了一起设计好坏直接决定模型能做多细粒度的事。我个人的体会是不要默认拿纯文本模型的掩码套路去套VLM也不要觉得掩码只是训练细节。把它当一个独立的设计模块来对待先画矩阵、再写代码、最后做消融整个模型的质量会上一个台阶。