Attention Rollout 机制解析：如何优化Transformer模型的可解释性与计算效率

1次阅读

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

Transformer 模型在 NLP 和 CV 领域取得了巨大成功，但其内部工作机制往往被视为黑箱。这种可解释性不足的问题，限制了模型在医疗、金融等高风险领域的应用。同时，随着模型规模的扩大，计算资源消耗呈指数级增长，导致训练和推理成本居高不下。

可解释性差 ：传统 Attention 权重仅反映单层关系，难以追踪跨层的全局注意力路径。
计算效率低 ：多头注意力机制的计算复杂度随序列长度平方增长，在长文本处理时尤为明显。

当前主流的可解释性方法各有优劣：

Grad-CAM：依赖梯度信息，适合 CNN 但难以直接应用于 Transformer
LIME：通过局部逼近解释预测，但可能产生不一致的结果
Attention Flow：需要构建额外图结构，实现复杂度高

Attention Rollout 的独特优势在于：

仅需原始 Attention 权重即可计算
保持模型原生的计算图结构
可直观展示跨层注意力传播路径

Attention Rollout 的核心思想是通过矩阵连乘聚合各层注意力信息：

单层注意力计算 ：
$$A_i = softmax(\frac{QK^T}{\sqrt{d_k}})$$
跨层聚合 ：
$$Rollout = \prod_{i=1}^n (0.5I + 0.5A_i)$$

其中 I 是单位矩阵，0.5 系数用于平滑处理

关键实现步骤：

提取各层原始 Attention 权重
应用残差连接平滑（公式中的 0.5 系数）
按层序进行矩阵乘法
归一化最终注意力分布

import torch
import numpy as np

def attention_rollout(model, input_ids, attention_threshold=0.5):
    """
    Compute attention rollout for given input
    Args:
        model: Pretrained Transformer model
        input_ids: Tokenized input sequence
        attention_threshold: Smoothing factor (0-1)
    Returns:
        rollout: Aggregated attention matrix (L,L)
    """
    with torch.no_grad():
        outputs = model(input_ids, output_attentions=True)

    # Stack all layer attentions (n_layers, n_heads, L, L)
    attentions = torch.stack(outputs.attentions)

    # Average over heads
    avg_attentions = attentions.mean(dim=2)

    # Initialize rollout matrix
    rollout = torch.eye(attentions.shape[-1])

    # Apply recursive multiplication
    for attn in avg_attentions:
        attn = attention_threshold * attn + (1-attention_threshold) * torch.eye(attentions.shape[-1])
        rollout = torch.matmul(attn, rollout)

    # Normalize final matrix
    rollout = rollout / rollout.sum(dim=-1, keepdim=True)

    return rollout.numpy()

我们在 GLUE 基准上对比了不同方法的性能：