Understanding In-Context Learning Beyond Transformers: An Investigation of State Space and Hybrid Architectures

这是一篇关于大语言模型（LLM）上下文学习（In-Context Learning, ICL）机制的深度研究论文。该研究不仅关注传统的 Transformer 架构，还扩展到了状态空间模型（SSM，如 Mamba）及其混合架构（Hybrid Architectures），旨在揭示不同架构在实现 ICL 时的内部机制差异。

以下是该论文的详细技术总结：

1. 研究问题 (Problem)

• 背景：ICL 是 LLM 在测试时仅通过少量示例（few-shot）进行推理而不更新参数的能力。既往研究主要集中在 Transformer 架构上，发现**函数向量（Function Vectors, FVs）和归纳头（Induction Heads）**是 ICL 的关键机制。
• 缺口：随着 Mamba 等高效状态空间模型（SSM）及其与 Transformer 的混合模型（如 Hymba, Zamba2）的兴起，这些新架构是否具备 ICL 能力？如果具备，其内部机制（特别是 FVs 的作用）是否与 Transformer 相同？目前尚缺乏针对这些非 Transformer 架构的深入机理分析。
• 核心问题：不同架构（Transformer, SSM, Hybrid）在 ICL 任务中的表现是否一致？其内部负责 ICL 的机制（如 FVs 的分布和重要性）有何异同？

2. 方法论 (Methodology)

研究采用了**行为分析（Behavioral Analysis）与机理可解释性分析（Mechanistic Interpretability）**相结合的方法。

A. 模型与任务

• 模型：选取了参数量级相近（约 1B-1.5B）的多种架构：
  • Transformer: Gemma-3-1B, Llama-3.2-1B, Qwen2.5-1.5B。
  • SSM: Mamba-1.4B, Mamba2-1.3B。
  • 混合架构: Hymba-1.5B (并行混合), Zamba2-1.2B (串行混合)。
• 任务分类：将 ICL 任务细分为两类（基于 Jin et al., 2025）：
  1. 参数化知识检索 (Parametric Knowledge Retrieval)：如国家 - 首都、货币转换，依赖模型内部存储的知识。
  2. 上下文知识理解 (Contextual Knowledge Understanding)：如情感分析、仇恨言论检测，依赖对提示词中上下文信息的理解。

B. 行为实验

• 标签随机化 (Label Randomization)：测试不同比例正确示例（0% 到 100%）对模型性能的影响，评估模型对噪声的鲁棒性。
• 标签翻转 (Label Flipping)：在上下文中将正确标签映射为错误标签，测试模型是否能学习反事实关系（即忽略预训练知识，仅学习上下文的新规则）。

C. 机理分析

• FV 头识别：复现了 Todd et al. (2024) 的方法，计算平均间接效应 (AIE) 来识别负责 ICL 的“函数向量头”（FV Heads）。对于 SSM 和混合模型，将 SSM 层视为类比于注意力头。
• 干预实验 (Intervention)：
  • 导向 (Steering)：在推理时向特定头的输出添加计算出的 FV，观察性能提升。
  • 消融 (Ablation)：将特定头的输出置零或替换为均值，观察性能下降。
  • 分层分析：分别对混合模型中的 Attention 流和 Mamba 流进行独立干预。

3. 关键发现与结果 (Key Findings & Results)

A. 行为层面：表现相似，内部不同

• 所有架构在定性上表现出相似的 ICL 能力（即随着正确示例增加，性能提升）。
• 差异点：在上下文知识理解任务中，Mamba2 和 Zamba2 在提供黄金标签（Gold）示例时，性能提升幅度远小于 Transformer 和 Hymba，显示出较弱的 ICL 能力。但在标签翻转实验中，所有模型（包括 Mamba2）都能学习新的反事实关系，表明它们具备 ICL 潜力，只是在不同任务类型上的表现存在差异。

B. 机理层面：FVs 的分布与重要性

1. FVs 的主要位置：
   • 在 Transformer、Mamba 和混合模型中，FVs 主要位于自注意力（Self-Attention）层。
   • 在混合模型（Hymba, Zamba2）中，ICL 能力主要由自注意力层中的 FVs 驱动，无论这些层是并行还是串行堆叠。
2. Mamba2 的特殊性：
   • Mamba2 似乎不使用 FVs 机制进行 ICL。对其 FV 头进行导向（Steering）并未带来显著性能提升，但其性能在消融实验中仍会下降。
   • 推测：Mamba2 可能使用了一种不同于 FVs 的机制来实现 ICL，或者由于其多头设计导致单头隐藏维度较小，难以在头级别捕捉特定任务信息。
3. 任务类型依赖性：
   • 参数化知识检索：FVs 起主导作用。不同任务识别出的 Top FV 头高度集中且一致。
   • 上下文知识理解：FVs 的作用较小，且不同任务识别出的 FV 头重叠度低（重叠率通常很低，如图 7 所示）。
   • 混合模型中的 Mamba 流：在上下文理解任务中，对混合模型中的 Mamba 流进行导向甚至会导致性能下降，暗示 Mamba 流可能承担其他功能（如长序列建模），而非 ICL 的核心机制。

C. 混合架构的结论

• 混合模型的 ICL 能力主要由自注意力层控制。
• 在 Hymba 中，虽然 Mamba 流对整体性能有贡献（消融 Mamba 流会导致性能下降），但 ICL 的特定机制（FVs）仍集中在 Attention 流中。

4. 主要贡献 (Contributions)

1. 扩展了 FV 研究的边界：首次将 FV 分析从 Transformer 扩展到 SSM 和混合架构，揭示了不同架构下 ICL 机制的异同。
2. 任务类型的细分：明确了“参数化知识检索”与“上下文知识理解”在 ICL 机制上的本质区别（前者依赖 FVs，后者依赖更分散的机制）。
3. 发现 Mamba2 的独特性：提出了 Mamba2 可能不依赖传统 FVs 机制进行 ICL 的假设，为理解 SSM 的 ICL 能力提供了新视角。
4. 方法论贡献：强调了结合行为实验与机理干预（Steering/Ablation）对于全面理解 LLM 能力的重要性。

5. 意义与局限性 (Significance & Limitations)

• 意义：
  • 为设计更高效的混合模型提供了理论依据（例如，为了增强 ICL，应优化自注意力层）。
  • 加深了对非 Transformer 架构（如 Mamba）内部工作原理的理解，打破了"ICL 仅由 Attention 机制驱动”的单一认知。
  • 指出了不同任务类型可能需要不同的内部机制，这对模型微调和应用场景选择具有指导意义。
• 局限性：
  • 研究仅限于预训练模型，未控制预训练数据差异（不同模型的预训练语料不同）。
  • 主要关注 FVs 机制，未深入探讨其他可能的机制（如归纳头 Induction Heads 在 SSM 中的表现）。
  • 计算资源限制导致实验规模较小（约 1B 参数模型）。

总结

该论文通过严谨的行为和机理实验，揭示了ICL 并非单一机制。在 Transformer 和混合模型中，自注意力层中的函数向量（FVs）是参数化知识检索任务的核心驱动力；而在上下文理解任务中，机制更为分散。特别值得注意的是，Mamba2表现出与 Transformer 和 Mamba-1 不同的行为，暗示其可能采用了不同的 ICL 实现路径。这一发现对于未来设计兼具高效性与强大 ICL 能力的下一代大模型具有重要参考价值。