ABNAT：面向模拟存内计算系统中鲁棒 Transformer 的基于注意力的噪声感知训练

arXiv:2610.22065v1 发布类型：新 摘要：基于 Transformer 的模型已成为众多任务的 backbone，但其密集的计算和数据搬运会带来显著的能耗开销。模拟存内计算（CIM）提供了一种高能效替代方案，但其固有的硬件非理想性会扰动存储权重，从而显著降低模型性能。尽管现有的噪声感知训练（NAT）方法通过在训练期间注入噪声来提升鲁棒性，但它们在很大程度上忽视了注意力机制在 Transformer 推理中的关键作用。为弥补这一不足，我们分析了 CIM 引入的噪声如何重塑 Transformer 中的注意力，并得出两点关键观察：（1）在更深的层中，干净注意力图与噪声注意力图之间的差异显著增大；（2）即使其注意力路径与干净参考结果存在明显偏离，传统 NAT 仍可能表现得过度自信。受此启发，我们提出 ABNAT，一种基于注意力的噪声感知训练框架，使噪声学生模型学习保持干净教师模型的逐头注意力路径，而不仅仅是复现答案；该方法无需改动硬件，也不会增加推理开销，其选择性变体还可进一步降低训练成本。在视觉、语言和多模态 Transformer 上的大量实验表明，相较于对应的 NAT 基线，ABNAT 能够持续提升鲁棒性，使受噪声影响的准确率最高提升 11.28%，并将注意力偏差最高降低 85.77%。