1. 论文基本信息

1.1. 标题

混合注意力跨度：利用异构滑动窗口长度优化 LLM 推理效率 (Mixture of Attention Spans: Optimizing LLM Inference Efficiency with Heterogeneous Sliding-Window Lengths)

1.2. 作者

Tianyu Fu, Haofeng Huang, Xuefei Ning, Genghan Zhang, Boju Chen, Tianqi Wu, Hongyi Wang, Zixiao Huang, Shiyao Li, Shengen Yan, Guohao Dai, Huazhong Yang, Yu Wang 作者团队主要来自 清华大学 (Tsinghua University) 和 Infinigence-AI，部分作者来自 斯坦福大学 (Stanford University) 和 上海交通大学 (Shanghai Jiao Tong University)。

1.3. 发表期刊/会议

预印本文章，发表于 arXiv。虽然是预印本，但其主题与大型语言模型 (LLM) 推理优化紧密相关，该领域是当前人工智能研究的热点，具有重要的实际应用价值。

1.4. 发表年份

2024年

1.5. 摘要

大型语言模型 (LLM) 在长上下文情境下，其注意力机制在内存和吞吐量方面面临巨大挑战，而滑动窗口注意力 (sliding-window attention) 提供了一种硬件高效的解决方案。现有方法通常在所有注意力头 (attention heads) 和输入大小上采用单一的窗口长度，但这种统一的方法未能捕捉 LLM 固有的异构注意力模式，也忽略了它们独特的准确性-延迟权衡。为解决此问题，论文提出了 混合注意力跨度 (Mixture of Attention Spans, MoA)，它能自动为不同的注意力头和层 (layers) 定制不同的滑动窗口长度配置。MoA 构建并探索了一个包含各种窗口长度及其相对于输入大小的缩放规则的搜索空间。它通过对模型进行剖析 (profiles the model)，评估潜在配置，并精确找出每个注意力头的最佳长度配置。MoA 能够适应不同的输入大小，揭示了一些注意力头会扩展其关注范围以适应更长的输入，而其他注意力头则始终专注于固定长度的局部上下文。实验表明，在相同的平均滑动窗口长度下，MoA 将有效上下文长度提高了 3.9 倍，在 $Vicuna-{7B, 13B}$ 和 $Llama3-{8B, 70B}$ 模型上，相较于统一窗口基线，检索准确率提高了 1.5-7.1 倍。此外，MoA 缩小了与全注意力 (full attention) 之间的性能差距，在三个长上下文理解基准测试中，最大相对性能下降从 9%-36% 降至 5% 以内。MoA 还实现了 1.2-1.4 倍的 GPU 内存减少，并将解码吞吐量 (decode throughput) 相较于 FlashAttention2 和 vLLM 分别提高了 6.6-8.2 倍和 1.7-1.9 倍，同时对性能影响极小。

1.6. 原文链接

原文链接: https://arxiv.org/abs/2406.14909 PDF 链接: https://arxiv.org/pdf/2406.14909v3.pdf

2. 整体概括

2.1. 研究背景与动机

大型语言模型 (LLM) 在各种应用中展现出卓越的能力，其核心在于注意力机制，它通过计算词元 (token) 之间的交互来实现上下文理解。为了增强 LLM 的能力，如事实检索、摘要、少样本学习和问答等，扩展输入长度至关重要。然而，随着输入长度的增长，注意力计算和键值缓存 (Key-Value Cache, KV-Cache) 会带来显著的效率挑战，包括内存消耗和吞吐量下降。

现有的滑动窗口注意力方法通常在所有注意力头和不同输入长度上采用统一、固定的注意力窗口长度。这种设计虽然能将注意力限制在局部范围，从而在有界定的计算和 KV-Cache 开销内处理长输入，但存在以下具体挑战和空白：

• 未能捕捉异构注意力模式： LLM 的不同注意力头具有异构的注意力模式。一些注意力头可能专注于局部上下文，而另一些则需要覆盖整个输入序列的全局上下文（如原文 Figure 2 所示）。统一的窗口长度限制了全局上下文注意力头的有效范围，同时又可能对局部上下文注意力头分配了过多的计算和内存预算。

• 无法有效扩展有效上下文长度： 尽管理论上固定跨度的局部注意力可以通过多层模型逐渐聚合全局信息，从而产生比单个注意力跨度更长的有效上下文长度，但论文发现，如 StreamingLLM 等统一滑动窗口方法，其有效上下文长度很难超出其设定的窗口跨度（如原文 Figure 1(b) 所示）。

• 缺乏弹性缩放规则： 随着输入长度的增加，某些注意力头需要比其他头更快地增加注意力跨度，以避免严重的性能下降（如原文 Table 1 所示），而统一方法缺乏针对不同注意力头差异化缩放注意力跨度的异构规则。

• 压缩配置与长上下文任务不匹配： 现有模型压缩方法（如量化和稀疏注意力）通常依赖通用语言建模语料库来确定压缩配置，但这未能准确捕捉其对长上下文任务的影响。

  这篇论文的切入点是识别并解决上述统一滑动窗口注意力方法的局限性，通过引入一种能够自动为每个注意力头和层定制异构、弹性滑动窗口长度的方法，以更好地平衡准确性与效率。

2.2. 核心贡献/主要发现

论文的主要贡献和发现可以总结如下：

• 提出了异构弹性规则 (Heterogeneous Elastic Rules)： 论文设计了 Mixture of Attention Spans (MoA) 方法，它能为每个注意力头自动定制不同的滑动窗口长度配置，并使其具有弹性，即窗口长度可以根据输入长度进行自适应缩放。这种异构弹性规则将内容检索准确率从 25% 提高到 98%，显著优于统一窗口的基线方法。
• 优化了校准数据集构建 (Calibration Dataset Construction)： 论文强调了校准数据集设计在 LLM 压缩中的重要性。研究发现，使用具有长距离依赖性、并以原始 LLM 生成的响应作为参考的校准数据集，比使用通用语言建模数据集和人类编写的摘要能更准确地剖析注意力影响，从而实现更好的性能。
• 提出自动配置搜索流水线 (Automatic Configuration Search Pipeline)： 论文开发了一个自动流水线，用于高效地搜索最佳的 MoA 配置，即为每个注意力头找到最优的异构弹性规则。这个流水线能在数小时内（例如 Vicuna-13B 仅需两小时）完成搜索，大大降低了配置成本。
• 显著的性能提升： MoA 在相同的平均滑动窗口长度下，将有效上下文长度提高了 3.9 倍，在 $Vicuna-{7B, 13B}$ 和 $Llama3-{8B, 70B}$ 模型上，检索准确率相较于统一窗口基线提高了 1.5-7.1 倍。
• 缩小与全注意力模型的性能差距： 在三个长上下文理解基准测试中，MoA 将与全注意力模型的最大相对性能下降从 9%-36% 降低到 5% 以内，平均相对性能下降低于 1%，远优于其他统一滑动窗口方法。
• 显著的效率提升： MoA 实现了 1.2-1.4 倍的 GPU 内存减少，并将解码吞吐量相较于 FlashAttention2 和 vLLM 分别提高了 6.6-8.2 倍和 1.7-1.9 倍，同时对性能影响极小。即使在 25% 的密度下，MoA 也能实现超过 90% 的检索准确率，显著优于需要 75%-100% 密度才能达到类似性能的基线。

3. 预备知识与相关工作

3.1. 基础概念

理解本文需要掌握以下核心概念：

