iFAN: Inference-Aware Learning for Plain Mask Transformers

2026-08-04

AI Summary

This work addresses a critical misalignment in existing query-based masked Transformers, where training objectives do not reflect inference goals, leading to high-confidence queries that may correspond to low-quality masks and the discarding of superior intermediate predictions. To resolve this, the authors propose iFAN, an inference-aware training framework that aligns query ranking with mask quality via Adjusted Probability-Mask Ranking (APMR) and transfers high-quality predictions from intermediate layers to the final layer through Cross-Layer Self-Distillation (CLSD). Notably, iFAN requires no modification to the inference pipeline and consistently improves performance across diverse benchmarks—yielding average gains of 1.20 PQ, 1.30 AP, and 0.63 mIoU on COCO, ADE20K, and Cityscapes—while being compatible with various architectures, backbone scales, and input resolutions, all with negligible increases in parameters, computational cost, or latency.

Abstract

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.

Problem

• mask transformers
• inference-aware learning
• query competition
• segmentation
• training-inference mismatch

Innovation

• Inference-Aware Learning
• Adjusted Probability-Mask Ranking
• Cross-Layer Self-Distillation
• Mask Transformer
• Segmentation