iFAN：面向推理感知的純遮罩Transformer學習

摘要

基於查詢的遮罩 Transformer 透過最終層查詢預測之間的像素級競爭來組裝分割輸出，然而此推理過程並未在訓練期間被明確最佳化。我們識別出兩個關鍵的不匹配：具有最高機率-遮罩分數的查詢不一定產生最準確的遮罩，且最終層解碼可能丟棄來自中間層的更優預測。為了解決這些問題，我們提出推理感知學習（iFAN），一個適用於普通遮罩 Transformer 的通用訓練框架。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.