iFAN: Inference-Aware Learning for Plain Mask Transformers

The Problem: Training and Inference Don't Align

Query-based mask transformers (the standard architecture for modern segmentation) make predictions through a competition process at inference time: multiple query predictions are scored, and the one with the highest probability-mask score "wins" and gets used. The catch is that during training, models are never explicitly taught to optimize for this inference procedure. This creates two concrete failures: a query can have a high confidence score paired with a low-quality mask (winning the competition despite being wrong), and excellent predictions from earlier layers get thrown away because the model only looks at the final layer when assembling outputs.

The authors demonstrate these mismatches empirically, then propose iFAN (Inference-Aware Learning), a training framework that teaches models how to actually compete during inference without changing the efficient inference process itself.

How iFAN Works

The solution has two complementary parts, both operating only during training:

Adjusted Probability-Mask Ranking (APMR) directly targets the query competition problem. Instead of letting the raw probability-mask score determine winners, APMR learns to rank queries by actual mask quality. It penalizes high-confidence predictions that are actually wrong—suppressing misleading competitors—while promoting queries with good masks even if their raw scores are lower. This retrains the model's scoring mechanism to match real prediction accuracy.

Cross-Layer Self-Distillation (CLSD) recovers information from intermediate layers. The method treats strong predictions from earlier transformer layers as a teaching signal, forcing the final layer to absorb and reproduce that knowledge. This is "self-distillation" because both teacher and student come from the same model. During inference, only the final layer is used, so this adds zero latency or parameters—it's purely a training regularizer.

Both components are training-only objectives. Inference remains unchanged: it still runs the fast, single-pass final-layer decoding that makes transformers practical.

Why It Matters

The gains are consistent and meaningful across diverse settings. On COCO panoptic segmentation, iFAN improves by 1.20 PQ (panoptic quality). Instance segmentation improves 1.30 AP (average precision), and semantic segmentation by 0.63 mIoU (mean intersection over union). These improvements hold across different backbone sizes and input resolutions, and across three major benchmarks (COCO, ADE20K, Cityscapes).

Crucially, these gains come with no practical cost: negligible additional parameters, no increase in FLOPs (floating-point operations), and no inference latency overhead. The framework is also general—it works across different plain mask transformer architectures, not tied to a specific design.

The core insight is simple but powerful: if inference makes decisions based on a particular mechanism (query competition), then training should explicitly optimize for that mechanism. By closing the gap between what training optimizes and what inference actually does, iFAN extracts performance without architectural changes or runtime penalties. For practitioners, this means dropping iFAN into existing mask transformer training pipelines for immediate gains.