注意力展开（Attention Rollout）是一种用于解释基于Transformer (architecture)的模型（主要是大型语言模型）的方法，通过计算从每个输入令牌到网络中所有其他令牌在每一层中的累积注意力来实现。作为对单层注意力权重通常带有噪声且不能直接反映模型最终决策这一观察的回应，注意力展开跨层聚合注意力矩阵，以生成一个连贯的重要性图。该技术广泛应用于可解释人工智能领域，用于可视化输入的哪些部分对模型预测贡献最大，有助于调试、偏差检测和理解模型行为。
其核心思想建立在这样一个概念上：在transformer中，每一层的注意力机制计算来自前一层表示的加权和。通过将连续层的注意力矩阵相乘，可以追踪信息如何从早期令牌通过网络传播到最终输出。这类似于在有向图中沿路径追踪，其中节点是令牌，边是注意力权重。得到的展开矩阵提供了令牌影响的全局视图，这通常比查看任何单层的注意力图更具可解释性。
背景与动机
Transformer由Google DeepMind和University of Toronto的研究人员在2017年论文《Attention Is All You Need》中引入，依赖自注意力机制来捕捉令牌之间的依赖关系。每一层计算注意力分数，指示每个令牌对其他令牌的关注程度。然而，这些分数不能直接解释为重要性。例如，一个令牌可能关注一个语义权重很小的标点符号，或者注意力可能分散在许多令牌上。此外，后层的注意力是基于已经由前层转换的表示计算的，因此原始权重并未考虑这种复合效应。
早期的可解释性尝试集中于可视化单个头或层的注意力图，但这些往往产生相互矛盾或令人困惑的结果。像Anima Anandkumar和Jakob Uszkoreit这样的研究人员指出，注意力模式在层和头之间可能高度可变，使得难以得出结论。这促使了聚合整个网络信息的方法的发展。
算法
注意力展开在假设注意力权重代表一种信息路由形式的前提下运行。算法步骤如下：
- 对于每一层 \( l \)，获取形状为（序列长度，序列长度）的注意力矩阵 \( A_l \)，其中 \( A_l[i][j] \) 是从令牌 \( i \) 到令牌 \( j \) 的注意力权重。该矩阵通常在该层所有注意力头上取平均。
- 将单位矩阵添加到每个注意力矩阵中，以考虑残差连接，这允许每个令牌保留自身信息。得到的矩阵为 \( \tilde{A}_l = A_l + I \)。
- 对 \( \tilde{A}_l \) 的每一行进行归一化，使其和为1，确保矩阵保持随机性。
- 通过将所有层的归一化矩阵相乘来计算展开矩阵 \( R \)：\( R = \tilde{A}_1 \cdot \tilde{A}_2 \cdot ... \cdot \tilde{A}_L \)，其中 \( L \) 是层数。
得到的矩阵 \( R \) 给出了从每个令牌到其他每个令牌的总注意力流，考虑了通过网络的所有路径。条目 \( R[i][j] \) 可以解释为来自令牌 \( j \) 的信息影响输出处令牌 \( i \) 表示的比例。
该方法首次在2020年由OpenAI和Stanford AI Lab的研究人员提出，他们展示了其在GPT-2和其他模型上的有效性。他们表明，注意力展开可以识别情感分析和代词消解等任务中的显著令牌，通常在与人判断的一致性方面优于单层注意力。
在模型解释中的应用
注意力展开已应用于自然语言处理及其他领域的各种任务。在Machine learning研究中，它用于：
- 识别关键令牌：通过检查展开矩阵，从业者可以看到哪些输入词或子词对模型输出影响最大。例如，在情感分类任务中，展开可能会突出“terrible”或“amazing”等词为高度影响。
- 检测偏差：如果模型持续以导致有偏预测的方式关注人口统计术语，展开可以揭示这些模式，帮助研究人员缓解不公平行为。
- 调试模型错误：当模型做出错误预测时，展开可以显示它是否关注了输入的不相关部分，指导训练数据或架构的改进。
- 比较模型：通过计算不同模型的展开矩阵，可以比较其内部信息流，这有助于模型选择和理解架构差异。
在文本之外，注意力展开已被改编用于计算机视觉中的视觉transformer（ViT）。例如，Google Cloud和Amazon Web Services的研究人员已将其应用于图像分类，以可视化图像的哪些区域对预测有贡献，辅助医学影像分析和自动驾驶系统。
局限性与批评
尽管受欢迎，注意力展开有几个局限性。首先，将注意力权重视为信息流的概率分布的假设并不总是有效。注意力权重是基于学习到的表示计算的，它们不一定反映因果影响。像Melanie Mitchell和Aleksander Madry这样的批评者认为，注意力不是解释，像展开这样的方法可能产生误导性解释。
其次，跨层乘法注意力矩阵可能导致数值问题，尤其是在深层网络中。展开矩阵可能变得过于分散，所有令牌具有相似重要性，或者由于重复乘法而集中在少数令牌上。这可能在实践中使结果不太有用。
第三，注意力展开不考虑注意力层之间发生的非线性变换（前馈网络、层归一化）。这些变换可能显著改变信息流，因此简单的乘法模型可能遗漏重要效应。
第四，该方法在头上平均注意力，这可能掩盖不同头的不同角色。有些头可能捕捉句法关系，而其他头捕捉语义关系；平均它们会失去这种细微差别。
变体与扩展
为了解决其中一些局限性，研究人员提出了几种变体：
- 注意力流：该方法由MIT CSAIL的研究人员引入，将注意力视为流问题，并使用最大流算法计算令牌重要性，避免乘法假设。
- 带残差权重的展开：一些实现通过反映残差连接强度的因子加权单位矩阵，而不是均匀添加。
- 逐层相关性传播（LRP）：这种来自Deep learning社区的技术通过网络向后传播相关性分数，提供基于注意力方法的替代方案。
- 集成梯度（Integrated Gradients）和SHAP：这些是不依赖注意力权重的特征归因方法，而是基于梯度信息或博弈论概念。它们通常用作对基于注意力方法的健全性检查。
尽管有这些替代方案，注意力展开仍然是一个流行的基线，因为其简单性和计算效率。它不需要额外训练或梯度计算，使其易于应用于任何预训练的transformer。
实现考虑
使用现代深度学习框架实现注意力展开是直接的。在PyTorch或TensorFlow中，可以挂钩前向传播以捕获注意力矩阵。关键步骤是：
- 在每个注意力层注册前向钩子以存储注意力权重。
- 在前向传播后，在头上平均注意力权重。
- 添加单位矩阵并归一化行。
- 顺序相乘矩阵。
一个实际考虑是，对于长序列，注意力矩阵可能很大，导致内存开销。对于长度为1024的序列，每个矩阵是1024x1024，相乘许多这样的矩阵可能计算密集。然而，对于典型输入，这是可管理的。
另一个考虑是，该方法假设标准transformer架构。对于具有Cross-Attention的模型（如编码器-解码器模型），展开必须适应处理编码器和解码器注意力之间的交互。在这种情况下，可以分别计算编码器和解码器的展开，或以更复杂的方式组合它们。
与其他可解释性方法的关系
注意力展开位于Artificial intelligence系统可解释性技术的更广泛格局中。它通常与以下方法比较：
- 基于梯度的方法：这些计算输出相对于输入嵌入的梯度，提供敏感性度量。它们在理论上更扎实，但需要反向传播。
- 基于扰动的方法：这些涉及修改输入令牌（例如，移除或掩码）并观察输出变化。它们是模型无关的，但计算昂贵。
- 代理模型：像LIME或SHAP这样的技术局部训练可解释模型以近似黑盒模型。它们对表格数据有用，但对文本不太适用。
注意力展开的独特之处在于它纯粹基于模型的内部注意力权重，除了前向传播外不需要额外计算。这使其特别适合实时应用，如交互式调试工具。
未来方向
随着transformer模型在规模和复杂性上继续增长，对稳健可解释性方法的需求增加。注意力展开可能在几个方面演变：
- 与因果方法集成：将展开与因果推断技术结合可以提供更准确的归因。
- 处理多模态模型：将展开扩展到同时处理文本、图像和音频的模型，如OpenAI和Anthropic开发的模型。
- 自动化分析：使用展开输出自动生成模型决策的自然语言解释，这可能对合规和审计有用。
- 基准测试：开发标准化基准来评估可解释性方法的忠实性，如Carnegie Mellon University和BAIR (Berkeley AI Research)的研究人员所提议的。
尽管有局限性，注意力展开已成为可解释性工具包中的基础工具。其简单性和有效性使其成为许多研究论文和实际应用中的标准基线。随着领域的发展，它可能会被精炼并与其他技术结合，以提供对神经网络内部运作的更深入见解。
结论
注意力展开提供了一种实用且直观的方式来理解transformer如何处理信息。通过跨层传播注意力权重，它提供了令牌影响的全局视角，这通常比检查单个层更有用。虽然并非没有缺陷，但其易用性和可解释性已巩固了其在可解释性文献中的地位。对于任何使用transformer模型的人来说，理解注意力展开是揭开这些强大但不透明系统神秘面纱的宝贵一步。