跳转至

Head Pursuit: Probing Attention Specialization in Multimodal Transformers

会议: NeurIPS2025
arXiv: 2510.21518
代码: GitHub
领域: 图像生成
关键词: [attention head specialization, Matching Pursuit, Logit Lens, model editing, vision-language model]

一句话总结

将经典稀疏信号恢复算法(SOMP)重新解释为一种多样本可解释性工具,发现 LLM 和 VLM 中注意力头存在细粒度语义专业化现象,仅通过翻转约 1% 的头即可可靠地抑制或增强特定概念(如国家名、毒性内容、颜色等)的生成。

研究背景与动机

领域现状:大规模生成模型(LLM、VLM)在各类任务上表现出色,但内部机制仍未完全理解。已有研究表明注意力头存在功能角色(语法追踪、复制行为、事实回忆),但这些发现通常基于启发式方法,难以跨样本泛化。

现有痛点:(1) Logit Lens 等可解释性工具一次只分析一个样本,无法稳定量化头的重要性;(2) Attention Lens 需要为每个头训练单独的线性探针,计算代价高;(3) 现有的头编辑方法缺乏数学基础,依赖试错。

核心矛盾:需要一种既能多样本聚合分析又无需额外训练的方法来系统性地发现头的语义专业化。

本文目标:提供数学上有原则的方法来识别、量化和利用注意力头在特定语义领域的专业化。

切入角度:将 Logit Lens 推广为稀疏信号恢复问题,用 SOMP(Simultaneous Orthogonal Matching Pursuit)在 unembedding 字典上做多样本稀疏分解。

核心 idea:注意力头的输出可以用 unembedding 矩阵中少量语义方向的稀疏线性组合近似,方差解释比可以量化头对目标概念的专业化程度。

方法详解

整体框架

方法分三步:(1) 给定数据集,计算每个注意力头对 residual stream 的贡献矩阵 \(\mathbf{H}_{h,l} \in \mathbb{R}^{n \times d}\);(2) 对目标概念(如颜色、国家),将 unembedding 矩阵限制为相关 token 行,用 SOMP 做稀疏分解,按方差解释比排序头;(3) 对排名靠前的头进行干预(翻转符号/缩放),观察对目标概念生成的影响。

关键设计

  1. SOMP 稀疏分解作为多样本 Logit Lens:

    • 功能:用 SOMP 算法在 unembedding 矩阵 \(\mathbf{D} \in \mathbb{R}^{v \times d}\) 上对头激活 \(\mathbf{H} \in \mathbb{R}^{n \times d}\) 做稀疏分解,找到最能解释头行为的少量语义方向
    • 核心思路:SOMP 每步选择与所有样本残差最大关联的字典原子 \(p^t = \arg\max_j \|\mathbf{D}[j]\mathbf{R}^{tT}\|_1\),加入支撑集后做最小二乘重拟合 \(\mathbf{W}^t = \arg\min_{\mathbf{W}} \|\mathbf{H} - \mathbf{W}\mathbf{D}[\mathbb{S}^{t+1}]\|_F\)。与 Logit Lens 的联系:LL 等价于在单样本上做单步 Matching Pursuit,SOMP 是它的多样本多步推广
    • 设计动机:单样本 LL 结果噪声大且冗余(Table 6),多样本 SOMP 通过跨数据集聚合获得稳定的头语义特征
  2. 基于方差解释比的头选择与干预:

    • 功能:给定目标概念(如"颜色"),将字典限制为颜色相关 token 的 unembedding 行,用 SOMP 分解每个头,按方差解释比 \(\|\mathbf{H}_r\|_F^2 / \|\mathbf{H}\|_F^2\) 排序,选择 top-\(k\) 头进行干预
    • 核心思路:干预方式为缩放头对 residual stream 的贡献——抑制时 \(\alpha = -1\)(翻转符号),增强时 \(\alpha = 5\)(放大 5 倍)。关键发现是极少量头(8-32 个,约 0.8%-3%)即可显著影响目标概念
    • 设计动机:如果头的方差在概念限制字典上被 SOMP 解释得很好,说明该头主要输出该概念相关的信号,干预它应当针对性地影响概念生成

损失函数 / 训练策略

