[ICLR 2025] SEKA：突破 FlashAttention 限制，毫秒级的精准注意力引导

TL;DR

传统的注意力引导方法（Attention Steering）往往需要在 GPU 内存中展开巨大的注意力矩阵，这不仅拖慢了速度，还让 FlashAttention 等优化技术无用武之地。本文提出的 SEKA (Spectral Editing Key Amplification) 换了个思路：不在结果上动刀，而是在输入端干预。通过对 Key 向量进行谱编辑，SEKA 实现了几乎零成本的提示词高亮（Prompt Highlighting），在长文本任务中表现优异。

1. 痛点：为什么“后处理”注意力矩阵行不通？

在处理复杂指令或长文档时，我们常常希望模型“多看一眼”某些关键信息（即 Prompt Highlighting）。现有的 SOTA 方法如 PASTA 采用了**后处理（Post-hoc）**逻辑：

1. 计算出所有的 Attention Scores。
2. 手动放大某些 Token 的权重。
3. 重新归一化。

这种方案的本质瓶颈在于：它要求把 $TimesT$ 的注意力矩阵完整显存化。这与 FlashAttention（为了提速会避免生成完整矩阵）完全冲突。在长文本时代，这种内存爆炸和延迟增加（+1.03s 每样本）是不可接受的。

2. 核心直觉：Key 空间中的“相关性”轨迹

作者发现了一个有趣的物理直觉：当一个 Token 变得“相关”时，它的 Key 向量在空间中会发生一致的位移。

如图所示，某些特定的注意力头中，从“不相关”到“相关”的 Key 向量存在明显的方向性（Spectral Direction）。这意味着我们可以通过一个简单的线性投射，在注意力计算发生之前，就把 Key 向量“推”向相关性更高的区域。

3. 方法论：SEKA 与 AdaSEKA

SEKA：谱分解提取相关性

SEKA 的核心是 SVD（奇异值分解）。作者构建了中性、正向和负向的对比提示词，提取 Key 向量的交叉协方差矩阵，从而识别出最能代表“相关性”的方向： $k^{\prime }=k+\frac{g^{+}{P}^{+}k+g^{-}{P}^{-}k}{2}$ 这种操作在代数上等价于给注意力 Logits 增加了一个低秩偏置，但它直接发生在 Embedding 层面，因此完美适配 FlashAttention。

AdaSEKA：动态“专家”路由

面对不同的任务（如事实识别 vs 指令遵循），固定的投影矩阵可能不够。AdaSEKA 引入了多专家机制：

1. 预先学习多个任务的专家子空间。
2. 推理时，根据当前查询（Query）的语义意图，自动路由并加权组合专家投影。

4. 实验战果：反转“迷失在中间”

长文本模型普遍存在“Lost-in-the-Middle”现象，即对文本中间的信息召回率极低。SEKA 展现了强大的控制力：通过仅高亮中间段落，SEKA 成功将 U 型曲线反转成了倒 U 型，显著提升了模型对中间区域的召回（Exact Match）。

性能对比

在标准基准测试（CounterFact, Bias in Bios）中，SEKA 在 Qwen3 和 Gemma3 家族模型上均达到了顶级水平，尤其是在参数量较小的模型（如 4B）上提升尤为显著。

5. 深度洞察：为什么选择特定的头？

并非所有的注意力头都对“高亮”敏感。作者通过实验发现，处理相关性的“检索头（Retrieval Heads）”主要集中在模型的中后期层。SEKA 引入了一个阈值 ${\delta }_{min}$，只在检测到 Key 向量偏移明显的敏感头上施加干预。这种选择性干预避免了对模型原生能力的破坏。

6. 总结

SEKA 的成功告诉我们，大模型的注意力机制不仅可以从外部“观测”，还可以通过底层的几何变换进行“导航”。它以极低的计算代价，解决了长文本处理中的关键痛点。

局限性：尽管目前在事实召回和简单指令上效果惊人，但对于需要深层推理的长逻辑链，单纯依靠 Key 向量的放大是否足够，仍需进一步验证。