arXivDaily arXiv每日学术速递 周一至周五更新
arXiv 2502.02406cs.CVcs.AIcs.DCcs.LG

LV-XAttn:多模态大语言模型中针对长视觉输入的分布式交叉注意力

LV-XAttn: Distributed Cross-Attention for Long Visual Inputs in Multimodal Large Language Models

University of Wisconsin-Madison(威斯康星大学麦迪逊分校)

Tzu-Tao Chang, Shivaram Venkataraman

AI总结:

针对多模态大模型长视觉输入下交叉注意力分布式计算的通信瓶颈,提出LV-XAttn机制,通过本地保留key-value块并跨GPU交换较小query块及激活重计算技术,实现高达10.62倍的端到端加速。

摘要:

多模态大语言模型(MLLMs)通常采用交叉注意力将视觉信息整合到语言主干中。然而,在视频理解等大型视觉输入应用中,在交叉注意力层处理大量视觉token会导致高内存需求,并通常需要在多个GPU间进行分布式计算。现有的分布式注意力机制面临显著的通信开销,使得交叉注意力层成为MLLMs高效训练和推理的关键瓶颈。为此,我们提出LV-XAttn,一种通信开销极小的分布式精确交叉注意力机制。我们观察到,在涉及大型视觉输入的应用中,query块的大小通常远小于key-value块。因此,在LV-XAttn中,我们将大型key-value块保留在每个GPU本地,并在GPU间交换较小的query块。我们还引入了高效的激活重计算技术以支持更长的视觉上下文。我们从理论上分析了LV-XAttn的通信优势,并表明它能在广泛的模型中实现加速。我们使用Llama 3-V、mPLUG-Owl3和OpenFlamingo模型进行评估,发现与现有方法相比,LV-XAttn实现了高达10.62倍的端到端加速。

英文摘要:

Cross-attention is commonly adopted in multimodal large language models (MLLMs) for integrating visual information into the language backbone. However, in applications with large visual inputs, such as video understanding, processing a large number of visual tokens in cross-attention layers leads to high memory demands and often necessitates distributed computation across multiple GPUs. Existing distributed attention mechanisms face significant communication overheads, making cross-attention layers a critical bottleneck for efficient training and inference of MLLMs. To address this, we propose LV-XAttn, a distributed, exact cross-attention mechanism with minimal communication overhead. We observe that in applications involving large visual inputs, the size of the query block is typically much smaller than that of the key-value blocks. Thus, in LV-XAttn we keep the large key-value blocks locally on each GPU and exchange smaller query blocks across GPUs. We also introduce an efficient activation recomputation technique to support longer visual context. We theoretically analyze the communication benefits of LV-XAttn and show that it can achieve speedups for a wide range of models. Our evaluations with Llama 3-V, mPLUG-Owl3 and OpenFlamingo models find that LV-XAttn achieves up to 10.62$\times$ end-to-end speedup compared to existing approaches.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