ABNAT: Attention-Based Noise-Aware Training for Robust Transformers on Analog Compute-in-Memory Systems

📄 arXiv: 2610.22065v1 📥 PDF

作者: Wenshuai Yao, Wenyong Zhou, Hanyong Shao, Zhiyuan Ning, Yizhe Chen, Yuannuo Feng, Kechao Tang, Ru Huang

分类: cs.LG

发布日期: 2026-10-15

备注: 8 pages, 10 figures, 2 tables. Accepted for publication at the 32nd Asia and South Pacific Design Automation Conference (ASP-DAC 2027)

DOI: 10.1145/3840404.3848194

💡 一句话要点

提出ABNAT以解决模拟计算内存系统中的噪声影响问题

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 噪声感知训练 Transformer 模拟计算内存 注意力机制 鲁棒性提升 深度学习 多模态学习

📋 核心要点

  1. 现有的噪声感知训练方法未能充分考虑注意力机制在Transformer推理中的重要性，导致模型在噪声环境下表现不佳。
  2. ABNAT框架通过关注注意力机制，教会模型在噪声条件下保持与干净模型的注意力路线一致，从而提高鲁棒性。
  3. 实验结果显示，ABNAT在多种任务上均优于传统的噪声感知训练方法，噪声降级准确率提升高达11.28%，注意力偏差降低85.77%。

📝 摘要（中文）

基于Transformer的模型已成为多种任务的核心，但其高强度计算和数据移动导致显著的能耗。模拟计算内存（CIM）提供了一种节能的替代方案，但其固有的硬件非理想性会扰动存储权重，显著降低模型性能。现有的噪声感知训练方法虽然通过在训练中注入噪声来提高鲁棒性，但往往忽视了注意力机制在Transformer推理中的关键作用。为此，本文提出了ABNAT，一个基于注意力的噪声感知训练框架，旨在教会噪声学生保留干净教师的头部注意力路线，而不仅仅是重现答案。实验结果表明，ABNAT在视觉、语言和多模态Transformer上均显著提高了鲁棒性，噪声降级准确率提升高达11.28%。

🔬 方法详解

问题定义：本文旨在解决模拟计算内存系统中噪声对Transformer模型性能的影响。现有的噪声感知训练方法未能有效应对注意力机制的变化，导致模型在噪声环境下的鲁棒性不足。

核心思路：ABNAT通过分析CIM引入的噪声如何重塑Transformer中的注意力，提出了一种新的训练策略，强调保持干净教师模型的注意力路线，而不仅仅是输出结果。

技术框架：ABNAT的整体架构包括噪声学生和干净教师两个模型，训练过程中通过对比注意力路线来优化学生模型的表现。该框架不需要硬件更改或推理开销，且具有选择性变体以降低训练成本。

关键创新：ABNAT的主要创新在于将注意力机制与噪声感知训练结合起来，解决了传统方法在注意力偏差上的过度自信问题，从而显著提高了模型在噪声环境下的鲁棒性。

关键设计：在训练过程中，ABNAT使用了特定的损失函数来量化学生模型与教师模型之间的注意力路线差异，并通过动态调整训练参数来优化模型性能。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，ABNAT在多个任务上均显著优于传统的噪声感知训练基线，噪声降级准确率提升高达11.28%，同时注意力偏差降低了85.77%，展示了其在提高模型鲁棒性方面的有效性。

🎯 应用场景

该研究的潜在应用领域包括自然语言处理、计算机视觉和多模态学习等，尤其是在资源受限的环境中，ABNAT能够有效提升模型在噪声条件下的鲁棒性，具有重要的实际价值和广泛的应用前景。

📄 摘要（原文）

Transformer-based models have become the backbone of a wide range of tasks, but their intensive computation and data movement incur substantial energy overhead. Analog compute-in-memory (CIM) offers an energy-efficient alternative, yet its inherent hardware non-idealities perturb stored weights and can significantly degrade model performance. Although existing noise-aware training (NAT) methods improve robustness by injecting noise during training, they largely overlook the critical role of attention mechanisms in Transformer inference.To bridge this gap, we analyze how CIM-induced noise reshapes attention in Transformers and identify two key observations: (1) the divergence between clean and noisy attention maps increases substantially in deeper layers, and (2) conventional NAT can remain overconfident even when its attention routes deviate markedly from the clean reference. Motivated by these findings, we propose ABNAT, an attention-based noise-aware training framework that teaches the noisy student to preserve the clean teacher's headwise attention routes rather than merely reproduce answers, without hardware changes or inference overhead; selective variants further reduce training cost.Extensive experiments across vision, language, and multimodal Transformers show that ABNAT consistently improves robustness over corresponding NAT baselines, improving noise-degraded accuracy by up to 11.28% and reducing attention deviation by up to 85.77%.