Attention based training improves transformer resilience on analog chips
ABNAT: Attention-Based Noise-Aware Training for Robust Transformers on Analog Compute-in-Memory Systems
Machine Learning
Summary
Transformers are powerful models used in many tasks but use a lot of energy. Analog compute-in-memory (CIM) chips can save energy but introduce noise that hurts transformer performance. The authors found that existing training methods don't focus enough on how noise affects the transformer's attention mechanism. They developed a new training approach, ABNAT, that teaches the model to keep its attention patterns more like a noise-free version, leading to better performance when running on noisy CIM hardware.
What this means in practice
- •For hardware engineers: Improve transformer model accuracy running on energy-efficient analog CIM chips by using attention-focused noise-aware training.
- •For machine learning engineers: Train transformer models to maintain stable attention patterns under hardware noise, enhancing robustness without extra inference cost.
Authors
Wenshuai Yao, Wenyong Zhou, Hanyong Shao, Zhiyuan Ning, Yizhe Chen, Yuannuo Feng, Kechao Tang, Ru Huang
Abstract
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%.