iFAN: Training Plain Mask Transformers for the Way They Actually Infer

Lucas

xxlucas

Follow

Paper: iFAN: Inference-Aware Learning for Plain Mask Transformers
Keywords: Image Segmentation, Mask Transformer, Inference-Aware Learning, Self-Distillation
Project Page: iFAN
Code: GitHub

Query-based Mask Transformers have become a widely adopted framework for semantic, instance, and panoptic segmentation. Each query jointly predicts a class probability and a segmentation mask, while the final output is assembled through pixel-level competition based on probability-mask scores.

However, the objectives used to train these models do not fully reflect how their predictions are selected at inference time. A model may already produce a high-quality mask, yet fail to expose it in the final output because the corresponding query is ranked incorrectly or because useful information from an intermediate decoder layer is not preserved by the last layer.

Two Overlooked Training–Inference Mismatches

Query-level mismatch: the highest-scoring query is not necessarily the best one

On COCO val2017, the query selected by the highest probability-mask score differs from the query with the highest mask IoU in 64.45% of cases. Moreover, 29.27% of predictions receive a high probability-mask score despite having a relatively low mask IoU.

This suggests that the standard inference score often favors the most confident query rather than the query with the best segmentation quality.

Layer-level mismatch: the last layer is not always the best layer

Mask Transformers are typically supervised at multiple decoder layers, but only the final layer is used at inference time. Our analysis shows that the last layer produces the best prediction for only about 39.8% of objects. In many cases, a higher-quality prediction already appears in an intermediate layer, but is weakened or forgotten in later decoding stages.

Directly ensembling predictions from multiple layers may improve accuracy, but it also introduces extra post-processing, computation, and latency. A more desirable solution is to preserve useful intermediate-layer information during training while keeping the original inference pipeline unchanged.

The Core Contribution of iFAN

To address these mismatches, we introduce iFAN, an inference-aware learning framework that redesigns the training process of Plain Mask Transformers around the decisions that actually determine the final output.

iFAN contains two complementary components:

• Adjusted Probability-Mask Ranking (APMR): makes inference scores better reflect mask quality and explicitly trains matched queries to outrank confusing hard-negative queries.
• Cross-Layer Self-Distillation (CLSD): dynamically selects intermediate-layer predictions that outperform the final layer and transfers their inference-relevant knowledge to the last layer.

Both the ranking objective in APMR and the cross-layer distillation in CLSD are used only during training. At inference time, the model still reads predictions from the final layer alone, introducing almost no additional parameters, computation, or latency.

APMR: Helping High-Quality Queries Win the Competition

The conventional probability-mask score is mainly determined by class confidence and pixel-level mask responses. As a result, a query with high classification confidence but poor mask quality may still dominate the competition at inference time.

APMR adds a lightweight Mask Quality Head to estimate the global mask quality of each query. Let ${\widehat{p}}_i\backslash hat\{p\}\_i$ denote the class probability of query $ii$, ${\widehat{M}}_i\backslash hat\{M\}\_i$ its mask response, and ${\widehat{q}}_i\backslash hat\{q\}\_i$ its predicted mask quality. For class $c_{k}c\_k$, the adjusted inference score is defined as

$${\widehat{S}}_i={({\widehat{q}}_i\cdot {\widehat{p}}_i(c_k))}^{\gamma }\cdot {\widehat{M}}_i,\backslash hat\; S\_i\; =\; \backslash left(\backslash hat\; q\_i\backslash cdot\; \backslash hat\; p\_i(c\_k)\backslash right)^\backslash gamma\; \backslash cdot\; \backslash hat\; M\_i,$$

where $\gamma \backslash gamma$ sharpens the distinction among high-scoring predictions and is set to $22$ in our experiments. Under this formulation, a query receives a high score only when its class prediction, overall mask quality, and pixel-level responses are all reliable.

Score calibration alone, however, is not sufficient. During training, unmatched queries are supervised as “no-object” predictions, but they may still retain strong foreground responses and become misleading competitors at inference time. APMR therefore introduces an explicit ranking objective between matched queries and hard-negative queries.

For object $kk$, let $ii$ be its matched query and $jj$ a hard-negative query. The ranking loss is

