Kornia SigLip2 注意力掩码极性修复深度解析：SDPA 布尔掩码语义与文本特征 NaN/全零问题

• 计算机视觉
• 深度学习
• 图像处理

导读

本文聚焦 Kornia 仓库中 SigLip2 视觉-语言模型的注意力掩码（attention_mask）极性反转修复（对应 changelog.d/+migration-120.fixed.md，issue #4043）。修复前，SigLip2Attention在将掩码交给torch.nn.functional.scaled_dot_product_attention（SDPA）前错误地取了反，导致"想 attend 的位置被 mask 掉、想 mask 的位置反而被 attend"，最终SigLip2.get_text_features(..., attention_mask=...)在 PyTorch 2.5 及以后返回全零文本特征、在更早版本返回nan。读完本文，你将掌握 SDPA 布尔掩码的真实语义、三种掩码维度的正确构造方式、该 bug 的版本相关症状，以及修复后与 Hugging Facetransformers对齐的 key 轴广播行为，并能直接照抄出可用的推理代码。

一、问题本质：SDPA 布尔掩码的语义与直觉相反

torch.nn.functional.scaled_dot_product_attention接受两种掩码形式：

• 布尔形式：True表示attend（参与注意力），False表示mask out（屏蔽）；
• 数值形式：float("-inf")表示屏蔽，0表示参与。

而在用户侧约定（如 Hugging Facetransformers的习惯）里，attention_mask通常是1/True表示有效 token，0/False表示 padding。也就是说，用户侧掩码与 SDPA 布尔掩码的语义在"取反"关系上是同向的：用户给1的位置，SDPA 布尔掩码也要给True才能 attend。

Kornia 的SigLip2Attention旧实现在把用户掩码转交给 SDPA 之前，对 2-D、3-D、4-D 三个分支的掩码都执行了一次取反（~mask或等价操作），于是语义被彻底颠倒：

• 每一个调用者希望 attend 的位置都被 mask 掉；
• 每一个 padding 位置反而被 attend；
• 当传入全 1 掩码（即不 padding 任何位置）时，每一行 query 的所有 key 位置全被屏蔽，整行变成"fully masked row"。

二、版本相关症状：0、NaN，还是看似正常？

修复文档明确指出，该 bug 的可见症状依赖 PyTorch 版本，因为 SDPA 对"全掩码行"的处理在版本间不一致：

这意味着在 Kornia 声明的torch>=2.0下限（见 pyproject.toml）范围内，这个 bug 几乎必然以"特征异常"的形式被观测到：要么全零、要么 NaN。对下游任务（图像-文本匹配、零样本分类、检索）而言，这两者都会静默地破坏排序结果——因为归一化后的全零向量与其他向量点积恒为 0，NaN 则会污染任何包含它的相似度计算。

三、修复实现：三个分支的掩码构造逐行解析

修复后的SigLip2Attention.forward位于 kornia/models/siglip2/attention.py，其掩码处理逻辑对三种维度分别做了不再取反的转换，统一成 SDPA 需要的(B, 1, N, N)布局（关键代码见 attention.py#L120-L137）：

attn_mask = None if attention_mask is not None: if attention_mask.dim() == 2: # (batch_size, seq_len) -> (batch_size, 1, 1, seq_len) # 只 mask key 轴：padding 的 query 行不会产生 fully masked row attn_mask = attention_mask.bool().unsqueeze(1).unsqueeze(2) elif attention_mask.dim() == 3: # (batch_size, seq_len, seq_len) -> (batch_size, 1, seq_len, seq_len) attn_mask = attention_mask.bool().unsqueeze(1) elif attention_mask.dim() == 4: # 已经是 SDPA 布局 (batch_size, 1, seq_len, seq_len)；True = attend attn_mask = attention_mask.bool() else: raise ValueError(f"Unsupported attention_mask dimension: {attention_mask.dim()}")

三种掩码形态的语义与来源如下：

• 2-D(B, N)：经典的 padding 掩码，由分词/批处理阶段生成，1表示有效 token、0表示 padding。修复后经bool().unsqueeze(1).unsqueeze(2)变为(B, 1, 1, N)——注意广播位置在key 轴（最后一维），而非 query 轴；
• 3-D(B, N, N)：逐 query 的二维掩码（例如因果掩码、自定义注意力模式），经unsqueeze(1)补上 head 维度，广播到所有注意力头；
• 4-D(B, 1, N, N)：已经是 SDPA 原生布局，直接bool()即可。

其中 2-D 分支的处理是该修复的核心设计决策：只在 key 轴广播。这样一条 query 行只在自己对应的 mask 位为 0 时才会被屏蔽；而 padding 序列中"query 位置虽然属于 padding、但 key 轴仍能看到有效 token"的行，不会整行被屏蔽，也就不会触发 SDPA 的 fully masked row 行为。

四、与 Hugging Facetransformers的对齐：create_bidirectional_mask

修复后的 2-D padding 掩码行为与 Hugging Facetransformers中SiglipTextTransformer.create_bidirectional_mask的实现保持一致。在 HF 中，create_bidirectional_mask生成的是(batch_size, 1, seq_len, seq_len)的双向掩码，且同样采用"True/1= 可 attend、False/0= 屏蔽"的约定。

对齐的意义在于：

1. 语义一致：同一份attention_mask在 Kornia 与 HF 两条实现路径下产生一致的注意力分布；
2. 行为一致：对部分 padding的序列，不再产生全掩码行——修复前"部分 padding 的序列会整行被屏蔽"的错误行为被消除；
3. 结果可移植：基于 HF 调通的权重与预处理流程，可以直接迁移到 Kornia 的纯 PyTorch 实现上。

