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.

翻译：视觉Transformer（ViT）已成为众多工业级视觉解决方案的事实选择。但其推理成本在许多场景下可能过高，因为它们在每一层计算自注意力，其计算复杂度随令牌数量呈二次方增长。另一方面，图像中的空间信息与视频中的时空信息通常具有稀疏性和冗余性。本工作中，我们提出LookupViT，旨在利用这种信息稀疏性以降低ViT推理成本。LookupViT提供了一种新颖的通用视觉Transformer模块，其通过将高分辨率令牌信息压缩至固定数量的令牌来运作。这些少量压缩令牌经过精细处理，而高分辨率令牌则通过计算成本更低的层传递。两组令牌间的信息共享通过双向交叉注意力机制实现。该方法具有多重优势——（a）可通过标准高级算子轻松在标准ML加速器（GPU/TPU）上实现，（b）适用于标准ViT及其变体，因而可泛化至多种任务，（c）能处理不同的令牌化与注意力方法。LookupViT还为压缩令牌提供灵活性，支持在单一训练模型中实现性能与计算量的权衡。我们在多个领域验证了LookupViT的有效性——（a）图像分类（ImageNet-1K与ImageNet-21K），（b）视频分类（Kinetics400与Something-Something V2），（c）使用冻结编码器的图像描述生成（COCO-Captions）。LookupViT在这些领域实现FLOPs降低$2\times$的同时保持或提升了精度。此外，LookupViT在图像分类（ImageNet-C,R,A,O）上还展现出开箱即用的鲁棒性与泛化能力，较ViT最高提升$4\%$。