方法完全无训练——不修改模型权重,仅在推理时对选定头的输出做缩放干预。所有实验中头选择基于训练数据完成,评估在不相交的测试数据上进行。对照实验使用随机选择的同等数量、同层分布的头作为 baseline。

实验关键数据

主实验

问答(Mistral-7B on TriviaQA,F1 score)

干预头数 目标(国家)概念性能↓ 非目标性能↓ 随机头影响
8头 (0.8%) 显著下降 轻微下降 无显著影响
16头 大幅下降 中等下降 无显著影响
32头 严重下降 中等下降 无显著影响

毒性缓解(归一化毒性生成计数↓)

数据集 8头 SOMP 8头 LL 8头 Random 32头 SOMP 32头 LL 32头 Random
RTP 0.83 0.91 1.02 0.66 0.71 1.13
TET 0.83 0.81 0.97 0.49 0.68 0.95

消融实验

LLaVA 图像分类(翻转头后的归一化准确率)

数据集 16头 SOMP↓ 16头 Random 32头 SOMP↓ 32头 Random
MNIST 大幅下降 无变化 严重下降 无变化
SVHN 大幅下降 无变化 严重下降 无变化
GTSRB 大幅下降 无变化 严重下降 无变化
EuroSAT 大幅下降 无变化 严重下降 无变化

Flickr30k 图像描述(抑制/增强颜色,16头)

干预 颜色关键词频率 CIDEr 保持率
抑制 (\(\alpha=-1\)) 接近零 >80%
增强 (\(\alpha=5\)) +60%以上 >80%

关键发现

  • SOMP 选择的头具有高度概念针对性:干预目标概念性能下降远大于非目标概念
  • Logit Lens 选择的头相关但非特异——等量降低目标和非目标性能
  • Jaccard 相似度分析揭示语义相似的任务共享头(MNIST/SVHN 重叠高,EuroSAT/RESISC45 重叠高)
  • 增强干预同样有效:\(\alpha=5\) 可以使颜色、情感、数量词频增加 60%+
  • 跨 VLM 验证(LLaVA-13B、Gemma3-12B、Qwen2.5-VL-7B)均确认相同趋势

亮点与洞察

  • 优雅的理论联系:将 Logit Lens 重新解释为 Matching Pursuit 的单步单样本特例,SOMP 是其自然推广
  • 极高的干预效率:仅 0.8%-3% 的头即可显著控制生成,暗示注意力层存在高度结构化的线性语义子空间
  • 双向可控:抑制和增强都有效,且不严重损害整体生成质量(CIDEr >80%)
  • 跨模态一致性:文本和视觉-语言任务上观察到一致的头专业化模式,支持"概念在 residual stream 中以线性方式编码"的假说

局限与展望

  • SOMP 假设线性分解,可能无法捕捉头表示的非线性结构
  • 语义字典的质量和覆盖度直接影响发现质量——不完整的关键词列表可能导致偏差
  • 干预机制较粗犷(全局缩放),未区分不同位置或模态 token 上的干预效果
  • 未探索用于图像生成(如 VLM 的图像解码阶段)的控制应用

相关工作与启发

  • Logit Lens / Tuned Lens:前者不需训练但只看单步单样本,后者需要为每层训练探针;本文方法在两者之间取得平衡
  • CLIP 的头分解:Gandelsman et al. (2024) 和 HeadPursuit 类似地分解 CLIP 的注意力头,但后者适用于生成式模型
  • 事实编辑(ROME/MEMIT):针对 MLP 层的知识编辑,本文转向注意力头层面的语义控制
  • 启发:注意力头的专业化结构可能为轻量级模型对齐和安全控制提供新途径——不需要微调,只需识别并缩放少量头

评分

  • 新颖性: ⭐⭐⭐⭐ 将 SOMP 引入 LLM 可解释性的视角新颖,理论联系清晰
  • 实验充分度: ⭐⭐⭐⭐⭐ 覆盖 QA、毒性缓解、图像分类、图像描述四类任务和五个模型,对照充分
  • 写作质量: ⭐⭐⭐⭐ 方法和实验层层递进,论证充分
  • 价值: ⭐⭐⭐⭐ 为模型理解和轻量级控制提供了实用工具