仍然保留的边界行为：如果某一行掩码全部为 0（例如 batch 中的空 caption，整行都是 padding），该行依旧是 fully masked row，仍遵循上文 SDPA 的版本相关行为（2.5+ 返回0，2.4 及更早返回nan）。这是 SDPA 自身的语义，修复并未试图消除，使用时应通过批处理逻辑避免"整行为 0"的输入。

五、源码纵深：attention_mask 的完整调用链

attention_mask从 API 入口到注意力内核的传递路径（全部位于 kornia/models/siglip2 目录）：

1. 入口：model.py 中SigLip2Model.get_text_features(input_ids, attention_mask=...)，将掩码透传给SigLip2TextModel；
2. 文本模型：text_encoder.py 的SigLip2TextModel.forward把掩码传给SigLip2TextEncoder，后者再逐层传给每个SigLip2TextLayer（pre-norm + 残差结构，self_attn与mlp分别经过layer_norm1/layer_norm2）；
3. 注意力模块：attention.py 的SigLip2Attention.forward完成掩码维度归一后调用：

attention_output = F.scaled_dot_product_attention( query, key, value, attn_mask=attn_mask, dropout_p=dropout_p, # 训练时使用配置的 attention dropout，推理时为 0 )

注意力输出经transpose(1, 2).contiguous().view(...)还原为(B, N, hidden_size)后由out_proj输出。文本侧在最后一层后做final_layer_norm，并以**最后一个 token（EOS 位置）**做池化（text_encoder.py#L248-L252）——这解释了为什么get_text_features对掩码错误如此敏感：池化位置本身位于序列末端，极易落入 padding 区。

顺带一提，视觉侧SigLip2VisionModel也复用了同一个SigLip2Attention（见 vision_encoder.py），因此本次修复同时惠及get_image_features(..., attention_mask=...)路径。

六、测试验证：修复后的行为有据可查

tests/models/test_siglip2.py 中直接覆盖了本次修复的语义：

• test_attention：验证"全 1 掩码必须是 no-op"——attention(hidden_states, attention_mask=全1)与attention(hidden_states)的输出在容差内一致（半精度下放宽到rtol=atol=5e-2，因为传入attn_mask会禁用 CUDA flash 后端、两者可能走不同 SDPA kernel，仅存在 kernel 级舍入差异）；
• test_attention_mask_handling：构造两条长度不同的序列（5 token 与 8 token，其余 padding），断言输出torch.isfinite(features).all()且两条特征不相等——正是"部分 padding 不再产生全掩码行"的回归测试；
• test_get_text_features：全 1 掩码下特征为有限值且单位范数（归一化后norm == 1）。

这些测试同时验证了SigLip2Model的test_dynamo（torch.compile兼容性）与test_gradcheck，说明修复没有引入额外的编译或数值问题。

七、实战：正确使用 attention_mask 获取文本特征

结合 builder.py 与 config.py 提供的接口，一个可直接运行的推理示例（纯 PyTorch，无transformers依赖，见 changelog.d/+migration-011.added.md）：

import torch from kornia.models.siglip2 import SigLip2Builder # 加载预训练权重（内部经 kornia 自研下载器 + safetensors 读取， # 缓存路径与其它 kornia 检查点一致，可用 cache_dir 指定） model = SigLip2Builder.from_pretrained_hf("google/siglip2-base-patch16-224") model.eval() # 构造一个 batch：第一条 5 个 token，第二条 8 个 token，其余为 padding(0) batch_size, seq_len = 2, 10 input_ids = torch.zeros(batch_size, seq_len, dtype=torch.long) input_ids[0, :5] = torch.arange(1, 6) input_ids[1, :8] = torch.arange(11, 19) attention_mask = torch.ones(batch_size, seq_len) attention_mask[0, 5:] = 0 attention_mask[1, 8:] = 0 with torch.no_grad(): features = model.get_text_features(input_ids, attention_mask=attention_mask) print(features.shape) # (2, projection_dim) print(torch.isfinite(features).all()) # True：不再出现 NaN/全零

使用要点与注意事项：

1. 掩码取值约定：1= 有效 token（attend），0= padding（屏蔽）。直接传入即可，SigLip2Attention会负责转成 SDPA 布尔布局，调用方不要预先取反；
2. 支持维度：(B, N)、(B, N, N)、(B, 1, N, N)三种均可用；2-D 掩码只屏蔽 key 轴，因此 padding 的 query 行输出虽然有限但有意义缺失——按 attention.py 中 docstring 的说明，调用方应自行丢弃这些位置的输出；
3. 空序列边界：某行掩码全 0（空 caption）时，该行仍是 fully masked row，遵循 SDPA 的版本相关行为（2.5+ 为0，更早为nan），请在批处理层面规避；
4. 注意力输出本身：修复只影响attention_mask分支，SigLip2Attention的结构（独立的 Q/K/V 投影、scale = 1 / sqrt(head_dim)、head_dim缺省取hidden_size // num_heads且强制整除校验）与修复前一致，权重无需重新加载。

八、结语：一个极性错误，两种版本症状

本次修复（+migration-120.fixed.md）的技术要点可归纳为一句话：SDPA 的布尔掩码以True表示 attend，任何在交给它之前取反用户掩码的实现都是错的。修复后 Kornia 的SigLip2实现与 Hugging Facecreate_bidirectional_mask语义对齐，get_text_features/get_image_features在传入掩码时不再输出全零或 NaN 的特征。若你在旧版本 Kornia 上观察到文本特征异常，升级并配合tests/models/test_siglip2.py中的用例即可快速回归验证。

• 计算机视觉
• 深度学习
• 图像处理