FlexAttention，通过高效处理高分辨率图像输入增强大型视觉-语言模型的能力 ！

FlexAttention，通过高效处理高分辨率图像输入增强大型视觉-语言模型的能力 ！

原创 AI 引擎  2024年08月08日 09:01 上海

当前的高分辨率视觉-语言模型将图像编码为高分辨率图像标记，并穷举所有这些标记以计算注意力，这显著增加了计算成本。为了解决这个问题，作者提出了FlexAttention，这是一种用于高效高分辨率视觉-语言模型的灵活注意力机制。

具体来说，一个高分辨率图像同时被编码为高分辨率标记和低分辨率标记，其中仅使用低分辨率标记和少数选定的高分辨率标记来计算注意力图，这大大减少了计算成本。

高分辨率标记是通过一个高分辨率选择模块来选取的，该模块可以根据输入的注意力图检索相关区域的标记。

选定的高分辨率标记随后与低分辨率标记和文本标记连接，并输入到分层自注意力层中，该层生成可用于下一步高分辨率标记选取的注意力图。分层自注意力和高分辨率标记选取过程在每个注意力层上迭代进行。

在多模态基准测试上的实验证明，作者的FlexAttention优于现有高分辨率视觉-语言模型（例如，在V* Bench上相对提高了约9%，在TextVQA上相对提高了约7%），同时计算成本也显著降低了近40%。

1 Introduction

大型的视觉-语言模型（VLMs），如文献[28, 30]中描述的，在包括图像标注、视觉问答、图像文本匹配等多模态任务中表现出卓越的能力。然而，这些模型通常以相对较低的分辨率处理图像（例如， 或 ），因而在需要仔细审视小区域（例如，小文本或小物体）的场景中表现不佳。这一局限在图1（a）中尤为明显，由于低分辨率输入的限制，这些模型无法辨识印刷标志上的文字。

为了解决这个问题，已经提出了几种高分辨率VLMs（例如，LLaVA-1.5-HD [35] 和 CogAgent [20]），它们能够接受高分辨率图像作为输入并将其编码为高分辨率标记。尽管这些模型对小区域进行了更详细的检查，但它们会穷尽地处理所有高分辨率标记以计算注意力，这给计算资源带来了沉重的负担。这些模型偏离了人类进行视觉推理的方式。作者不是记住所有像素级完美的细节，而是倾向于首先保持一个粗糙的表示，并在受到外部刺激时仅关注相关区域以获取更多细节。高分辨率VLMs能够根据低分辨率特征灵活、动态地关注感兴趣的区域，以便检索高分辨率细节，这一点至关重要。

为此，作者提出了FlexAttention，一种新颖的注意力机制，它能够无缝地嵌入大多数视觉-语言模型中，以高效地提高它们对高分辨率图像的感知能力。具体来说，如图1（c）所示，FlexAttention接受高分辨率图像作为输入，并将图像同时编码为高分辨率图像标记和低分辨率图像标记。为了计算效率，作者只将低分辨率标记和文本标记输入前几层以大致理解整个图像。对于后续层，只利用低分辨率标记和非常小部分的高分辨率标记来计算注意力，这显著减少了计算成本。在带有FlexAttention的每个解码层中，作者有一个高分辨率特征选择模块和一个分层自注意力模块。高分辨率特征选择模块根据输入的注意力图检索相关区域的高分辨率标记。所选择的高分辨率标记与低分辨率标记和文本标记一起连接，并输入到分层自注意力模块中。分层自注意力模块生成一个注意力图，该图可用于高分辨率标记选择，选择输入到下一层分层自注意力的标记。这两个模块迭代处理，直到最后一层，通过投影仪产生最终答案。

作者在几个高分辨率多模态基准上评估了作者的FlexAttention，包括通用的V* Bench [53] 和 Magnifierbench [27]，以及特定领域的TextVQA [45] 文本理解和RSVQA [38] 遥感。作者展示了与其他高分辨率方法相比，几乎减少40%的计算成本的同时性能更优，证明了作者方法的效率。更重要的是，与商业聊天机器人如GPT-4V [1]相比，作者在V* Bench上获得了更高的分数。

2 Related Works

视觉语言模型。作者的工作与尝试训练大型多模态模型的研究密切相关，这些模型用于各种视觉语言任务，如视觉问答、参照表达式理解以及基于文本的图像检索[56, 32]。传统方法[19, 29]通常收集大规模的视觉-语言数据集，并从零开始学习视觉与语言之间的联合表示以处理不同任务。这些模型通常在领域内数据上表现良好，但在需要常识理解和外部世界知识[40, 27]的基准测试中表现不佳。

