iFAN: 推論を考慮したプレーンなマスクTransformerの学習

要旨

クエリベースのマストランスフォーマーは、最終層のクエリ予測間のピクセル単位の競争を通じてセグメンテーション出力を構築するが、この推論プロセスは学習中に明示的に最適化されない。我々は、2つの重要なミスマッチを特定する：最高の確率マスクスコアを持つクエリが必ずしも最も正確なマスクを生成するとは限らないこと、そして最終層のデコードが中間層からの優れた予測を破棄し得ることである。これらの問題に対処するため、我々はプレーンなマストランスフォーマーのための汎用的な学習フレームワークである推論を考慮した学習（iFAN）を提案する。iFANは、クエリ競争を予測マスク品質と整合させ、高信頼度だが不正確な競合予測を抑制する調整確率マスクランキング（APMR）を導入する。さらに、クロスレイヤー自己蒸留（CLSD）を用いて、より強力な中間層予測を最終層に転送する。ランキングと蒸留の目的関数は学習時のみに適用され、推論時には効率的な最終層デコードが保持される。COCO、ADE20K、Cityscapesでの実験により、パノプティック、インスタンス、セマンティックセグメンテーション、および異なるアーキテクチャ、バックボーンスケール、入力解像度にわたって一貫した改善が実証される。全体として、iFANは平均1.20 PQ、1.30 AP、0.63 mIoUの性能向上を達成し、追加パラメータ、FLOPs、推論レイテンシは無視できるほど僅かである。

English

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.