FlexAttention for Efficient High-Resolution Vision-Language Models (2407.20228v1)

Abstract: Current high-resolution vision-LLMs encode images as high-resolution image tokens and exhaustively take all these tokens to compute attention, which significantly increases the computational cost. To address this problem, we propose FlexAttention, a flexible attention mechanism for efficient high-resolution vision-LLMs. Specifically, a high-resolution image is encoded both as high-resolution tokens and low-resolution tokens, where only the low-resolution tokens and a few selected high-resolution tokens are utilized to calculate the attention map, which greatly shrinks the computational cost. The high-resolution tokens are selected via a high-resolution selection module which could retrieve tokens of relevant regions based on an input attention map. The selected high-resolution tokens are then concatenated to the low-resolution tokens and text tokens, and input to a hierarchical self-attention layer which produces an attention map that could be used for the next-step high-resolution token selection. The hierarchical self-attention process and high-resolution token selection process are performed iteratively for each attention layer. Experiments on multimodal benchmarks prove that our FlexAttention outperforms existing high-resolution VLMs (e.g., relatively ~9% in V* Bench, ~7% in TextVQA), while also significantly reducing the computational cost by nearly 40%.

• The paper presents FlexAttention, a novel mechanism that reduces computational cost for high-resolution vision-language models by selectively processing tokens.
• It employs a high-resolution selection module and hierarchical self-attention layer, achieving a 9% improvement on V* Bench and 7% on TextVQA.
• These advances enable detailed image analysis for applications like remote sensing and medical imaging, paving the way for adaptive attention research.

FlexAttention for Efficient High-Resolution Vision-LLMs

Vision-LLMs (VLMs) are integral to a wide array of tasks involving both image and text processing. They excel at tasks such as visual question answering and image-text matching. However, these models have operated predominantly at low resolutions, imposing constraints on their ability to scrutinize fine details within images. This limitation is particularly evident in instances where recognition of minor text or small objects is crucial.

The paper introduces FlexAttention, a novel attention mechanism devised to enhance the efficiency of high-resolution vision-LLMs. The primary objective of FlexAttention is to reduce computational costs while maintaining, or even improving, the model's performance by handling high-resolution inputs more effectively.

Methodological Framework

FlexAttention's approach departs from traditional exhaustive token-based attention mechanisms by integrating a hierarchical process. It encodes images into both low- and high-resolution tokens but utilizes only a fraction of the high-resolution tokens during attention computation. This is achieved through:

1. High-Resolution Selection Module: This module identifies relevant high-resolution tokens by analyzing an input attention map, thereby preserving computational resources by focusing only on areas of interest.
2. Hierarchical Self-Attention Layer: Following token selection, this layer concatenates the selected high-resolution tokens with low-resolution and text tokens. It then performs self-attention, iteratively refining the attention map for subsequent layers.

Through these innovations, FlexAttention significantly reduces computational demand — by nearly 40% — while outperforming existing high-resolution VLMs, achieving improvements of around 9% on the V* Bench benchmark and about 7% on TextVQA.

Implications and Future Work

The implications of FlexAttention in the high-resolution context are substantial. Practically, it allows for more efficient processing of images with finer details, which has direct applications in fields requiring detailed visual analysis, such as remote sensing or medical imaging.

Theoretically, FlexAttention challenges the conventional approach of exhaustive computation in self-attention mechanisms. By demonstrating that selective token processing does not detract from — and can even enhance — model performance, it sets the stage for further research into adaptive attention mechanisms.

Looking ahead, the principles underlying FlexAttention could inform the development of attention mechanisms in other domains. For instance, video or audio data, which also inherently involve long sequences, may benefit from a similar approach to improve computational efficiency without sacrificing detail-oriented accuracy. Continued exploration in this direction could lead to more scalable and resource-efficient models across various modalities.

In conclusion, while existing VLMs are limited by computational inefficiencies at high resolutions, FlexAttention offers a promising alternative path, underscoring the potential for thoughtful architectural innovation to drive the next frontier in AI research.