• 大型语言模型 (Large Language Models, LLMs)： 指的是参数量巨大、在海量文本数据上预训练的神经网络模型，如 GPT、Llama 等。它们在各种自然语言处理任务中表现出色，但计算和内存开销巨大。
• 注意力机制 (Attention Mechanism)： 由 Vaswani 等人于 2017 年在论文《Attention Is All You Need》中提出，是 Transformer 架构的核心。它允许模型在处理序列数据时，动态地权衡不同词元之间的关联强度，从而捕获长距离依赖关系。
  • 自注意力 (Self-Attention)： 在序列内部计算词元之间的注意力。每个词元都可以关注序列中的所有其他词元。
  • 多头自注意力 (Multi-Head Self-Attention, MHA)： 将自注意力机制并行地运行多次，每次使用不同的线性投影（即不同的 Query (Q)、Key (K) 和 Value (V) 矩阵），然后将多个“头”的输出拼接起来并再次线性投影。这使得模型能够从不同的表示子空间中学习信息，并关注序列的不同方面。
  • 注意力计算公式： $$\mathbf{S} = \mathbf{Q}\mathbf{K}^T \\ \mathbf{A} = \mathrm{softmax}(\mathbf{S} + \mathbf{M}) \\ \mathbf{O} = \mathbf{A}\mathbf{V}$$ 其中：
    • $\mathbf{Q}$ (Query)、$\mathbf{K}$ (Key)、$\mathbf{V}$ (Value) 是输入序列通过线性变换得到的矩阵。它们通常是由输入词嵌入经过不同的权重矩阵投影而来。
    • $\mathbf{K}^T$ 是 $\mathbf{K}$ 的转置。
    • $\mathbf{S}$ 是 $Q$ 和 $K$ 的点积，代表查询和键之间的相似度分数。
    • $\mathbf{M}$ 是一个掩码矩阵 (mask matrix)，用于在自回归生成任务中防止模型关注未来的词元（因果掩码 causal mask），或用于实现稀疏注意力模式。对于因果掩码，未来词元对应的 $\mathbf{S}$ 值会被设置为负无穷，使其在 softmax 后变为 0。
    • $\mathrm{softmax}(\cdot)$ 函数将相似度分数转换为注意力权重，使其和为 1。
    • $\mathbf{A}$ 是注意力权重矩阵。
    • $\mathbf{O}$ 是注意力机制的输出，是 $A$ 和 $V$ 的加权和，代表了结合了上下文信息的词元表示。
• 键值缓存 (Key-Value Cache, KV-Cache)： 在 LLM 的自回归推理 (autoregressive inference) 过程中，为了避免重复计算先前词元对应的 Key 和 Value 矩阵，这些矩阵会被缓存下来。随着生成序列的增长，KV-Cache 的大小也会线性增长，成为内存和计算的瓶颈。
• 预填充阶段 (Prefill Stage)： 在 LLM 推理的初始阶段，模型处理整个输入提示 (prompt) 序列以生成第一个响应词元。此阶段通常涉及大量的并行计算。
• 解码阶段 (Decode Stage)： 在预填充之后，模型进入迭代生成阶段。它每次生成一个新词元，并使用新生成的词元和先前缓存的 KV 矩阵（即 KV-Cache）来生成下一个词元。此阶段计算量相对较小，但 KV-Cache 的增长是主要瓶颈。
• 滑动窗口注意力 (Sliding-Window Attention)： 一种静态稀疏注意力 (static sparse attention) 方法。它通过对注意力矩阵应用一个预定义的掩码来限制每个词元只能关注其周围固定长度的局部上下文（即一个“窗口”内的词元）。这种方法可以显著减少内存和计算开销，因为 KV-Cache 的大小被限制在窗口长度。
• 注意力汇聚点 (Attention Sink)： 在滑动窗口注意力中，为了确保模型能够访问一些重要的全局信息（例如，提示的开头部分），通常会保留序列开头的少数词元（例如 64 个）始终可见，不被窗口掩码限制。这些始终可见的词元被称为 attention sink。

3.2. 前人工作

3.2.1. 高效注意力方法分类

现有高效注意力方法主要分为两类：

• 动态稀疏注意力 (Dynamic Sparse Attention)：
  • 动态跳过预填充计算 (Dynamic Skip Prefill Computations)： 在预填充阶段动态地跳过部分注意力计算，例如 Pagliardini et al., 2023; Qu et al., 2022。这类方法通常需要复杂的控制流和特定硬件支持才能实现显著的加速。
  • 动态丢弃 KV-Cache (Dynamic KV-Cache Pruning)： 在解码阶段根据输入序列动态地修剪 KV-Cache，例如 Anagnostidis et al., 2023; Zhang et al., 2023; Ge et al., 2023。这可能需要大量的重新训练、额外的修剪分数计算，或复杂的内存交换策略。
• 静态稀疏注意力 (Static Sparse Attention)：
  • 预定义掩码 (Predefined Masks)： 采用预定义的掩码在所有输入序列上一致地应用。由于计算流固定，这类方法通常更高效且对 GPU 友好。
  • 用于语言理解模型： 对于像 BERT (Devlin et al., 2018) 这样的语言理解模型，已经使用了各种掩码，如 BigBird (Zaheer et al., 2020)、Longformer (Beltagy et al., 2020)、Sparse Transformers (Child et al., 2019)。
  • 用于生成式 LLM： 主流方法是结合少量初始词元的全局注意力 (global attention) 的统一滑动窗口掩码，例如 StreamingLLM (Xiao et al., 2024c) 和 LM-Infinite (Han et al., 2023)。这些方法可以丢弃超出局部注意力范围的 KV-Cache，从而节省长序列场景下的内存。

3.2.2. LLM 加速框架与内核优化

• LLM 加速框架： DeepSpeed-Inference (Aminabadi et al., 2022)、Accelerate (Gugger et al., 2022)、FlexGen (Sheng et al., 2023)、vLLM (Kwon et al., 2023) 等框架旨在提高 LLM 的训练和推理效率。
• 内核级优化： FlashAttention (Dao et al., 2022) 及其后续版本 FlashAttention-2 (Dao, 2024) 和 FlashAttention-3 (Shah et al., 2024) 等通过优化 GPU 内存访问模式和并行计算，显著提高了注意力机制的计算速度和内存效率。

3.3. 技术演进

LLM 的效率优化从早期对 Transformer 架构的稀疏化探索开始，如 Longformer、BigBird 等，主要通过设计固定的稀疏注意力模式来处理长序列。随着 LLM 规模的爆炸式增长，KV-Cache 的内存瓶颈和自回归解码的吞吐量问题变得突出。StreamingLLM 和 LM-Infinite 提出了统一的滑动窗口注意力与 attention sink 结合的方案，以限制 KV-Cache 的增长。同时，硬件加速（如 FlashAttention）和系统级优化（如 vLLM 的 PagedAttention）也在不断推动效率边界。然而，这些方法大多采用“一刀切”的统一策略，未充分考虑 LLM 内部注意力头异构性的问题。

3.4. 差异化分析

本文提出的 MoA 方法与相关工作的主要区别和创新点在于：

• 异构性 (Heterogeneity)： 现有滑动窗口方法（如 StreamingLLM）通常对所有注意力头和层使用统一的窗口长度。MoA 则针对不同的注意力头和层自动定制不同的窗口长度。这是基于论文观察到的 LLM 注意力模式的异构性（如原文 Figure 2 所示）。

• 弹性规则 (Elastic Rules)： 现有方法通常采用固定长度的窗口，或者窗口长度的缩放规则是统一的。MoA 引入了“弹性规则”，即注意力跨度 $S$ 可以是输入长度 $N$ 的线性函数 $S_h = \alpha_h + \beta_h \cdot N$，其中 $α$ 和 $β$ 是每个注意力头独有的超参数。这意味着 MoA 的窗口长度不仅因头而异，还能根据输入长度自适应地变化。

• 自动搜索流水线 (Automatic Search Pipeline)： MoA 提供了一个训练无关 (training-free) 的自动流水线，通过梯度剖析 (gradient-based profiling) 和多目标优化 (multi-objective optimization) 来寻找每个注意力头的最佳异构弹性规则，而不是手动设置或通过通用数据集预设。

• 校准数据集设计： MoA 强调使用具有长距离依赖性、并与原始 LLM 响应对齐的校准数据集进行剖析，以更准确地捕捉注意力影响，这与许多依赖通用语言建模数据集的方法不同。