后来，大型语言模型（LLM）在自然语言理解和推理方面显示出惊人的能力，这为视觉和语言研究带来了新的可能性。一系列大型视觉-语言模型被提出，这些模型通常将预训练的视觉编码器[13, 44]与预训练的大型语言模型[50, 9]相连接。Flamingo[2]首次使用交叉注意力机制将视觉上下文编码到LLM中。BLIP2[31]提出了QFormer，它使用BERT模型[24]将视觉特征转换为一组学习到的标记供LLM使用。Fuyu[4]直接将图像块投影成LLM的输入，以摆脱预训练的视觉编码器。

尽管这些模型在常识理解方面表现出色，并且在传统视觉语言任务上表现惊人，但它们通常无法处理需要高分辨率输入[53, 27]的任务，原因有两个：

1）大多数VLM使用CLIP[44]作为其视觉编码器，这限制了这些VLM的输入图像大小为CLIP训练的固定且相对较小的分辨率（例如，224x224）；

2）它们缺乏有效处理长图像块序列的模型机制，当图像分辨率增加时，图像块数量将成倍增加，导致计算成本过高。

高分辨率VLM。为了提高VLM处理高分辨率输入的能力，已经提出了几种VLM[14, 20, 36]。DocPedia[14]将图像转换到频域以保持高分辨率图像的更好语义。LLaVA 1.6[36]设计了不同尺度的输入以满足不同任务的需求，并在效率和性能之间取得平衡。尽管这些模型缓解了密集计算的问题，但它们与作者的方法是正交的，并未设计任何新的注意力机制来应对由自注意力机制引入的二次计算成本增加的挑战。最近，CogAgent[20]为高分辨率图像输入设计了一个新的视觉编码器。与作者的方法不同，它需要在大型语言模型的每一层通过交叉注意力计算隐藏状态与整个高分辨率图像特征之间的密集对应关系，这使得它的效率较低。此外，训练模型的数据并不公开，而作者已经计划为整个研究社区发布所有数据、代码和模型。

序列建模的高效机制。作者的工作与序列建模高效机制的发展相关。一种方法处理标准注意力机制关于序列长度的二次复杂度问题，这通过使用结构化近似或线性注意力[7]来实现。另一种方法完全用循环风格的模型替换注意力，如递归神经网络（RNN）和状态空间模型。特别相关的是杨等人[55]的工作，他们引入了一种用于文档分类的分层注意力网络。他们的模型使用分层结构和两层注意力机制来改进文档表示。作者的工作通过专注于为高分辨率图像输入设计的高效机制而与之区分开来，确保与大型语言模型的计算无缝协作。

3 Preliminary

符号表示. 作者定义了一些将在全文中使用的一些术语。对于高分辨率视觉-语言模型，作者将其高分辨率图像输入定义为 ，文本输入定义为 。此外，作者将低分辨率图像标记定义为 ，高分辨率图像标记定义为 ，文本标记定义为 。视觉-语言模型的隐藏状态表示为 ，序列长度为 ，隐藏状态大小为 。隐藏状态  包括  个低分辨率图像标记，随后是  个文本标记。作者定义  为  个高分辨率图像标记  的选定子集。

自回归大型语言模型. 自回归大型语言模型（LLM），如 LLaMA [50]，在大多数视觉-语言模型中扮演着关键角色，因为它们负责将图像和文本标记作为输入，并生成答案序列。自回归LLM由多个堆叠的解码器层组成。每个解码器层有两个子层：一个是自注意力模块，另一个是前馈神经网络（FFN）层。在每个子层周围都使用了跳过连接，并跟随层归一化（LN）。简而言之，每个子层的输出为 LN(子层)。为了简单起见，在后续讨论中将省略层归一化。

自注意力和注意力图. 自注意力 [51] 是解码器层的基本模块。对于自注意力，给定输入隐藏状态 ，它首先使用线性投影层将  投影到 、 和 ，即 Query 、键和值矩阵，并执行以下计算：

其中 ，，，且  是可学习的线性投影矩阵。特别地，注意力图  在 softmax 操作之后获得：

注意力图  是一个 x 矩阵，它衡量标记间的重要性：注意力图中的 (i, j) 注意力值表示第 j 个标记对第 i 个标记的重要性，值越高意味着第 j 个标记对第 i 个标记越重要。

