Dynamic Token-Pass Transformers for Semantic Segmentation
Yuang Liu Qiang Zhou Jing Wang Fan Wang Jun Wang Wei Zhang
Vision transformers (ViT) usually extract features via forwarding all the
tokens in the self-attention layers from top to toe. In this paper, we
introduce dynamic token-pass vision transformers (DoViT) for semantic
segmentation, which can adaptively reduce the inference cost for images with
different complexity. DoViT gradually stops partial easy tokens from
self-attention calculation and keeps the hard tokens forwarding until meeting
the stopping criteria. We employ lightweight auxiliary heads to make the
token-pass decision and divide the tokens into keeping/stopping parts. With a
token separate calculation, the self-attention layers are speeded up with
sparse tokens and still work friendly with hardware. A token reconstruction
module is built to collect and reset the grouped tokens to their original
position in the sequence, which is necessary to predict correct semantic masks.
We conduct extensive experiments on two common semantic segmentation tasks, and
demonstrate that our method greatly reduces about 40% $\sim$ 60% FLOPs and the
drop of mIoU is within 0.8% for various segmentation transformers. The
throughput and inference speed of ViT-L/B are increased to more than 2$\times$
on Cityscapes.