• 性能和效率： 结合上述创新点，MoA 在保持与全注意力模型相当的性能（显著缩小性能差距）的同时，实现了比现有先进稀疏注意力方法更高的吞吐量和更低的内存消耗，并在更低的 KV-Cache 密度下达到更高的准确率。

  简而言之，MoA 的核心在于从“统一、静态”向“异构、弹性、自动优化”的注意力窗口管理转变，从而更精细地匹配 LLM 内部复杂的注意力需求，实现性能和效率的显著提升。

4. 方法论

Mixture of Attention Spans (MoA) 旨在通过自动为每个注意力头和层定制异构且弹性的滑动窗口长度配置，以优化 LLM 的推理效率。其核心思想是识别不同注意力头对上下文长度的需求差异，并据此分配注意力资源。

4.1. 方法原理

MoA 的核心直觉来源于对预训练 LLM 注意力模式的观察：不同注意力头表现出异构的注意力模式（如原文 Figure 2 所示），一些头专注于局部上下文，而另一些则需要全局上下文。此外，这些注意力头对输入长度变化的响应也不同，即具有异构的弹性行为。例如，某些头可能随着输入长度的增加而扩展其关注范围，而另一些则始终保持固定长度的局部关注。

基于这些观察，MoA 认为统一的滑动窗口长度无法有效满足所有注意力头的需求，导致性能瓶颈。因此，MoA 提出为每个注意力头定制其注意力跨度 (attention span) 的规则，使其能够根据输入长度灵活调整，从而在保证模型性能的同时，最大限度地减少 KV-Cache 的开销。

4.2. 核心方法详解

MoA 的方法主要包括三个部分：定义异构弹性规则的搜索空间、自动配置搜索流水线（注意力影响剖析和自动优化）以及校准数据集的构建。

4.2.1. 异构弹性规则搜索空间

MoA 采用硬件友好的滑动窗口掩码作为基础注意力掩码。与现有工作类似，序列开头的少量词元（MoA 默认 64 个）不被掩码，作为注意力汇聚点 (attention sink)。注意力跨度 (attention span) 等于滑动窗口跨度加上这些初始未被掩码的词元数量。

MoA 为每个注意力头 $h$ 定义了一个注意力跨度 $S_h$，它是一个关于输入长度 $N$ 的线性函数： $$S_h = \alpha_h + \beta_h \cdot N$$ 其中：

• $S_h$ 是注意力头 $h$ 在输入长度为 $N$ 时的注意力跨度。

• $\alpha_h$ 是一个超参数，控制注意力头 $h$ 的基本（或最小）注意力跨度。

• $\beta_h$ 是一个超参数，控制注意力头 $h$ 的注意力跨度随输入长度 $N$ 变化的膨胀率（或缩放系数）。

• $\alpha_h$ 和 $\beta_h$ 的值从多个离散选项中选择。例如，MoA 默认使用 6 个 $\alpha$ 选项和 9 个 $\beta$ 选项。

  这个公式也被更一般地定义在附录 A.2 中，用于表示弹性规则 $f(\cdot)$： $$\mathcal{S}(N; \alpha, \beta) = \alpha \times N + \beta$$ 这里，$\mathcal{S}$ 表示给定输入长度 $N$ 和参数 $\alpha, \beta$ 下的注意力跨度。

每个注意力头采用不同的 $(\alpha, \beta)$ 组合，从而产生异构的注意力跨度。对于一个具有多头和多层的 LLM，所有注意力头的弹性规则集合构成了 MoA 的配置。由于搜索空间可能非常大（例如，一个 7B 模型可能有 $54^{1024}$ 种组合），MoA 需要一个自动化的流水线来高效地找到最佳的 $\alpha$ 和 $\beta$。

注意力掩码的可视化与构成

原文 Figure 3(a) 形象地展示了 MoA 搜索空间中的注意力掩码。

• 垂直紫色条纹 (Vertical purple stripe): 代表初始未被掩码的词元（attention sink），固定为 64 个词元。
• 对角紫色带 (Diagonal purple band): 代表滑动窗口区域，其长度由弹性规则确定。
• 白色单元格 (White cells): 代表被修剪的词元，其注意力被设置为 0。 每个注意力掩码的形成是滑动窗口跨度与固定前缀（attention sink）之和。

4.2.2. 自动 MoA 配置搜索流水线

MoA 的自动配置搜索流水线如原文 Figure 3(b) 所示，包括三个主要步骤：校准数据集构建、注意力影响剖析和自动优化。

4.2.2.1. 校准数据集与监督

在剖析步骤之前，MoA 需要一个精心设计的校准数据集。传统上，LLM 压缩方法常使用通用语言建模数据集（如 RedPajama），并以下一个词元预测为监督目标。然而，这种方法缺乏长上下文依赖性，且人类编写的监督文本与模型自身的响应可能存在偏差。

MoA 强调使用具有长距离依赖性且与模型对齐 (model alignment) 的校准数据集。

• 长距离依赖性： 采用 MultiNews (Fabbri et al., 2019) 等包含长文本摘要的数据集，这些摘要高度依赖于长距离内容。
• 模型对齐： 使用原始密集模型生成的摘要作为监督文本，而非人类编写的摘要。这确保了在计算注意力值及其梯度时，模型自身注意力模式与监督文本之间的对齐，从而更准确地剖析注意力影响。

4.2.2.2. 注意力影响剖析 (Attention Influence Profiling)

此步骤旨在量化每个注意力值对预训练 LLM 最终预测损失的影响。它为后续优化步骤提供了关于掩盖每个注意力值所带来的准确性权衡信息。

注意力影响的计算基于注意力矩阵 $\mathbf{A}$ 及其梯度 $\partial L / \partial \mathbf{A}$，这些是在校准数据集上计算的。当应用滑动窗口注意力掩码时，模型预测损失的变化 $\Delta L$ 通过一阶泰勒展开近似。

对于注意力头 $h$、行 $i$ 和列 $j$ 的注意力值 $A_{h,i,j}$ 被掩码（即设为零）时，损失的估计变化 $E_{h,i,j}$ 定义为注意力影响矩阵。这个度量考虑了掩码的直接和间接效应：

1. 直接效应： 掩码直接将 $A_{h,i,j}$ 降至零，表示为 $\Delta A_{h,i,j | j} = -A_{h,i,j}$。

2. 间接效应： softmax 函数将注意力矩阵每行的和归一化为 1。因此，将列 $j$ 处的某个注意力值设为零，会导致同一行中其他注意力值 $A_{h,i,n | j}, n \neq j$ 的相对增加。

   这两种效应被整合到以下公式中（为简化表示，省略了头索引 $h$）： $$E_{i,j} = \sum_n \frac{\partial L}{\partial A_{i,n}} \cdot \Delta A_{i,n|j} = \frac{\partial L}{\partial A_{i,j}} \cdot (-A_{i,j}) + \sum_{n \neq j} \frac{\partial L}{\partial A_{i,n}} \cdot A_{i,n} \cdot \frac{A_{i,j}}{1 - A_{i,j}}$$ 其中：

• $E_{i,j}$ 是当注意力值 $A_{i,j}$ 被掩码时，模型预测损失的估计变化（注意力影响）。

• $L$ 是模型在校准数据集上的预测损失。

• $\partial L / \partial A_{i,n}$ 是损失 $L$ 对注意力矩阵中第 $i$ 行第 $n$ 列元素 $A_{i,n}$ 的梯度。

• $\Delta A_{i,n|j}$ 是当 $A_{i,j}$ 被掩码时，注意力矩阵中第 $i$ 行第 $n$ 列元素的变化量。

  在实际操作中，通过反向传播在校准数据集上计算每个注意力头 $\bar{\mathbf{E}}_h$ 的平均注意力影响，并针对不同的输入长度分别计算。这使得 MoA 能够根据应用候选弹性规则所掩码的注意力值之和来计算在特定输入长度下的准确性损失。

若将 $M_{r_h}$ 定义为与规则 $r$ 对应的头 $h$ 的二值掩码（被掩码位置为 1，其他为 0），则准确性损失 $\Delta L$ 形式化为： $$\Delta L = \sum_h \Delta L_{h,r_h} = \sum_h \sum_i \sum_j M_{r_h,i,j} \cdot \bar{E}_{h,i,j}$$ 其中：

