Inf-MLLM：单个 4090D 实现 4M Token 长序列问答

一、背景

本文中我们简单介绍一个新的解决长序列推理效率的新方案 Inf-MLLM，也是基于 Token 稀疏化。有关 Token 稀疏化的方案通常可以从如下几个方面了解：

怎么识别关键 Token，包括静态识别或动态识别，比如 Streaming LLM 里 Attention Sink 和 Local Token。
Token 的稀疏度为多少，Token 太多会影响推理效率、占用过多内存，Token 太少又可能影响质量。
在哪个阶段执行 KV Cache 驱逐和更新，比如很多方案只在 Prefill 阶段执行稀疏化计算，这种方式比较适合超长输入，输出很短的场景。
计算时怎么使用这些 Token，这通常会涉及更高效的推理。

二、摘要

多模态大语言模型（MLLM）以其综合的多模态能力著称，并且有非常广泛的应用场景。然而，多模态往往意味着会产生较长的上下文，这给推理带来极大的挑战。长上下文推理需要缓存之前 Token 的 Key、Value Cache（KV Cache），这会导致过多的内存消耗以及延迟的增加。这也在很大程度上限制了 MLLM 在实际应用中的发展。

本文中，作者提出 Inf-MLLM，这是一种适用于 MLLM 的高效推理框架，可在单个 GPU 上实现无限上下文长度的 MLLM 流式推理。Inf-MLLM 主要基于作者对 LLM 和 MLLM 中 “Attention Saddle” 的观察，基于这种观察，Inf-MLLM 通过动态缓存最近的 Token 和相关 Token 来维护有限的 KV Cache。此外，Inf-MLLM 提出了注意力偏差，这是一种使 MLLM 能够捕获长期依赖的新方法。

作者实验表明，基于 Inf-MLLM，多个 LLM 和 MLLM 都能够在单个 GPU 上执行 4M Token 或 1 小时长视频的多轮对话。此外，Inf-MLLM 表现出比 StreamingLLM、H2O 等现有方法更好的 Stream 推理质量，并且可以加速 2x。

PS：相关文章我们已经介绍很多，对应的背景信息也多次提到，比如长序列时：Attention 计算和序列的二次方关系；存储 KV Cache 的巨大挑战；长程记忆能力；Memory Bound 问题。以及各种常见的 Token 稀疏化手段等。这里不再展开。

三、方法

3.1 洞察

作者通过可视化观察 MLLM 不同层的 Attention Map，发现了一些有利于 KV Cache 选择和驱逐的特性。如下图 Figure 2 所示，可以发现：

模式 1：最近的 Token 有更高的 Attention Score。对角线部分，这个应该也是共识。
模式 2：Video 对应的 Token 通常有比较高的 Attention Score，Video 中的初始 Token 甚至占据了 Attention Score 的 40%。（PS：但是下文的方案中似乎也没有特殊处理）
模式 3：有些垂直线对应的 Token 拥有很大的 Attention Score，这中现象也与 StreamingLLM 的起始几个 Token 类似。不过本文中作者发现 MLLM 中这些 Token 可能是动态的，不一定在起始位置。
模式 4：随着多轮对话的进行，高 Attention Score 会向前偏移。

作者将上述与当前 Decoding Token 最相关的这些 Token 称作 Attention Saddle，借鉴了数学中的 Saddle Point。为了识别和存储这些关键 Token，作者提出两种技术手段：

设计了一种 KV Cache 驱逐机制，来驱逐不相关的 Token，用于维护上述 Pattern 1/2/3 中提到的 Token。
提出了 Attention 偏置来动态的更新和捕获偏移特性，以增强长程记忆能力，对应 Pattern 4。

3.2 KV Cache 驱逐和更新

如下图 Figure 3 所示，在 Streaming 推理场景，当一个新的 Prompt 到来时（PS：也就是只在多轮对话的每个 Prefill 阶段执行，不在单个 Token 的 Decoding Step 执行）:

计算 Prompt 对应的 Q、K 和 V 并存储在一个长度为 L 的检索窗口内。
使用窗口内的 Q 计算与历史（之前轮，假设共 t 个 Token）K Cache 的 Attention Score，对应中间灰色部分。
计算窗口内部 Q 与之前 K 的 Attention Score，对应中间蓝色部分。
最终得到一个 L * n 的分数矩阵。
按照 L 维度对 Attention Score 进行平均，相当于 L 中每个 Token 与窗口中每个 Token 的 Attention Score 的平均，得到一个 L 维的 Attention Score 向量。这样主要是为了提升计算效率。
从 L 个 Attention Score 中挑选出 Top t 个，并从 KV Cache 中驱逐其他不重要的 Token。

需要说明的是：上述过程中的 t 个 KV Cache 已经是稀疏化之后的，因此相比原始完整 Token 序列要短的多，因此相应的代价也就不是特别高。

针对所有自学遇到困难的同学们，我帮大家系统梳理大模型学习脉络，将这份 LLM大模型资料 分享出来：包括LLM大模型书籍、640套大模型行业报告、LLM大模型学习视频、LLM大模型学习路线、开源大模型学习教程等, 😝有需要的小伙伴，可以 扫描下方二维码领取🆓↓↓↓

👉[CSDN大礼包🎁：全网最全《LLM大模型入门+进阶学习资源包》免费分享（安全链接，放心点击）]()👈

3.3 Attention 偏置

为了进一步增强 KV Cache 驱逐能力，尤其是在长上下文和多轮对话中，需要解决一些问题。首先，由于 Softmax 的存在，随着序列长度增加，Attention Score（权重）的总和始终为 1，也就是每个 Token 的权重会逐渐降低，高分 Token 的识别难度变大。其次，经过几轮 KV Cache 驱逐之后，剩余 Token 之间的 Attention Score 分布变得不均匀，一些 Token 的 Attention Score 会因为多轮积累而得到提升。如下图 Figure 4 所示，这种现象会阻碍识别与当前对话轮次更相关的 Attention Saddle，导致不恰当的 KV Cache 驱逐，甚至在长上下文中经过几轮推理后 KV Cache 不再更新而导致性能大幅下降的情况。

为了在 Streaming 推理中持续更新 KV Cache，作者提出 Attention Bias，将注意力焦点转移到最新的上下文中。如下图 Figure 4 所示，Attention Bias 可以调整 Attention Score 的分布，并使得多轮视频对话可以持续进行。具体来说，在计算检索窗口中的平均 Attention Score 后，向其中添加 Attention Bias 以促使 KV Cache 丢弃更久以前的 Token。当 Attention Bias 较高时，KV Cache 倾向于保留更多新 Token；当 Attention Bias 较低时，KV Cache 倾向于保留更早之前的 Token，以捕获更长期的依赖关系。因此，适当调整 Attention Bias 可以在确保长上下文 Streaming 推理的同时保留长期依赖关系。