# iFAN: Inference-Aware Learning for Plain Mask Transformers

arXiv: https://arxiv.org/abs/2608.03216 · HTML: https://mlsanity.com/paper/2608.03216/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 25 of its 31 references that are in the map.
- First posted on arXiv 4 Aug 2026.

## Builds on (most cited first, 25 of 25)

- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Microsoft COCO: Common Objects in Context](https://mlsanity.com/paper/1405.0312/) (2014)
- [DINOv2: Learning Robust Visual Features without Supervision](https://mlsanity.com/paper/2304.07193/) (2023)
- [End-to-End Object Detection with Transformers](https://mlsanity.com/paper/2005.12872/) (2020)
- [The Cityscapes Dataset for Semantic Urban Scene Understanding](https://mlsanity.com/paper/1604.01685/) (2016)
- [Masked-attention Mask Transformer for Universal Image Segmentation](https://mlsanity.com/paper/2112.01527/) (2021)
- [DINO: DETR with Improved DeNoising Anchor Boxes for End-to-End Object Detection](https://mlsanity.com/paper/2203.03605/) (2022)
- [Per-Pixel Classification is Not All You Need for Semantic Segmentation](https://mlsanity.com/paper/2107.06278/) (2021)
- [Panoptic Segmentation](https://mlsanity.com/paper/1801.00868/) (2018)
- [Be Your Own Teacher: Improve the Performance of Convolutional Neural Networks via Self Distillation](https://mlsanity.com/paper/1905.08094/) (2019)
- [DN-DETR: Accelerate DETR Training by Introducing Query DeNoising](https://mlsanity.com/paper/2203.01305/) (2022)
- [Deeply-Supervised Nets](https://mlsanity.com/paper/1409.5185/) (2014)
- [MaX-DeepLab: End-to-End Panoptic Segmentation with Mask Transformers](https://mlsanity.com/paper/2012.00759/) (2020)
- [OneFormer: One Transformer to Rule Universal Image Segmentation](https://mlsanity.com/paper/2211.06220/) (2022)
- [Mask Scoring R-CNN](https://mlsanity.com/paper/1903.00241/) (2019)
- [Panoptic SegFormer: Delving Deeper into Panoptic Segmentation with Transformers](https://mlsanity.com/paper/2109.03814/) (2021)
- [Mask Transfiner for High-Quality Instance Segmentation](https://mlsanity.com/paper/2111.13673/) (2021)
- [Rank-DETR for High Quality Object Detection](https://mlsanity.com/paper/2310.08854/) (2023)
- [MP-Former: Mask-Piloted Transformer for Image Segmentation](https://mlsanity.com/paper/2303.07336/) (2023)
- [DETRDistill: A Universal Knowledge Distillation Framework for DETR-families](https://mlsanity.com/paper/2211.10156/) (2022)
- [Pseudo-label Alignment for Semi-supervised Instance Segmentation](https://mlsanity.com/paper/2308.05359/) (2023)
- [KD-DETR: Knowledge Distillation for Detection Transformer with Consistent Distillation Points Sampling](https://mlsanity.com/paper/2211.08071/) (2022)
- [OD-DETR: Online Distillation for Stabilizing Training of Detection Transformer](https://mlsanity.com/paper/2406.05791/) (2024)
- [The Missing Point in Vision Transformers for Universal Image Segmentation](https://mlsanity.com/paper/2505.19795/) (2025)
- [PMT: Plain Mask Transformer for Image and Video Segmentation with Frozen Vision Encoders](https://mlsanity.com/paper/2603.25398/) (2026)

## Abstract (arXiv)

Query-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.