iFAN: Inference-Aware Learning for Plain Mask Transformers
Fang Li, Yu He, Haoyang Tong, Lichen Ma, Jingling Fu, Wenxiao Fan, Tongxuan Liu, Luohang Liu, Ke Zhang, Junshi Huang
cs.CV
2026-08-04
指出 mask transformer 的 query 竞争在训练时未被优化:最高分 query 掩码未必最准、最后一层丢弃中间层更好预测。iFAN 用质量排序+跨层自蒸馏对齐训练与推理,推理零额外开销。
基于 query 的 mask transformer(做分割的那类)在推理时,靠最后一层各 query 之间的逐像素竞争拼出分割结果。问题是这个推理过程在训练里从没被显式优化过,论文量化出两个错配。第一,概率-掩码分数最高的 query 不见得掩码最准:按 max-prob-mask 选出来的结果里,有 64.45% 和「按最大 IoU 选」的理论上限对不上,其中 29.27% 是高分但 IoU 很低。第二,最后一层解码会丢掉中间层更好的预测:只有 39.8% 的真值目标其最佳预测出现在最后一层。
字节跳动团队的 iFAN(Inference-Aware Learning)用两个只在训练时起作用的目标对齐训练和推理,推理阶段保持高效的最后一层解码不变:
两个目标都只在训练时加,推理没有额外参数和延迟。
COCO、ADE20K、Cityscapes 上,跨全景/实例/语义分割、跨 EoMT 和 PMT 两种架构、跨 ViT-S/B/L/G 多种 backbone 都稳定涨,平均 +1.20 PQ、+1.30 AP、+0.63 mIoU,参数/FLOPs/延迟几乎不变。挑几个亮的:
| 设置 | 涨幅 |
| ADE20K 全景(PMT,ViT-L,1280²) | +3.0 PQ |
| ADE20K 全景(EoMT,ViT-L,1280²) | +2.2 PQ |
| COCO 实例(EoMT,ViT-L,1280²) | +1.8 AP |
| COCO 实例(PMT,ViT-L,640²) | +1.2 AP |
消融里还有个验证:top-query 冲突率从 64.45% 降到 27.49%。作者也观察到 backbone 越大涨幅越小(小模型留下的可纠错竞争错误更多),但到 ViT-G 仍有提升。
它是 plain mask transformer 这类(EoMT/PMT)架构上一个即插即用的训练框架,推理零成本。更值钱的是它把「训练和推理不对齐」这件事量化成了两个具体错配和一组数字,这个诊断对做分割的人独立有用。
绝对涨幅不算大,而且 backbone 越大收益越边际(ViT-G 在 COCO 全景只 +0.4 PQ)。只在全景/实例/语义三类分割、两种架构上验证过。在已经很强的骨干上,这点提升是否值得改训练流程,要看具体场景。