Progressive Token Length Scaling in Transformer Encoders for Efficient Universal Segmentation
This addresses efficiency issues for researchers and practitioners scaling segmentation models, though it is incremental as it builds on Mask2Former.
The paper tackles the high computational cost in transformer encoders for universal segmentation, specifically in Mask2Former, by proposing PRO-SCALE, a method that progressively scales token lengths to reduce encoder computations by ~52% and overall GFLOPs by ~27% with no performance drop on COCO.
A powerful architecture for universal segmentation relies on transformers that encode multi-scale image features and decode object queries into mask predictions. With efficiency being a high priority for scaling such models, we observed that the state-of-the-art method Mask2Former uses 50% of its compute only on the transformer encoder. This is due to the retention of a full-length token-level representation of all backbone feature scales at each encoder layer. With this observation, we propose a strategy termed PROgressive Token Length SCALing for Efficient transformer encoders (PRO-SCALE) that can be plugged-in to the Mask2Former segmentation architecture to significantly reduce the computational cost. The underlying principle of PRO-SCALE is: progressively scale the length of the tokens with the layers of the encoder. This allows PRO-SCALE to reduce computations by a large margin with minimal sacrifice in performance (~52% encoder and ~27% overall GFLOPs reduction with no drop in performance on COCO dataset). Experiments conducted on public benchmarks demonstrates PRO-SCALE's flexibility in architectural configurations, and exhibits potential for extension beyond the settings of segmentation tasks to encompass object detection. Code here: https://github.com/abhishekaich27/proscale-pytorch