PaperlayerView exact revision on arXiv

Source-linked AI summary

iFAN: Inference-Aware Learning for Plain Mask Transformers

Fang Li, Yu He, Haoyang Tong, Lichen Ma, Jingling Fu, Wenxiao Fan, Tongxuan Liu, Luohang Liu, Ke Zhang, Junshi Huang

arXiv:2608.03216v2cs.CV

TL;DR

Mask transformers do not explicitly optimize final-layer query competition, so high-scoring queries can yield inaccurate masks and stronger intermediate predictions may be discarded. iFAN aligns competition with mask quality and transfers intermediate knowledge during training, improving panoptic, instance, and semantic segmentation by average gains of 1.20 PQ, 1.30 AP, and 0.63 mIoU with virtually no inference overhead.

  • Problem

    Mask-transformer maxprob-mask inference can select high-scoring queries with inaccurate masks, while final-layer decoding may discard stronger intermediate predictions.

  • Method

    iFAN combines quality-aware Adjusted Probability-Mask Ranking with Cross-Layer Self-Distillation to improve query competition and final-layer predictions without multi-layer inference.

  • Results

    Across panoptic, instance, and semantic segmentation, iFAN achieves average gains of 1.20 PQ, 1.30 AP, and 0.63 mIoU with virtually no additional inference overhead.

  • Takeaways & Limitations

    iFAN consistently improves plain mask-transformer segmentation across datasets, architectures, backbone sizes, and input resolutions while preserving inference efficiency.

Abstract

from arXiv · show

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.

Introduction

The introduction identifies inference mismatches in query-based mask transformers: probability-mask selection can choose inaccurate queries, while final-layer decoding can discard stronger intermediate predictions. It presents inference-aware learning, combining quality-aware ranking with cross-layer self-distillation to improve final predictions without added inference cost.

  • Inference limitations: 64.45% of max-prob-mask selections are suboptimal relative to the max-IoU query, revealing a mismatch between query confidence and mask accuracy.The selected query may cover only salient object regions or fail to predict an object’s complete shape.
  • Inference limitations: The max-prob-mask strategy combines class distributions with class-agnostic foreground masks, but isolated mask predictions lack global object information.This incomplete competition can produce correct categories with incomplete or inaccurate regions.
  • Inference limitations: 39.8% of ground-truth targets obtain their best predictions at the final layer, showing that intermediate layers can contain stronger segmentation results.The penultimate layer achieves comparable overall performance, motivating cross-layer knowledge transfer.
  • Inference-aware learning: iFAN adjusts probability-mask scoring with a global mask-quality score and uses object-oriented ranking to produce more discriminative prediction scores.These training objectives address inaccurate competitors during query competition.
  • Inference-aware learning: iFAN uses self-distillation across layers to improve final-layer predictions while preserving efficient inference without additional inference cost.The framework exploits useful information from different mask-transformer layers during training.

Related Work

Prior work extends query-based set prediction from detection to segmentation and develops quality-aware ranking, intermediate supervision, and self-distillation. iFAN distinguishes itself by applying ranking to probability-mask competition and transferring stronger intermediate predictions without requiring a separate teacher.

  • Query-Based Mask Transformers: DETR introduced object queries and bipartite set prediction, which MaskFormer and Mask2Former extended to class–mask segmentation.OneFormer and Mask DINO further developed the formulation for universal segmentation and joint detection–segmentation modeling.
  • Quality Estimation and Ranking: Quality-aware methods calibrate instance scores, resolve mask overlaps, or rank detection queries, whereas APMR targets probability-mask competition.The cited methods include Mask Scoring R-CNN, Panoptic SegFormer, and Rank-DETR.
  • Intermediate Supervision: Intermediate supervision uses annotation-based objectives or mask-consistency losses to improve intermediate predictions, but does not explicitly transfer a better intermediate probability-mask prediction.Examples include Deeply-Supervised Nets, query-based segmenters, and MP-Former.
  • Self-Distillation: Existing self-distillation transfers knowledge between predictions, while CLSD requires neither a pre-trained nor an exponential-moving-average teacher.For each target, CLSD selects the intermediate matched query with the highest soft-IoU.

Method

