Attention Rollout 机制解析：如何提升人工智能模型的解释性

1次阅读

没有评论

共计 1863 个字符，预计需要花费 5 分钟才能阅读完成。

近年来，深度学习模型在各类任务上表现优异，但这些模型往往被视为“黑箱”。在实际应用中，尤其是在医疗、金融等关键领域，仅仅知道模型预测结果是不够的。我们需要理解模型为什么做出这样的预测，以及这些预测是否可靠。

现有方法如 LIME、SHAP 等虽然能提供局部解释，但对于基于注意力机制的模型（如 Transformer）解释力有限。
传统的注意力可视化通常只展示单层注意力权重，忽略了跨层的信息传播路径。
缺乏系统性的方法来追踪输入特征如何通过多层注意力机制影响最终预测。

Attention Rollout 的核心思想是通过矩阵乘法来聚合各层的注意力权重，从而追踪从输入到输出的完整注意力传播路径。

基础数学表达：假设我们有 L 层的 Transformer 模型，每层的注意力权重矩阵为 A_l ∈ R^{n×n}（n 是序列长度）
跨层传播计算：从最后一层开始，递归地计算前向注意力权重：R_l = A_l + A_l @ R_{l+1}
归一化处理：对最终的 rollout 矩阵进行行归一化，得到每个输入 token 对输出预测的影响程度

这种方法的优势在于：

保持了注意力权重的相对比例
能够捕捉长距离的依赖关系
计算过程可微，可以方便地与其他解释性方法结合

import torch
from torch import nn

def compute_attention_rollout(model, input_tokens, layer_start=0):
    """
    计算 Attention Rollout 权重
    Args:
        model: Transformer 模型
        input_tokens: 输入 token 的 embedding
        layer_start: 从哪一层开始计算
    Returns:
        rollout: 最终的 attention rollout 矩阵
    """
    with torch.no_grad():
        # 初始化各层注意力权重存储
        attentions = []

        # 前向传播收集各层注意力
        outputs = model(input_tokens, output_attentions=True)
        attentions = outputs.attentions[layer_start:]

        # 初始化 rollout 矩阵
        rollout = torch.eye(attentions[0].size(-1))

        # 反向计算 rollout
        for attention in reversed(attentions):
            # 平均多头注意力
            attention_heads = attention.mean(dim=1)
            # 添加残差连接
            rollout = attention_heads + attention_heads @ rollout

        # 行归一化
        rollout = rollout / rollout.sum(dim=-1, keepdim=True)

    return rollout

关键实现要点：