Search papers, labs, and topics across Lattice.
This paper introduces Inference-Aware Learning (iFAN), a training framework designed to optimize query-based mask transformers by addressing mismatches between query probability and mask accuracy. The authors implement Adjusted Probability-Mask Ranking (APMR) to align query competition with the quality of predicted masks, while Cross-Layer Self-Distillation (CLSD) enhances the final layer's predictions by leveraging stronger intermediate outputs. Experimental results on datasets like COCO, ADE20K, and Cityscapes show that iFAN consistently improves segmentation performance across various metrics without increasing computational costs.
iFAN boosts segmentation accuracy by aligning query competition with mask quality, achieving significant performance gains without extra computational overhead.
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.