• $\Delta L$ 是总的准确性损失。

• $\Delta L_{h,r_h}$ 是注意力头 $h$ 应用规则 $r_h$ 时的准确性损失。

• $M_{r_h,i,j}$ 是掩码 $M_{r_h}$ 中第 $i$ 行第 $j$ 列的元素。

• $\bar{E}_{h,i,j}$ 是注意力头 $h$ 的平均注意力影响矩阵中第 $i$ 行第 $j$ 列的元素。

  通过剖析阶段，MoA 获得了弹性规则在不同注意力头上的估计准确性影响，这有助于将更长的窗口分配给更敏感的头，将更短的窗口分配给不敏感的头。

4.2.2.3. 自动优化 (Automatic Optimization)

MoA 自动为每个注意力头选择最优的弹性规则，以在密度预算下最小化不同序列长度的准确性损失。

这个过程是一个多目标优化问题。首先，MoA 识别出帕累托前沿 (Pareto-front) 配置，即在某个剖析长度上改善准确性损失必然会恶化另一个长度的准确性损失。为了确保对未剖析长度的最佳泛化能力，MoA 随后从帕累托前沿解决方案中选择在未见验证长度上产生最小损失的配置作为最终配置。

具体来说，目标函数是在剖析长度上最小化总准确性损失，同时遵守用户定义的密度约束： $$\underset{r_h \in \mathbb{R}}{\arg\min} \Delta L^{(N_i)}, N_i \in \mathbb{N}_{\mathrm{profile}} \quad \mathrm{s.t.} \frac{1}{H} \sum_{h=1}^H d_{r_h}^{(N_i)} \leq d_{\mathrm{constr}}^{(N_i)}, \forall N_i \in \mathbb{N}_{\mathrm{constr}}.$$ 其中：

• 上标 (N) 表示在不同长度下的值。

• $\mathbb{N}_{\mathrm{profile}}$ 和 $\mathbb{N}_{\mathrm{constr}}$ 分别表示用于剖析和受密度约束的长度集合。

• $\mathbb{R}$ 表示候选规则集合。

• $\Delta L^{(N_i)}$ 表示应用注意力掩码引起的准确性损失。

• $d_{r_h}^{(N_i)}$ 表示注意力头 $h$ 的规则 $r_h$ 在长度 $N_i$ 下的密度 (density)。

• $d_{\mathrm{constr}}^{(N_i)}$ 是在长度 $N_i$ 下的密度约束。

• $H$ 是注意力头的总数。

  这个公式对应于经典的多目标混合整数规划 (Mixed-Integer-Programming, MIP) 问题，可以使用现有的线性求解器（如 Gurobi）在几分钟内有效求解。

在附录 E.3.1 中，优化问题被转换为标准的 MIP 框架，引入了二值变量 $X_{h,r_h} \in \{0, 1\}$，表示是否为注意力头 $h$ 选择规则 $r_h$。假设模型有 $H$ 个注意力头，每个头 $h$ 有 $R_h$ 条弹性规则： $$\begin{array}{c} \underset{\boldsymbol{X}}{\mathrm{argmin}} \frac{1}{H} \sum_{h=0}^{H-1} \sum_{r_h=0}^{R_h-1} \Delta L_{h,r_h} X_{h,r_h} \quad \mathrm{s.t.} \\ \sum_{r_h=0}^{R_h-1} X_{h,r_h} = 1, \quad h \in \{0, \cdots, H-1\} \quad \text{(16b)} \\ \frac{1}{H} \sum_{h=0}^{H-1} \sum_{r_h=0}^{R_h-1} d_{r_h} X_{h,r_h} \leq d_{\mathrm{constr}} \quad \text{(16c)} \\ X_{h,r_h} \in \mathbb{Z}, \quad \forall h \in \{0, \cdots, H-1\}, \forall r_h \in \mathbb{R} \quad \text{(16d)} \end{array}$$ 其中：

• (16a) 是目标函数，旨在最小化总损失。

• (16b) 约束条件确保每个注意力头精确选择一个 MoA 配置。

• (16c) 约束条件确保平均密度不超过预定义的 $d_{\mathrm{constr}}$。

• (16d) 强制 $X_{h,r_h}$ 是二值变量。

  为了限制每个模型层中不同规则的数量，还可以对单层内选择的 $X_{h,r_h}$ 的范数进行约束。

多长度优化

对于多个长度的优化，MoA 采用了 epsilon-constraint 方法 (Yv et al., 1971)，将多目标 MIP 问题转化为多个单目标 MIP 问题。 原文 Figure 15 详细说明了这个过程：

1. 确定损失范围： 首先，对每个目标（即每个输入长度）单独进行单目标优化，并记录其他目标对应的损失值。这确定了每个输入长度的损失范围 $R^{(N_j)} = [\min_i \Delta L_i^{(N_j)}, \max_i \Delta L_i^{(N_j)}]$。
2. 迭代优化与约束： 然后，再次迭代每个目标。这一次，将其他目标的损失范围分割成 $M$ 个均匀区间 $S_k^{(N_j)}$，并将这些区间作为额外的约束条件加入到当前单目标优化中。 $$\underset{r_h \in \mathbb{R}}{\arg\min} \Delta L^{(N_i)} \quad \mathrm{~s.t.~} \frac{1}{H} \sum_{h=1}^H d_{r_h}^{(N_i)} \leq d_{\mathrm{constr}}^{(N_i)}, \forall N_i \in \mathbb{N}_{\mathrm{constr}}; \quad \Delta L^{(N_j)} \in S_{k_j}^{(N_j)}, \forall j \neq i.$$ 这个优化对每个 $i$ 和每个 $j$ 的 $k_j$ 独立变化进行。最终，移除不符合帕累托前沿要求的方案，得到最终的帕累托前沿集合。

4.2.3. 注意力头语义与规则的关联

MoA 通过统计分析和优化来保留原始密集模型学到的语义。注意力值 $A_{h,i,j}$ 受到位置编码 $P_h(i,j)$ 和词元嵌入 $T_h(t_i,t_j)$ 的影响： $$A_{h,i,j} = \mathbb{A}_h(t_i, t_j, i, j) = f_h\left(P_h(i, j), T_h(t_i, t_j)\right)$$ 其中 $t_i$ 和 $t_j$ 分别是第 $i$ 个和第 $j$ 个词元。 通过线性近似 $A_{h,i,j} = \alpha_h P_h(i,j) + (1-\alpha_h)T_h(t_i,t_j)$，并假设词元位置和内容不相关，可以推导出注意力值对不同句子 $s$ 的期望 $\mathbb{E}_t[A_{h,i,j}]$。进一步，通过计算期望的标准差 (Standard division of Expectation, SoE) $\sigma_p(\mathbb{E}_t[A_{h,i,j}])$，可以量化一个注意力头是多大程度上基于位置的。 $$\sigma_p(\mathbb{E}_t[A_{h,i,j}]) = \alpha_h \delta_p$$ 其中 $\delta_p$ 对所有头都相同。这个公式表明，位置影响因子 $\alpha_h$ 与注意力头 $h$ 的 SoE 成正比。这意味着 SoE 越高，该头对位置信息的依赖性越强。MoA 生成的掩码密度与 SoE 呈正相关（如原文 Figure 13 所示），这表明 MoA 成功地捕捉了密集语言模型的语义信息。

5. 实验设置

5.1. 数据集

MoA 实验中使用了多种数据集，主要分为校准数据集和评估数据集。

5.1.1. 校准数据集

• MultiNews (Fabbri et al., 2019): 这是一个大规模多文档摘要数据集。它包含多个新闻文档，要求生成一个长篇摘要，天然具有长距离依赖性。MoA 采用此数据集构建校准集，并使用原始密集模型生成的摘要作为监督文本，以确保模型对齐。
• RedPajama (Together Computer, 2023): 这是一个通用语言建模数据集，由人类编写的文本语料库组成。在论文的消融实验中作为对比，代表缺乏长距离依赖性和模型对齐的传统校准数据集。

