Abstract
Vision Transformers (ViT) have emerged as the de-facto choice for numerous industry grade vision solutions. But their inference cost can be prohibitive for many settings, as they compute self-attention in each layer which suffers from quadratic computational complexity in the number of tokens. On the other hand, spatial information in images and spatio-temporal information in videos is usually sparse and redundant. In this work, we introduce LookupViT, that aims to exploit this information sparsity to reduce ViT inference cost. LookupViT provides a novel general purpose vision transformer block that operates by compressing information from higher resolution tokens to a fixed number of tokens. These few compressed tokens undergo meticulous processing, while the higher-resolution tokens are passed through computationally cheaper layers. Information sharing between these two token sets is enabled through a bidirectional cross-attention mechanism. The approach offers multiple advantages - (a) easy to implement on standard ML accelerators (GPUs/TPUs) via standard high-level operators, (b) applicable to standard ViT and its variants, thus generalizes to various tasks, (c) can handle different tokenization and attention approaches. LookupViT also offers flexibility for the compressed tokens, enabling performance-computation trade-offs in a single trained model. We show LookupViT's effectiveness on multiple domains - (a) for image-classification (ImageNet-1K and ImageNet-21K), (b) video classification (Kinetics400 and Something-Something V2), (c) image captioning (COCO-Captions) with a frozen encoder. LookupViT provides $2\times$ reduction in FLOPs while upholding or improving accuracy across these domains. In addition, LookupViT also demonstrates out-of-the-box robustness and generalization on image classification (ImageNet-C,R,A,O), improving by up to $4\%$ over ViT.
Abstract (translated)
Vision Transformers (ViT) 已经成为许多行业级视觉解决方案的首选,但对于许多设置,其推理成本可能过高,因为它们在每层计算自注意力,这会导致在token数量上的平方计算复杂度。另一方面,图像和视频中的空间信息通常是稀疏和冗余的。在这项工作中,我们介绍了一个名为LookupViT的工作,旨在利用这种信息稀疏性降低ViT的推理成本。LookupViT提供了一个新颖的一般性视觉 transformer 块,通过压缩高分辨率token的信息,将其压缩到固定数量的路上。这几个经过精心处理的压缩token在较低计算成本的层中通过双向注意力机制进行信息共享。这种方法具有多个优点-(a)通过标准的高层操作很容易实现于标准 ML 加速器(GPUs/TPUs)上,(b)适用于标准的ViT及其变体,所以可以应用于各种任务,(c)可以处理不同的token化和注意力方法。LookupViT还具有压缩token的灵活性,在单个训练模型上实现性能计算权衡。我们在多个领域证明了LookupViT的有效性 -(a)图像分类(ImageNet-1K 和 ImageNet-21K),(b)视频分类(Kinetics400 和 Something-Something V2),(c)图像摘要(COCO-Captions)使用冻结编码器。LookupViT在FLOPs方面提供了2倍的增长,同时在这些领域保持或提高准确性。此外,LookupViT还展示了在图像分类(ImageNet-C,R,A,O)上的稳健性和泛化能力,提高了4%以上。
URL
https://arxiv.org/abs/2407.12753