自注意力的局限性. 自注意力机制的计算成本与隐藏状态  的序列长度  成二次增长关系。当整合高分辨率图像时，这种计算复杂性进一步放大，因为它显著增加了图像标记的数量，从而延长了隐藏状态的长度。因此，自注意力机制的计算要求显著增加，使得处理高分辨率图像输入变得不切实际，因为计算开销过大。

4 Vision-language Model with FlexAttention

Overall Architecture

为了解决自注意力在处理高分辨率图像时的局限性，作者引入了FlexAttention，它通过动态关注高分辨率图像的重要区域，高效地分析高分辨率图像。FlexAttention可以通过替换其自注意力模块为作者的提出的FlexAttention模块，而被集成到大多数视觉-语言模型中。

如图2所示，修改后的视觉-语言模型由个解码器层组成，其中前层使用标准的自注意力模块，最后层使用作者提出的FlexAttention模块。给定一个高分辨率图像，作者首先将其下采样为一个低分辨率图像，并将这两个图像输入图像编码器，分别得到高分辨率和低分辨率的图像标记。为了计算效率，作者仅将低分辨率图像标记和文本标记输入前层，以大致理解整个图像。对于接下来的个带有FlexAttention的解码器层，为了更高效地感知更多图像细节，作者还向其输入所选的高分辨率图像标记。具体来说，FlexAttention包括两个模块：一个高分辨率特征选择模块和一个分层自注意力模块。高分辨率特征选择模块不是将所有高分辨率标记进行前向传播，而是根据注意力图灵活地为下一层选择重要的标记。分层自注意力模块被设计用来将所选的高分辨率信息融合到原始隐藏状态中。最后，作者使用一个投影线性层来生成文本输出。

High-resolution Feature Selection Module

对于自回归的大语言模型（LLM），下一个标记是通过前一个标记的最后一个隐藏状态预测的。通过检查等式2中与最后一个标记对应的注意力图上所有其他标记的注意力值，作者可以找出模型在生成下一个预测标记时关注哪些标记。对于视觉语言模型，这也适用于图像标记。那些具有高注意力值的图像标记在生成下一个标记时可以被视为与重要图像区域相关。尽管低分辨率图像标记中包含的细节有限，但作者能够检索到模型所关注的高分辨率图像区域的详细信息。因此，作者不是将所有高分辨率标记输入到注意力模块中，这会导致计算成本过高，而是动态选择非常小的一部分（大约10%）的高分辨率标记，即，并且只将这一部分传递到注意力模块。

如图3所示，作者从注意力图的最后一列中取出前个值，这对应于低分辨率图像标记对最后一个文本标记的重要性，并将这个1-D向量 Reshape 为一个2-D图，作者将其称为注意力 Mask 。该 Mask 中的每个值与低分辨率图像中的一个块相关联，指示该块的重要性。这个 Mask 被归一化、二值化，然后调整到与高分辨率特征块标记相同的大小，形成高分辨率选择 Mask ，该 Mask 作为是否选择块标记的决策。最后，作者将这个 Mask 应用于高分辨率图像标记，以获得选定的高分辨率特征。

Hierarchical Self-attention Module

分层自注意力是将选定的较高分辨率标记的信息融合到隐藏状态中的核心机制，该隐藏状态由较低分辨率标记和文本标记组成。它以选定的较高分辨率标记和隐藏状态作为输入，输出注意力图和更新后的隐藏状态。分层自注意力的计算可总结为

(4)  (5)  其中是可学习的线性投影矩阵。和是融合了高分辨率特征信息的键和值矩阵。与自注意力类似，经过softmax操作后，作者可以得到一个注意力图：

与自注意力不同，这个注意力图的形状为，因为它还包含了高分辨率标记对应其他标记的注意力值。作者仅保留等式7中矩阵的前个注意力值作为选择下一层中使用的高分辨率特征的注意力图。关于带有作者FlexAttention的视觉-语言模型如何工作的伪算法在算法1中描述。

Complexity Analysis

FlexAttention提供了类似于传统自注意力计算的优点，因此在尽量减少对模型架构修改的同时，促进了多粒度特征的高效融合。设选择的高分辨率特征长度为，原始隐藏状态长度为，隐藏状态大小为。作者提出的分层自注意力的计算复杂度为

如果不使用作者的分层自注意力，而是直接将高分辨率图像与低分辨率图像相加，计算复杂度将为

对于标准的自注意力机制，额外添加一个高分辨率特征将导致计算时间呈二次增长，因为需要处理一个更大的矩阵，序列中的每个额外元素都会按元素基础增加计算负担。然而，FlexAttention采用的分层自注意力机制巧妙地缓解了这一问题，通过在添加高分辨率特征时保持线性关系，从而显著降低了计算负担。