5.1.2. 评估数据集

• LongEval (Li et al., 2023a): 用于测试长上下文检索能力。该基准测试将大量的键值对作为输入，并根据给定的键测试从不同位置检索值的准确性。包含 100 个数据项，用于评估不同注意力跨度和输入长度下的检索能力。
  • LongEval 输入示例 (Format 1):

    Below is a record of lines I want you to remember. Each line begins with 'line <line index>' and contains a ' <REGISTER_CONTENT>' at the end of the line as a numerical value. For each line index, memorize its corresponding <REGISTER_CONTENT>. At the end of the record, I will ask you to retrieve the corresponding <REGISTER_CONTENT> of a certain line index. Now the record start:
    
    line delightful-incandescence: REGISTER_CONTENT is <19147>
    line cloistered-presence: REGISTER_CONTENT is <8862>
    ... (大量键值对) ...
    
    Now the record is over. Tell me what is the <REGISTER_CONTENT> in line cloistered-presence? I need the number.
    

• Needle-In-A-Haystack (NIAH) (Kamradt, 2024): 另一个常见的检索基准测试。它包含一个“针 (needle)”句子，通常不适合无关上下文。模型需要根据这个“针”句来回答问题。
  • NIAH 输入示例 (Format 2):

    People who are powerful but uncharismatic will tend to be disliked. Their power makes them a target for criticism that they don't have the charisma to disarm. That was Hillary Clinton's problem.
    
    The best thing to do in San Francisco is eat a sandwich and sit in Dolores Park on a sunny day.
    
    It also tends to be a problem for any CEO who is more of a builder than a schmoozer.
    
    What is the best thing to do in San Francisco?
    

• LV-Eval (Yuan et al., 2024): 用于评估长上下文理解能力，包含 11 个子数据集。在实验中使用了 16k 分割，输入截断至 15500 词元。
• LongBench (Bai et al., 2023): 另一个用于评估长上下文理解能力的基准，包含 13 个子数据集，具有平衡的各长度级别数据项。使用了 LongBench-E 分割，输入截断阈值分别为 3,500（0-4k）、7,500（4-8k）和 15,500（8k+）词元。
• 困惑度 (Perplexity) 评估数据集：
  • Qasper (Dasigi et al., 2021): 问答数据集。
  • MultiNews (Fabbri et al., 2019): 摘要数据集。
  • TREC (Li & Roth, 2002; Hovy et al., 2001): 少样本学习数据集（通常用于分类任务）。
  • LCC (Mohler et al., 2016): 代码补全数据集。 这些数据集每个长度级别采样 $50 \times 4$ 个数据项的测试集，困惑度仅在回答部分计算。
• AlpacaEval 2.0 (Li et al., 2023b; Dubois et al., 2024): 用于评估模型在通用指令遵循任务上的性能，通过与 gpt4_turbo 进行比较来计算胜率 (win rate)。

校准数据集输入示例 (Format 3):

You are given several news passages. Write a one-page summary of all news. <News1
News2>

Now, write a one-page summary of all the news.

Summarization (原始密集模型生成的摘要作为监督)

5.2. 评估指标

论文中使用的评估指标及其说明如下：

1. 检索准确率 (Retrieval Accuracy)

   • 概念定义： 衡量模型在给定长上下文和特定查询时，从上下文中正确提取目标信息的比例。它直接反映了模型在长文本中定位和回忆关键事实的能力。
   • 数学公式： 通常表示为正确检索的数量与总查询数量之比。 $$\text{Retrieval Accuracy} = \frac{\text{Number of Correct Retrievals}}{\text{Total Number of Queries}}$$
   • 符号解释：
     • Number of Correct Retrievals：模型正确找到并返回所需信息的次数。
     • Total Number of Queries：模型尝试检索信息的总次数。

2. LV-Eval 分数 (LV-Eval Score)

   • 概念定义： LV-Eval 是一个长上下文理解基准测试，其分数反映了模型在多方面长文本理解任务（如问答、摘要等）上的综合表现。分数的具体计算可能根据子任务的不同而有所调整，但总体目标是量化模型对长文本复杂语义的把握能力。
   • 数学公式： 论文中未提供 LV-Eval 的具体计算公式，但通常此类基准分数是根据多个子任务的特定评估指标（如 F1-score、Rouge-score、Exact Match 等）加权或平均得出。
   • 符号解释： 无特定公式符号，分数越高表示性能越好。

3. LongBench 分数 (LongBench Score)

   • 概念定义： LongBench 是另一个综合性的长上下文理解基准，涵盖了多种任务和长度级别。其分数反映了模型在处理不同类型和长度的长文本任务时的综合能力。
   • 数学公式： 论文中未提供 LongBench 的具体计算公式，同 LV-Eval，通常是各项子任务评估指标的聚合。
   • 符号解释： 无特定公式符号，分数越高表示性能越好。

4. 困惑度 (Perplexity, PPL)

   • 概念定义： 困惑度是评估语言模型性能的常用指标，衡量模型预测序列中下一个词元的难度。困惑度越低，表示模型对训练数据的建模越好，对序列的预测越自信和准确。它本质上是模型对文本的平均不确定性。
   • 数学公式： $$\text{PPL}(W) = \left( \prod_{i=1}^N \frac{1}{P(w_i | w_{1 \dots i-1})} \right)^{\frac{1}{N}} = \exp\left(-\frac{1}{N} \sum_{i=1}^{N} \log P(w_i | w_{1 \dots i-1})\right)$$
   • 符号解释：
     • $W = (w_1, w_2, \dots, w_N)$：一个由 $N$ 个词元组成的序列。
     • $w_i$：序列中的第 $i$ 个词元。
     • $P(w_i | w_{1 \dots i-1})$：语言模型预测给定前 i-1 个词元时，第 $i$ 个词元出现的概率。
     • $N$：序列中的词元总数。

5. 解码吞吐量 (Decode Throughput)

   • 概念定义： 衡量模型在推理阶段每秒可以生成的词元数量。高吞吐量意味着更快的响应速度，对实际应用至关重要。
   • 数学公式： $$\text{Throughput} = \frac{\text{Number of Tokens Generated}}{\text{Total Decoding Time (seconds)}}$$
   • 符号解释： 无特定公式符号，单位通常为 tokens/s，值越高表示效率越好。

6. GPU 内存减少 (GPU Memory Reduction)

   • 概念定义： 衡量与基线方法相比，MoA 在 GPU 上节省的内存比例。内存效率是处理长上下文和部署大型模型时的关键考量。
   • 数学公式： 通常表示为倍数，例如“1.2-1.4 倍内存减少”意味着 MoA 使用的内存是基线的 $\frac{1}{1.2}$ 到 $\frac{1}{1.4}$。
   • 符号解释： 无特定公式符号，倍数越大表示内存节省越多。

