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
iFAN fixes two inference-time mismatches in mask transformers, top-probability queries do not yield the best masks and final-layer decoding discards stronger intermediate predictions, via a quality-ranking and cross-layer self-distillation objective that is training-only.
Query-based mask transformers, the family used for segmentation, assemble their output at inference through pixel-wise competition among queries of the final layer. That inference process is never optimized explicitly during training, and the paper quantifies two mismatches. First, the query with the highest probability-mask score does not give the most accurate mask: 64.45% of results picked by max-prob-mask differ from the max-IoU oracle, and 29.27% are high-score but low-IoU. Second, final-layer decoding discards better predictions from intermediate layers: only 39.8% of ground-truth targets get their best prediction at the final layer.
iFAN (Inference-Aware Learning), from a ByteDance team, aligns training with inference through two training-only objectives while keeping efficient final-layer decoding unchanged at inference:
Both objectives are training-only; inference adds no parameters or latency.
On COCO, ADE20K, and Cityscapes, gains are consistent across panoptic, instance, and semantic segmentation, across the EoMT and PMT architectures, and across ViT-S/B/L/G backbones, averaging +1.20 PQ, +1.30 AP, and +0.63 mIoU with near-zero extra parameters, FLOPs, or latency. A few highlights:
| Setting | Gain |
| ADE20K panoptic (PMT, ViT-L, 1280 squared) | +3.0 PQ |
| ADE20K panoptic (EoMT, ViT-L, 1280 squared) | +2.2 PQ |
| COCO instance (EoMT, ViT-L, 1280 squared) | +1.8 AP |
| COCO instance (PMT, ViT-L, 640 squared) | +1.2 AP |
The ablation also confirms the fix: top-query conflict drops from 64.45% to 27.49%. The authors note gains shrink as the backbone scales (smaller models leave more competition errors for iFAN to correct), but even ViT-G improves.
A plug-and-play training framework for the plain-mask-transformer family (EoMT/PMT), with zero inference cost. The more valuable part is quantifying the train-inference mismatch into two concrete errors with numbers, a diagnosis useful to segmentation researchers on its own.
Absolute gains are modest and shrink as the backbone grows (ViT-G adds only +0.4 PQ on COCO panoptic). Verified only on panoptic, instance, and semantic segmentation and two architectures. Whether the gain justifies changing the training recipe on already-strong backbones depends on the use case.