$${\mathcal{L}}_{\mathrm{r}\mathrm{a}\mathrm{n}\mathrm{k}}^{\mathrm{\prime }}=\frac{1}{\mathrm{\mid }{M}_k\mathrm{\mid }}\sum \mathrm{softplus}({\widehat{S}}_j-{\widehat{S}}_i)\cdot M_k\mathrm{.}\backslash mathcal\; L\_\{\backslash mathrm\{rank\}\}\text{'}\; =\; \backslash frac\{1\}\{|M\_k|\}\; \backslash sum\; \backslash operatorname\{softplus\}\; \backslash left(\backslash hat\; S\_j-\backslash hat\; S\_i\backslash right)\; \backslash cdot\; M\_k.$$

The loss compares the two queries only inside the ground-truth foreground region. It increases when the hard-negative score approaches or exceeds the matched-query score, and gradually decreases once the correct query establishes a clear margin.

The key contribution of APMR is therefore not merely the addition of a quality-prediction branch. It unifies mask-quality estimation, inference-time scoring, and relative query ranking into a single objective that directly optimizes the competition used to produce the final segmentation.

CLSD: Preserving Strong Intermediate Predictions in the Final Layer

While APMR resolves competition among queries within the same layer, CLSD addresses information loss across decoder layers.

During training, CLSD compares predictions for the same object across intermediate layers and the final layer. For object $kk$, the teacher layer can be summarized as

where $LL$ denotes the final layer. An intermediate prediction is selected as the teacher only when it genuinely outperforms the final-layer prediction.

CLSD then transfers the teacher query's adjusted probability-mask score map to the corresponding query in the final layer. The central relationship can be written as

$${\widehat{S}}_k^{({\mathrm{\ell }}_k^{\ast })}⟶{\widehat{S}}_k^{(L)}\mathrm{.}\backslash hat\; S\_\{k\}^\{(\backslash ell\_k^\{*\})\}\; \backslash longrightarrow\; \backslash hat\; S\_\{k\}^\{(L)\}.$$

The goal is not to force unconditional agreement across all layers. Instead, knowledge is transferred only when an intermediate layer already contains better evidence than the final layer.

CLSD requires neither a separately trained teacher network nor multi-layer fusion at inference time. Intermediate layers are used only during training to identify stronger predictions, while the final model retains efficient single-layer inference.

Where Does the Novelty Lie?

The novelty of iFAN does not come from building a more complicated decoder. Instead, it lies in redefining what a Mask Transformer should optimize during training.

1. Learning the inference-time selection process itself

Conventional objectives mainly supervise whether a matched query predicts the correct class and mask. Yet the final output also depends on whether that query wins the competition at inference time. iFAN therefore goes beyond producing strong candidates: it explicitly trains the model so that the correct candidate is selected under the actual inference rule.

2. Combining quality-aware scoring with hard-negative ranking

Predicting mask quality alone can improve score calibration, but it does not guarantee that the correct query will outrank a specific misleading competitor. APMR combines quality estimation with pairwise ranking, directly optimizing the relative order between matched and hard-negative queries.

3. Turning intermediate layers into dynamic self-distillation teachers

Auxiliary supervision can improve predictions at intermediate layers, but it does not ensure that their strongest results survive in the final layer. CLSD performs knowledge transfer only when an intermediate prediction is better, making self-distillation selective and target-aware rather than unconditional.

4. Improving training without changing inference

The ranking loss in APMR and the distillation mechanism in CLSD are confined to training. Testing requires no multi-layer ensemble, additional post-processing, or more complex decoder, so the accuracy gains introduce almost no extra inference cost.

Overall, iFAN focuses on two decisions that directly shape the final output: which query should win, and which layer's useful evidence should be preserved.

Experimental Results

We apply iFAN to two Plain Mask Transformer architectures: EoMT, built on DINOv2 backbones, and PMT, built on DINOv3 backbones. Experiments cover COCO, ADE20K, and Cityscapes across panoptic, instance, and semantic segmentation.

On average, iFAN achieves:

• +1.20 PQ for panoptic segmentation;
• +1.30 AP for instance segmentation;
• +0.63 mIoU for semantic segmentation.

The gains remain consistent across different architectures, backbone sizes, input resolutions, datasets, and segmentation tasks, while largely preserving the original parameter count, FLOPs, and inference speed.

Conclusion

iFAN shifts the focus from merely predicting a good mask to ensuring that the mask is actually selected at inference time and preserved in the final decoder layer.

Through APMR and CLSD, iFAN improves query competition and cross-layer information retention, respectively. The result is a simple, general, and effective training enhancement for Plain Mask Transformers that delivers consistent accuracy gains with almost no additional inference cost.