7. 长度控制胜率 (Length-controlled Win Rate) (AlpacaEval 2.0)

   • 概念定义： 在 AlpacaEval 2.0 基准中，胜率衡量的是模型生成的响应在与另一个参考模型（如 GPT-4 Turbo）进行比较时，被裁判（通常是另一个 LLM）判断为“更好”的比例。长度控制 (Length-controlled) 指的是在评估过程中对生成文本的长度进行某种形式的归一化或考虑，以避免简单地偏向于生成更长的响应。
   • 数学公式： $$\text{Win Rate} = \frac{\text{Number of Times Model A's Response is Preferred}}{\text{Total Number of Comparisons}}$$
   • 符号解释： 无特定公式符号，百分比越高表示模型的指令遵循和生成质量越好。

8. 每词元能量消耗 (Energy Per Token)

   • 概念定义： 衡量生成一个词元所需的能量（焦耳），反映了模型的能源效率。在可持续 AI 和大规模部署中是一个重要指标。
   • 数学公式： $$\text{Energy Per Token} = \frac{\text{Total GPU Energy Consumption (Joules)}}{\text{Total Number of Tokens Generated}}$$
   • 符号解释： 无特定公式符号，单位通常为 Joule/token，值越低表示能源效率越高。

5.3. 对比基线

论文将 MoA 与一系列最先进的静态和动态高效注意力方法、以及 LLM 加速框架进行了比较：

• StreamingLLM (Xiao et al., 2024c): 一种静态高效注意力方法，采用统一的滑动窗口掩码和 attention sink，旨在通过限制 KV-Cache 大小来处理长上下文。它使用稀疏预填充。
• InfLLM (Xiao et al., 2024a): 一种训练无关的记忆增强方法，旨在通过训练无关的记忆来处理极长序列。它使用稀疏预填充，但动态决定 KV-Cache。
• H2O (Zhang et al., 2023): Heavy-Hitter Oracle，一种动态高效注意力方法，基于列和的注意力矩阵计算来动态决定 KV-Cache 的驱逐。它使用原始的密集预填充，并在解码阶段进行额外计算。
• BigBird (Zaheer et al., 2020): 静态稀疏注意力方法，为 Transformer 引入了块稀疏注意力机制，以处理更长序列。
• SnapKV (Li et al., 2024c): 一种近期提出的高效 KV-Cache 压缩方法，旨在在生成前预测和保留关键 KV 条目。
• PyramidKV (Cai et al., 2024): 另一种近期提出的动态 KV-Cache 压缩方法，基于金字塔式信息漏斗。
• FlashAttention2 (Dao, 2024): 一种高度优化的 CUDA 内核实现，用于加速密集注意力计算，显著提高内存效率和吞吐量。
• vLLM (Kwon et al., 2023): 一个高效的 LLM 服务框架，其核心是 PagedAttention，通过分页管理 KV-Cache，实现了高吞吐量。

预填充方式差异：

• MoA 和 StreamingLLM 使用高效稀疏预填充。
• H2O 使用原始密集预填充，因为它依赖于注意力矩阵的列和来计算词元的 KV-Cache 驱逐重要性。
• InfLLM 使用稀疏预填充，但需要额外的计算来动态确定 KV-Cache。

5.4. MoA 设置

• 注意力汇聚点： MoA 使用块滑动窗口注意力模式，块大小为 64。第一个词元块不被掩码，作为 attention sink。
• 剖析阶段：
  • 使用 MultiNews 数据集作为校准数据集，并以模型响应作为监督文本（如第 5 节所述）。
  • 在 2k、4k 和 8k 长度下使用 $50 \times 3$ 个数据项进行剖析。数据项被填充到相应的长度级别，以确保注意力影响张量的形状统一。
  • 在剖析过程中，采用块粒度，计算每个块内的平均注意力影响。
• 超参数搜索空间：
  • $\alpha$ 默认有 6 个选项，均匀采样自 $[-2048, 8192]$。
  • $\beta$ 默认有 9 个选项，均匀采样自 [0, 1]。
  • 这为每个注意力头创建了 $6 \times 9 = 54$ 对 $(\alpha, \beta)$ 的搜索空间。
  • 最终的注意力跨度长度被限制在 0 和当前输入长度之间。
• 优化阶段：
  • 使用多目标优化在同一组长度上进行。
  • 限制每个模型层最多使用两种不同的规则，以确保推理效率。
  • 在帕累托前沿解决方案中，选择在 12k 长度的验证数据集上具有最低困惑度的配置作为最优解决方案。
• 模型未微调： 所有实验均在未进行任何微调的模型上进行。

6. 实验结果与分析

6.1. 核心结果分析

6.1.1. 准确性-吞吐量权衡

原文 Figure 4 展示了 MoA 在检索准确率和解码吞吐量方面相对于六个基线方法的显著优势，将帕累托前沿向前推进。在相同密度下，MoA 的吞吐量比 H2O、InfLLM、BigBird、SnapKV 和 PyramidKV 提高了 1.6-18.1 倍。这得益于其高效的静态注意力设计和定制的 GPU 内核。

该图像是图表，展示了在不同密度下七种注意力方法在Vicuna-7B模型中进行的解码吞吐量与检索准确率之间的权衡。方法包括InfLLM、BigBird、H2O、SnapKV、PyramidKV、StreamingLLM和MoA，X轴表示解码吞吐量（token/s），Y轴表示检索准确率。

6.1.2. 性能评估 (50% 密度)

原文 Table 4 比较了不同模型在 50% 密度下各种注意力方法的检索准确率、LV-Eval 分数、LongBench 分数和困惑度。

关键发现：

• 长上下文检索： MoA 在检索准确率方面表现卓越，最大相对准确率下降仅为 8%，远低于 StreamingLLM (87%)、InfLLM (58%) 和 H2O (44%)。平均而言，MoA 的相对准确率下降低于 1%。特别是在 Llama3-8B 模型上，MoA 在 4k、8k、16k 长度下均达到了与 Original 模型几乎相同的 0.99-1.00 的检索准确率。
• 长上下文理解： 在 LV-Eval 和 LongBench 基准上，MoA 将最大相对性能下降分别降低到 5% 和 3%，显著优于 StreamingLLM (36% 和 18%)。虽然 H2O 和 InfLLM 也有较小的性能下降，但它们通常伴随着更高的效率成本（如密集预填充或额外的动态计算）。
• 困惑度： MoA 的相对困惑度增加不到 1%，而其他方法则增加了 4%-13%。这表明 MoA 在保持模型生成连贯性方面表现出色。
• 大规模模型表现： Llama3-70B 模型上，H2O 在 16k 长度时出现内存溢出 (OOM)，而 MoA 依然能保持 0.94 的检索准确率和与 Original 模型相近的 LV-Eval 和 LongBench 分数，以及较低的困惑度。

6.1.3. 有效上下文长度与检索准确率

原文 Figure 1 展示了 Vicuna-7B 模型在 LongEval 基准上不同注意力方法的检索准确率。

该图像是表述不同的注意力方法在 LongEval 基准测试中的检索准确性图，包含三个部分：(a) 原始模型，注意力范围为输入长度的 100%；(b) StreamingLLM，注意力范围为输入长度的 50%，超出范围检索效果降低；(c) MoA，注意力范围同样为输入长度的 50%，在范围外保持有效检索。

• 全注意力模型 (a)： 在所有输入长度和检索位置上均保持高准确率。

• StreamingLLM (b)： 当注意力跨度为输入长度的 50% 时，在超出其注意力跨度后（即更早的输入部分）检索准确率急剧下降，甚至在更长输入长度时表现更差。

• MoA (c)： 同样在 50% 的注意力跨度下，MoA 能够有效地保持超出其物理跨度的检索准确率，表明其能够捕捉到更长距离的上下文信息。

  原文 Figure 7(a) 进一步量化了 MoA 在 LongEval 上的有效上下文长度。

  该图像是图表，展示了在LongEval上使用Vicuna-7B进行的检索准确性测试。图(a)显示了不同注意力跨度下的有效长度，图(b)展示了在特定密度下的检索准确性变化。

• StreamingLLM 和 H2O 的有效上下文长度（检索准确率保持在 90% 以上的最大输入长度）超出其注意力跨度不超过 2k 词元。

• 相比之下，MoA 将其有效上下文长度扩展到注意力跨度的大约 3.9 倍，在 Vicuna-7B 模型上最高可达 12k 词元，与原始密集模型相当。 原文 Figure 7(b) 显示，在 8k 固定输入长度下，MoA 仅需 25% 的密度即可达到 0.9 以上的检索准确率，而 StreamingLLM 和 H2O 分别需要 100% 和 75% 的密度才能达到类似性能。

原文 Figure 8 验证了 MoA 在 Needle-In-A-Haystack (NIAH) 检索任务上的表现。

该图像是一个图表，展示了不同注意力方法在 Llama-3-8B 模型中，针对输入长度从 8k 到 256k 的 Needle-In-A-Haystack (NIAH) 检索准确率。图中显示了 MoA、H2O、StreamingLLM 和 InfLLM 方法在不同 Token 限制下的深度表现，颜色深浅表示不同的准确性评分。

• MoA 在 Llama-3-8B 模型上，输入长度从 8k 到 256k 均实现了完美检索准确率。
• StreamingLLM 的有效上下文长度有限。
• InfLLM 在 64k 输入长度内检索准确率下降。
• H2O 和 InfLLM 在极端长度下出现 OOM 和 OOT 错误。

6.1.4. 长上下文理解与泛化能力

原文 Figure 5(b) 和 Table 9 展示了 MoA 在更长上下文长度（32k-256k）上的泛化能力。

该图像是图表，展示了Vicuna-7B模型在不同注意力跨度和输入长度下的检索准确率。X轴表示不同的注意力跨度，Y轴表示不同的输入长度，图中的热图显示了不同配置下的检索准确性。

• 通过在 12k 长度内进行压缩，MoA 成功泛化到 32k-256k 的长度。
• 在这些扩展长度下，MoA 在检索准确率方面比 InfLLM 和 StreamingLLM 提高了 1.9-3.3 倍，在 LV-Eval 分数方面提高了 1.2-1.4 倍，并展现出与原始密集模型相当的性能。
• 与 SnapKV 和 PyramidKV 等最新基线相比，MoA 在 32k 到 256k 的长上下文上持续表现出优越的检索准确率。

6.1.5. 指令遵循生成

原文 Table 10 在 AlpacaEval 2.0 基准上评估了 MoA 在通用指令遵循任务上的性能。

• MoA 实现了最高的长度控制胜率 (9.83%)，优于所有高效注意力基线和原始模型。
• 这表明 MoA 的弹性设计在短上下文（平均 50 词元输入，450 词元输出）的通用指令遵循任务中也表现良好。

6.1.6. 性能评估 (75% 密度)

原文 Table 6 展示了在 75% 密度下 MoA 与 StreamingLLM 的性能对比。

• 在 75% 密度下，MoA 依然持续优于 StreamingLLM，进一步验证了其异构弹性规则的有效性。

6.1.7. 长上下文理解子任务表现

原文 Figure 10 和 Table 8 详细展示了 MoA 在 LongBench 和 LV-Eval 子任务上的表现。

该图像是图表，展示了 MoA、H2O 和 StreamingLLM 在不同模型（7B、13B、70B）上的性能对比。图中展示了各种注意力方法在总结、代码完成和多文档问答任务中的得分，指标均经过归一化处理。

• MoA 实现了与原始密集模型以及 H2O（尽管效率成本更高）相当的全面性能。
• StreamingLLM 和 InfLLM 表现出不一致的性能：在某些任务中可能超过原始模型，但在其他任务中则遭受显著性能下降。这凸显了 MoA 在不同子任务上的稳健性和一致性。

6.2. 消融实验/参数分析

6.2.1. 校准数据集设计消融研究

原文 Table 2 验证了校准数据集设计选择对 MoA 性能的重要性。

• 长距离依赖性： 使用 MultiNews 摘要进行损失计算，相比于 RedPajama 或 MultiNews 仅用上下文作为监督，检索准确率显著提高 60%，困惑度降低 0.98。这强调了长距离依赖性数据的重要性。
• 模型对齐： 使用原始密集模型生成的摘要作为监督，相比于人类编写的摘要，性能进一步提升（检索准确率从 0.87 提高到 0.95）。这证明了模型响应与监督文本对齐的重要性。

6.2.2. 搜索空间异构性消融研究

原文 Table 3 评估了不同滑动窗口注意力搜索空间对性能的影响，通过逐步引入异构性进行。

• 从基本的统一滑动窗口方法开始，逐步引入层级异构性、注意力头异构性，最后是弹性规则，性能都得到了显著提升。
• 特别是引入注意力头异构性后，检索准确率从 0.31 跃升至 0.95，困惑度也大幅下降。
• 进一步引入弹性规则，使检索准确率达到 0.98，困惑度也略有改善。这证实了 MoA 中异构性（在层和头之间）和弹性规则的必要性。

6.3. 效率评估

6.3.1. 内存和吞吐量分解

原文 Table 5 比较了 MoA 在不同 LLM 框架和注意力方法下的解码吞吐量，并对 MoA 的效率提升进行了因素分解。

吞吐量提升来源：

• 静态 KV-Cache 尺寸： 仅保留每个头跨度内的词元，防止 KV-Cache 增长，从而避免额外的内存分配（约 3.0 倍提升）。
• 减少注意力计算： 注意力跨度内的密度降低，减少了注意力计算所需的计算量和内存访问（约 1.5 倍提升）。
• 增加批处理大小： KV-Cache 内存使用减少，支持更大的批处理大小（约 1.4 倍提升）。
• 定制 GPU 内核： 为 MoA 的异构注意力模式设计定制 CUDA GPU 内核（约 1.2 倍提升）。

与基线的对比：

• MoA 在 50% 密度下，相比 FlashAttention2 实现了 6.6-8.2 倍的解码吞吐量提升。

• 相比 H2O 和 InfLLM 提高了 1.2-4.0 倍。

• 相比高度优化的 vLLM 框架，MoA 仍能实现 1.7-1.9 倍的吞吐量增益。

  原文 Table 11 详细比较了内存效率。

  Size	Framework	Memory (GB) 4k 8k 16k
  7B	FlashAttn2 H20 MoA	28.5 36.9 22.7	44.4 OOM 32.9	76.3 OOM 53.5
  13B	FlashAttn2 H20 MoA	36.8 40.4 32.0	49.2 77.9 39.6	74.0 OOM 55.0

• H2O 在较长输入长度下（7B 模型 8k，13B 模型 16k）出现 OOM。

• MoA 相较于 FlashAttention2 减少了 1.2-1.4 倍的内存消耗。

6.3.2. 更长输入长度下的效率

原文 Table 12 评估了 MoA 在 128k 输入长度下的运行时效率。

• MoA 仅使用一个 A100 GPU 即可高效处理 128k 输入，而 FlashAttention2 和 vLLM 基线需要至少两个 GPU。
• MoA 相比 FlashAttention2 实现了 4.7-4.9 倍的解码速度提升，同时 GPU 数量减半。
• 相比 vLLM，MoA 在每个 GPU 上提供了 1.3-1.4 倍的吞吐量，并显著节省了内存。

6.3.3. 能量效率

原文 Table 13 显示了 MoA 在不同序列长度下的每词元能量消耗。

• MoA 实现了 8.7-10 倍的每输出词元能量消耗降低，主要得益于略低的 GPU 功耗。

6.3.4. 自动搜索流水线开销

原文 Table 14 详细列出了 MoA 流水线各阶段的时间开销。

• MoA 能够在数小时内完成自动配置搜索（例如 Vicuna-13B 仅需 1.75 小时）。
• 剖析阶段是资源需求最大的部分，对于 13B 模型，需要两块 A100 GPU 约 25 分钟。
• 优化阶段主要在 CPU 上运行，耗时相对较少（约 25 分钟）。

6.4. MoA 配置分析

6.4.1. 掩码密度分布

原文 Figure 11 展示了 MoA 掩码密度在不同 LLM 的层级上的分布。

该图像是图表，展示了不同层级上 Vicuna-7B 和 Vicuna-13B 模型的 MoA 掩码密度。图中以不同颜色表示窗口长度（2k、4k、8k、16k），并显示了各层的密度变化趋势。

• 初始层和中间层的掩码密度较高，表明这些层通常需要较大的注意力跨度。这与关于 LLM 固有维度和层敏感性的研究结论一致。

• 在最后的层，大多数注意力头需要较低的密度，但少数异常头仍然需要高密度。 原文 Figure 12 进一步展示了平均密度和密度范围。

  该图像是散点图，展示了Vicuna-7B和Vicuna-13B模型中平均密度与密度范围之间的关系。每个点代表不同的配置，图中还包括线性回归趋势线及阴影区域，反映了两者的密度变化趋势。

• 平均密度较低的层（图中的线条值较小）通常在不同注意力头之间显示出更宽的密度范围（图中的阴影区域更宽）。这证实了在同一层内也需要异构注意力规则。

6.4.2. 规则与语义的关联

原文 Figure 13 展示了 MoA 掩码稀疏性与注意力头对位置依赖性 (SoE) 之间的正相关关系。

该图像是图表，展示了MoA的掩码稀疏性与头部对位置依赖性的正相关关系(SoE)。图表中，水平坐标轴表示SoE(E-2)，竖直坐标轴表示稀疏性百分比，伴随有两侧的边际分布图，显示不同SoE值的稀疏性变化情况。

• 这表明 MoA 成功捕获了密集语言模型的语义信息，将更高的稀疏性（更小的注意力跨度）分配给那些更多依赖于词元内容而非位置的注意力头，而将更低的稀疏性（更大的注意力跨度）分配给那些更依赖于位置信息的头。

  原文 Figure 14 显示了相同的注意力头在不同任务中（少样本学习、聊天、编码）表现出相似的注意力跨度。

  该图像是一个示意图，展示了Vicuna-7B模型中不同任务（少量学习、聊天和编码）下第17层第29个注意力头的注意矩阵。每个注意力矩阵通过256个数据项的平均值展示，反映了不同任务中的注意力模式。相同的注意力头在不同任务中表现出相似的注意力跨度，说明我们方法的跨数据集泛化能力。

• 这解释了 MoA 方法强大的跨数据集泛化能力，因为注意力头的内在模式在不同任务中相对稳定。

7. 总结与思考

7.1. 结论总结

论文提出了一种名为 混合注意力跨度 (Mixture of Attention Spans, MoA) 的创新方法，旨在通过为大型语言模型 (LLM) 的每个注意力头和层自动定制异构且弹性的滑动窗口长度，从而显著提高长上下文推理的效率。

MoA 的核心在于认识到 LLM 内部注意力模式的异构性，并设计了一个包含线性缩放规则的搜索空间，使得注意力跨度能够根据输入长度自适应调整。通过一个自动化的剖析-优化流水线，MoA 能够高效地找到最优的注意力配置，该流水线特别强调了使用具有长距离依赖性且与模型响应对齐的校准数据集。

实验结果强有力地证明了 MoA 的有效性：

• 性能方面： 在相同的平均滑动窗口长度下，MoA 将有效上下文长度提高了 3.9 倍，在多项长上下文检索和理解基准测试中，检索准确率相较于统一窗口基线提高了 1.5-7.1 倍。它将与全注意力模型的最大相对性能下降从 9%-36% 降低到 5% 以内，平均下降低于 1%。

• 效率方面： MoA 实现了 1.2-1.4 倍的 GPU 内存减少，并将解码吞吐量相较于 FlashAttention2 和 vLLM 分别提高了 6.6-8.2 倍和 1.7-1.9 倍。在 128k 的超长输入下，MoA 仅需一个 GPU 即可运行，并提供更高的每 GPU 吞吐量和更低的内存使用。此外，MoA 还将每词元能量消耗降低了 8.7-10 倍。

• 泛化能力： MoA 在 12k 长度内进行配置搜索，能够有效泛化到 32k-256k 的更长上下文长度。其发现的注意力头模式在不同任务中表现出一致性，解释了其强大的跨数据集泛化能力。

  总而言之，MoA 提供了一个训练无关、高效且高性能的解决方案，解决了 LLM 在长上下文推理中面临的效率与准确性权衡问题，为 LLM 的实际部署和应用带来了显著的效益。

7.2. 局限性与未来工作

论文作者指出了 MoA 的以下局限性：

• 极低密度预算下的性能： 在极低的密度预算下，MoA 未能保持良好的性能。

• 非线性弹性规则： MoA 目前使用线性弹性规则来定义注意力跨度，但可能存在更优的非线性规则。

  基于这些局限性，作者提出了以下未来研究方向：

• 动态 MoA 方法： 针对极低密度预算下的性能问题，设计一种动态 MoA 方法可能是一个有前景的方向。

• 非线性弹性规则： 探索使用具有有界注意力跨度的非线性弹性规则，可能会带来更好的性能。

• 推广到其他压缩方法： MoA 的剖析方法可以适应评估权重和其他激活的影响，从而促进其他压缩方法（如量化）的发展。

7.3. 个人启发与批判

7.3.1. 个人启发

1. 异构性是关键： 这篇论文最核心的启发在于，LLM 内部的注意力机制并非“同质化”的，不同注意力头和层有其独特的关注模式和对上下文长度的需求。简单粗暴地采用统一策略必然会造成资源浪费或性能瓶颈。未来在设计各种 LLM 优化（如稀疏化、剪枝、量化）时，应更深入地考虑这种内在的异构性，而不是采取“一刀切”的方案。
2. 剖析驱动的优化： MoA 通过梯度剖析来量化每个注意力值对损失的影响，这提供了一个数据驱动、模型感知的优化方向。这种“评估-生成-优化”的流水线思想具有很强的通用性，可以迁移到其他模型压缩和结构优化的任务中。例如，在权重剪枝中，也可以通过类似的影响剖析来确定哪些权重或神经元可以安全地移除。
3. 校准数据集的重要性： 论文强调了校准数据集的设计对压缩效果至关重要，特别是长距离依赖性和模型对齐。这提醒我们，在进行任何模型优化时，用于评估和指导优化的数据必须与目标任务和模型特性高度匹配，否则可能导致优化结果与实际需求脱节。
4. 工程与算法的结合： MoA 的成功不仅依赖于巧妙的算法设计（异构弹性规则、多目标优化），也离不开高效的工程实现（定制 CUDA 内核、内存管理优化）。这再次证明了在 LLM 时代，算法创新和系统级优化是并驾齐驱的。
5. 语义可解释性： 论文中对 MoA 规则与注意力头语义（如对位置依赖性）的关联分析，为稀疏化提供了额外的可解释性。了解哪些头专注于局部，哪些专注于全局，以及它们对位置或内容信息的偏好，有助于我们更好地理解 Transformer 的工作原理，并指导更智能的稀疏化策略。

7.3.2. 批判与潜在改进

1. 线性弹性规则的局限性： MoA 采用了简单的线性函数 $S_h = \alpha_h + \beta_h \cdot N$ 来定义注意力跨度。虽然这在一定程度上捕捉了弹性行为，但实际的注意力需求可能更为复杂，是非线性的。例如，在非常短或非常长的输入下，注意力跨度的增长可能呈现饱和或更复杂的非单调趋势。引入更灵活的（例如分段线性、对数或指数）或数据驱动的非线性弹性规则，可能进一步提升性能。
2. 计算开销的衡量： 尽管 MoA 在推理阶段效率显著，但其自动搜索流水线的开销（对于 70B 模型达到 34.7 小时 GPU 时间）对于一些资源受限的团队来说仍然较高。虽然是一次性开销，但如何进一步减少这个搜索时间是一个值得探索的方向。例如，可以探索更高效的梯度近似方法、更快的优化器，或者利用元学习 (meta-learning) 从现有模型配置中学习搜索策略。
3. 稀疏度的细粒度控制： MoA 主要关注滑动窗口长度的异构性。然而，稀疏注意力还有其他形式，例如基于内容的稀疏化（如 Routing Transformers）、或更复杂的局部-全局混合模式。MoA 可以在此基础上，探索更细粒度的稀疏模式，例如在滑动窗口内部，进一步对不重要的 KV 对进行动态剪枝。这可能有助于解决在极低密度预算下的性能下降问题。
4. 动态 MoA 的挑战： 作者提到未来工作将探索动态 MoA。这意味着在推理时根据输入内容动态调整注意力跨度，而非仅根据输入长度。这将带来更大的挑战，例如如何高效地计算动态重要性分数、如何管理动态变化的 KV-Cache 结构以及如何在不引入额外延迟的情况下实现这些动态调整。这需要更深层次的软硬件协同设计。
5. 普适性验证： 虽然论文在多个 Llama 系列模型上进行了验证，但 MoA 的配置搜索结果和其内在的语义关联是否能泛化到更广泛的 LLM 架构（例如非 Transformer 架构、不同的 tokenizer 等）仍需进一步验证。此外，在多模态 LLM 中，注意力机制的异构性可能更加复杂，MoA 如何适应这些场景也是一个有趣的研究方向。