5 Experiments

作者分别在高分辨率多模态基准测试以及一般多模态基准测试上评估FlexAttention方法，并将其与低分辨率的大型视觉-语言模型以及其他高分辨率方法[20, 35]进行比较。

Implementation

为了评估作者提出的FlexAttention在性能和效率方面的表现，作者将其整合到LLaVA-1.5-7b[35]模型中，形成了一个作者称之为LLaVA-FlexAttn的变体。输入分辨率设置为1008x1008，这是原始输入图像分辨率的 three 倍。然后，作者通过将这个变体与原始的LLaVA-1.5-7b模型进行比较，来展示利用高分辨率图像输入的优势。此外，作者还比较了FlexAttention与LLaVA-1.5-HD[35]和CogAgent[20]中用于实现高分辨率图像输入的方法，以展示作者提出方法的效率。

5.1.1 LLaVA-1.5-Hd [35]

在这个模型中，高分辨率图像标记与普通标记类似。它们与低分辨率图像标记连接在一起，并一同输入大型语言模型。由于该模型尚未公开发布，作者在LLaVA-1.5的代码库之上重新实现了它。作者以LLaVA-1.5-7b模型为基础模型。根据[35]中的设置，将高分辨率图像的输入分辨率设置为448x448。作者将这个 Baseline 称为LLaVA-HD。

5.1.2 CogAgent [20]

在这个模型中，作者使用跨注意力模块来感知高分辨率特征。在跨注意力模块中，高分辨率特征作为键（key）和值（value），而隐藏状态，包括低分辨率图像标记和文本标记，则作为 Query （query）。由于CogAgent是在文档和GUI风格数据上进行训练的，且数据处理和训练代码尚未公布，为了公平比较CogAgent中使用的的高分辨率操作符的有效性，作者将CogAgent推理代码库中的跨注意力模块转移到LLaVA-1.5上，并重新实现了训练代码。作者使用LLaVA-1.5-7b模型作为基础模型。将高分辨率图像的输入分辨率设置为1008x1008，以保持与作者的设置一致。作者将这个 Baseline 称为LLaVA-XAttn。

Datasets.

作者在四个高分辨率基准测试上进行了实验：V* Bench [53]，MagnifierBench [27]，TextVQA [45] 和 RSVQA-HRBEN [38]。前两个基准测试主要评估模型在一般高分辨率VQA上的能力，而后两个基准测试则专注于评估模型在特定领域高分辨率VQA上的表现，如TextVQA针对文本理解，RSVQA-HRBEN针对遥感应用。

Baselines.

作者对LLaVA-FlexAttn与两类视觉-语言模型（VLM）进行了比较分析：低分辨率VLM，具体包括InstructBLIP [11]、Otter [28]、MiniGPT-4 [62]、MiniGPTv2 [8]和LLaVA [35]；以及为本研究重新实现的 高分辨率VLM。此外，作者还与商用聊天机器人如GPT-4V [1]和专业VLM如GeoChat [26]进行了比较，以评估高分辨率图像输入能力的重要性。

Results.

表1展示了在两个高分辨率通用VQA基准测试上的评估结果。总体而言，三种高分辨率VLM的表现均优于低分辨率VLM，而作者的模型一致优于其他高分辨率VLM，在V* Bench上的总体准确率为54.5%，在MagnifierBench上的准确率为35.0%。与基础模型LLaVA-1.5-7b相比，V* Bench的总体准确率提高了6.9%，MagnifierBench的准确率提高了8.2%。与其他高分辨率方法相比，在计算量远低于其他高分辨率方法的情况下（TFLOPs降低了近30%，从LLaVA-HD的24.9降至17.1，以及比LLaVA-XAttn降低了超过37%，从27.1降至17.1），作者的方法取得了可比较甚至更高的准确率。关于TFLOPs和推理时间的详细讨论可在5.6节找到。得益于高分辨率特征选择和分层自注意力机制，作者的方法能够将输入图像分辨率在原始分辨率的基础上提升三倍，而计算成本的增长却呈次线性，实现了计算成本与准确率之间的更好权衡。与V* Bench上的GPT-4V相比，作者的方法显示出具有竞争力的性能，在空间类别上甚至取得了比GPT-4V更高的准确率，并且总体性能与GPT-4V相当。