iFAN aligns query competition with mask quality through APMR and transfers stronger intermediate predictions to the final layer through CLSD. Both objectives are training-only, preserving efficient final-layer inference with negligible overhead.

  • Adjusted Probability-Mask Ranking: APMR adds a shared quality head to estimate each query’s global mask quality and calibrate probability-mask scores for segmentation competition.The head is a single linear layer with sigmoid, shared across queries and prediction layers, introducing negligible parameter overhead.
  • Adjusted Probability-Mask Ranking: APMR raises matched-query scores above hard unmatched competitors by using adjusted scores (q̂_i ∗ p̂_i)^γ, with γ = 2, and an object-oriented ranking loss.Top N_H unmatched queries are selected as hard negatives, while detached gradients prevent oversuppression of class probabilities and mask predictions.
  • Overall Objective and Inference: The combined objective augments segmentation training with ranking and distillation losses, while inference retains only final-layer decoding and incurs no additional CLSD cost.The coefficients are λ_rank = 0.1 and λ_distill = 0.4; the distillation coefficient uses warm-up and annealing.
  • Cross-Layer Self-Distillation: CLSD selects the matched intermediate-layer query with the highest soft IoU as a teacher and distills stronger adjusted probability-mask evidence into the final-layer query.Matching is performed across layers ℓ ∈ [L−4, L], and the hinge loss activates when the selected teacher is stronger than the final-layer prediction.

Experiments

Across panoptic, instance, and semantic segmentation benchmarks, iFAN consistently improves PMT and EoMT with virtually no additional inference overhead. Ablations and diagnostics attribute these gains to better query ranking, cross-layer distillation, and robust performance across model scales.

  • Panoptic segmentation: 0.55 PQ and 2.15 PQ improve PMT on COCO and ADE20K, while 0.68 PQ and 1.58 PQ improve EoMT, respectively.iFAN reaches 141 FPS for PMT at 640^2 resolution; the strongest EoMT-iFAN is only 0.3 PQ below the best reported results.
  • Instance segmentation: 1.1 AP improves PMT and 1.5 AP improves EoMT on COCO, with the strongest EoMT result rising from 48.8 to 50.6 AP.The improvements occur with virtually no additional inference overhead.
  • Semantic segmentation: 0.3 mIoU improves EoMT on Cityscapes, while EoMT and PMT improve by 0.6 and 0.9 mIoU on ADE20K, respectively.PMT also improves by 0.7 mIoU on Cityscapes, demonstrating gains across tasks, architectures, and datasets.
  • Ablation: 57.00 PQ is achieved by complete iFAN, versus 56.00 PQ for EoMT; quality learning, ranking, and CLSD contribute progressively in the ablation.The tested intermediate configurations obtain 56.43, 56.68, and 56.54 PQ.
  • Cross-layer and scale analysis: 53 to 98 classes become final-layer-optimal with CLSD, while intermediate-layer-optimal classes decrease from 80 to 35.Across ViT-S, B, L, and G backbones, iFAN improves PQ by 1.6, 1.2, 1.0, and 0.8, respectively.
  • Inference diagnostics: 64.45% to 27.49% reduces disagreement between max-prob-mask and max-IoU queries, while 29.27% to 11.50% reduces high-score, low-IoU cases.APMR calibrates the quality used for segmentation competition without disrupting the original classification and mask objectives.

Conclusion

iFAN is an inference-aware learning framework for plain mask transformers that improves query competition and cross-layer knowledge transfer. It makes query selection more consistent with mask quality and distills stronger intermediate predictions to the final layer.

  • Conclusion: iFAN improves both query competition and cross-layer knowledge transfer in plain mask transformers.The framework combines Adjusted Probability-Mask Ranking (APMR) with Cross-Layer Self-Distillation (CLSD).
  • Conclusion: APMR augments probability-mask scores with global mask quality to align inference-time query selection with mask quality.It explicitly ranks matched queries above hard unmatched competitors.
  • Conclusion: CLSD identifies a better intermediate prediction for each target and distills its adjusted probability-mask score map to the final layer.This transfers cross-layer knowledge to the final prediction stage.

Supplementary Material

The supplementary material reports hyperparameter sensitivity, details a query-wise diagnostic for APMR, and shows that iFAN substantially improves score–quality alignment. It also provides expanded benchmark tables and qualitative predictions across segmentation tasks.

  • Hyperparameter sensitivity: iFAN remains strong for NH ∈[7, 15] and changes by only 0.18 PQ as γ varies from 0.5 to 3; the default is NH = 10 and γ = 2.γ = 2 suppresses low-valued responses without the excessive attenuation caused by larger exponents.
  • Query-wise diagnostic: APMR raises pairwise ordering accuracy from 76.07% to 79.42% and top-query accuracy from 35.55% to 72.51%.The diagnostic evaluates EoMT and EoMT-iFAN on all 5,000 COCO val2017 images.
  • Query-wise diagnostic: The diagnostic compares each retained query’s competition score with its best class-consistent soft-IoU reference quality, using many-to-one query–segment comparisons.The evaluation first filters queries by class-agnostic spatial overlap, then prevents different-class overlaps from increasing reference quality.
  • Query-wise diagnostic: APMR decreases top-query conflict from 64.45% to 27.49%, average mask-quality gap from 0.1070 to 0.0591, and high-score low-quality rate from 29.27% to approximately 11.5%.These results indicate improved alignment in the high-score region governing probability-mask competition.
Loading 2608.03216v2…