面向 Transformer-SSM 混合体的检索感知蒸馏

状态空间模型 (SSM) 在序列建模方面提供了高效性，但在需要原语检索的基准测试上落后于 Transformer。先前的工作将这一差距与一小组注意力头（称为 Gather-and-Aggregate (G&A)）相关联，这些头 SSM 难以复制。我们提出了*检索感知蒸馏*，将预训练的 Transformer 转换为仅保留这些检索关键头并将其余部分蒸馏到循环头的混合体。通过对合成检索任务进行消融实验，我们识别了关键头，产生稀疏、非均匀注意力布局的混合体。我们表明，**仅保留 2% 的注意力头即可在检索密集型任务上恢复超过 95% 的教师性能**（在 1B 模型中仅 10 个头），所需头数远少于保留至少 25% 的混合体。我们进一步发现，较大的循环状态往往能补偿缺失的检索：一旦检索由这些头处理，SSM 主干就可以在极少的损失下进行简化，甚至可以减少 $8\times$ 的状态维度。通过减少注意力缓存和 SSM 状态，该混合体比可比混合体高效 $5$--$6\times$，在较低的内存成本下弥合了 Transformer-SSM 之间的差距。