表2展示了在两个高分辨率特定领域VQA基准测试上的结果。在RSVQA-HRBEN和TextVQA上，作者的LLaVA-FlexAttn一致优于基础模型和其他高分辨率方法。此外，在RSVQA-HRBEN基准测试上，作者的方法在总体准确率上超越了专门为远程遥感视觉问题回答基准设计的GeoChat模型。这一结果强调了结合高分辨率图像输入的有效性，表明高分辨率输入提供的额外细节和清晰度可以显著提高模型在特定VQA任务中对复杂视觉模式的处理和理解能力。

Evaluation on General Multimodal Benchmarks

Datasets and Baseline.

作者在几个多模态任务上评估了通用视觉-语言模型的性能，包括GQA [22]、VQAv2 [3]、POPE [34]、RefCOCO [57]、MM-Bench [37]、MME [15]和MM-Vet [58]。这一系列基准测试评估了模型的综合能力，包括空间理解、定位、避免幻觉的能力以及在学术导向任务中的表现。作者将作者的方法与基础模型LLaVA-1.5-7b进行了比较，以分析模型通用能力的变化。

结果。在表3中，作者展示了采用作者的FlexAttention后，在RefCOCO上的性能有所提升。RefCOCO要求根据参照表达式定位一个目标。因此，融入高分辨率特征可以减少识别小目标挑战的难度，并提高其位置预测的精度。在POPE上，作者实现了类似的幻觉率，并在大规模VQA基准测试上保持了相似的性能。这表明融入FlexAttention并不会影响模型的总体能力。

Ablation Study

特征选择策略。首先，作者进行消融研究，以验证作者方法关键设计的有效性，即使用注意力图选择高分辨率特征的策略。

作者将作者的注意力图选择策略与两种简单的 Baseline 策略进行比较：

1）随机选择，意味着随机选择高分辨率特征的一些 Patch ；

2）中心选择，意味着选择高分辨率特征的中心区域。选择比例保持在约与作者注意力图选择策略相同的10%左右。作者分别使用相同的微调数据集，并遵循相同的训练设置进行微调，然后在Magnifierbench和TextVQA上评估它们的性能。

分辨率的影响。作者还探讨了高分辨率图像尺寸的效果。默认设置为1008x1008，是原始低分辨率图像分辨率的3倍。此外，作者还引入了另外两种设置：672x672和1344x1344，分别使原始分辨率翻倍和四倍。作者分别使用相同的微调数据集并遵循相同的训练设置进行微调。作者在Magnifierbench基准上测量了它们的平均TFLOPs，并在Magnifierbench和TextVQA上评估了它们的性能。

物体大小的影响。对于第5.4节评估的通用基准，大多数问题并不关注小细节，因此无法揭示作者模型处理高分辨率图像的能力。为了更好地在通用基准上评估作者的模型，作者根据问题相关物体的尺寸将基准分为两个子集，将大于图像5%的物体分类为大物体，其余的为小物体。作者在RefCOCO验证集上进行实验，因为它提供了物体大小。

表4显示，在小物体上的准确率提升要远高于大物体。这表明即使对于像RefCOCO这样的非高分辨率基准，当问题涉及到小物体或详细信息时，作者的方法仍然可以提高准确性。

Inference Time on Hardware

作者在硬件上测量推理时间，以评估作者的FlexAttention的效率。模型使用PyTorch实现，并在单个NVIDIA V100 32G GPU上测量推理时间。作者在两个基准测试中测量了平均TFLOPs和总推理时间：在Magnifierbench中，模型的答案是单个字母；在TextVQA中，模型的答案是一个短句。在推理前进行预热和CUDA同步，以确保测量结果的准确性。

测量结果如表5所示。在Magnifierbench中，推理时间的减少与TFLOPs中理论计算成本减少的线性比例一致，作者的方法比两个 Baseline 分别快约30%和40%。在TextVQA中，速度优势略有下降，但仍然比 Baseline 快约15%和25%。需要注意的是，TextVQA的平均输出长度比Magnifierbench长，因此推理时间会更受生成阶段的影响，这一阶段受内存限制而非计算限制。作者在附录中提供了相关讨论。

6 Conclusion

在本论文中，作者提出了FlexAttention方法，旨在通过高效处理高分辨率图像输入并从中获益，来增强大型视觉-语言模型的能力。

通过利用动态高分辨率特征选择和分层自注意力机制，FlexAttention在性能和效率上都超越了现有高分辨率方法。FlexAttention背后的思想可以扩展到其他长序列模态，如视频或音频，这可能是未来研究的重要方向。

参考

[1].FlexAttention for Efficient High-Resolution Vision-Language Models.