titleiFAN: Inference-Aware Learning for Plain Mask Transformers
authorsFang Li, Yu He, Haoyang Tong, Lichen Ma, Jingling Fu, Wenxiao Fan, Tongxuan Liu, Luohang Liu, Ke Zhang, Junshi Huang
date2026-08-07
venuearXiv (preprint, not peer reviewed)
idarXiv:2608.03216
upvotes5 (Hugging Face Daily Papers, 2026-11-01)
linksabs pdf hf code
abstractQuery-based mask transformers assemble segmentation outputs through pixel-wise competition among query predictions of the final layer, yet this inference process is not explicitly optimized during training. We identify two key mismatches: the query with the highest probability-mask score does not necessarily produce the most accurate mask, and final-layer decoding may discard superior predictions from intermediate layers. To address these issues, we propose Inference-Aware Learning (iFAN), a general training framework for plain mask transformers. iFAN introduces Adjusted Probability-Mask Ranking (APMR), which aligns query competition with predicted mask quality and suppresses high-confidence but inaccurate competitors. We further employ Cross-Layer Self-Distillation (CLSD) to transfer stronger intermediate predictions to the final layer. The ranking and distillation objectives are training-only, while inference retains efficient final-layer decoding. Experiments on COCO, ADE20K, and Cityscapes demonstrate consistent improvements across panoptic, instance, and semantic segmentation, as well as across different architectures, backbone scales, and input resolutions. Overall, iFAN improves performance by an average of 1.20 PQ, 1.30 AP, and 0.63 mIoU, with negligible additional parameters, FLOPs and inference latency.

cite

@misc{li2026ifan,
  title={iFAN: Inference-Aware Learning for Plain Mask Transformers},
  author={Fang Li and Yu He and Haoyang Tong and Lichen Ma and Jingling Fu and Wenxiao Fan and Tongxuan Liu and Luohang Liu and Ke Zhang and Junshi Huang},
  year={2026},
  eprint={2608.03216},
  archivePrefix={arXiv},
  url={https://arxiv.org/abs/2608.03216}
}