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

Wenshuai Yao; Wenyong Zhou; Hanyong Shao; Zhiyuan Ning; Yizhe Chen; Yuannuo Feng; Kechao Tang; Ru Huang

cs.LGOct 15, 2026

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

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

Organizations: School of Integrated Circuits, Peking University, Beijing, China · Department of Electrical and Computer Engineering, The University of Hong Kong, Hong Kong · School of Integrated Circuit Science and Engineering, Beihang University, Beijing, China · Beijing Advanced Innovation Center for Integrated Circuits, Beijing, China

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%.

Figures & tables

Figure 1. Attention maps from clean inference and from the noisy model, NAT, and ABNAT under CIM noise.

Figure 2. Layerwise attention drift across CIM noise intensities; shaded bands denote std across seeds.

Figure 3. Accuracy degradation versus mean attention drift across CIM noise settings.

Figure 4. Sensitivity scores $\mathrm{sens}_{l,h}$ across all Transformer blocks and attention heads.

Figure 5. ABNAT with sensitivity-guided late-block selection and task, logit-distillation, and attention-alignment objectives.

Variant	Task loss	Logit KD	Feature ( Zhou et al., 2025 )	Attn. Align. (Ours)	$\sigma=0.05$	$\sigma=0.10$	$\sigma=0.15$	$\sigma=0.20$	Avg.
Original+Noise	–	–	–	–	86.28 $\pm$ 0.37	64.34 $\pm$ 1.37	10.06 $\pm$ 1.51	2.09 $\pm$ 0.61	40.69
NAT ( Rasch et al., 2023 )	$\checkmark$	–	–	–	84.98 $\pm$ 0.35	79.23 $\pm$ 0.51	71.52 $\pm$ 0.67	58.86 $\pm$ 0.66	73.65
KD	–	$\checkmark$	–	–	87.21 $\pm$ 0.15	82.59 $\pm$ 0.32	75.00 $\pm$ 0.53	63.29 $\pm$ 0.74	77.02
FeatureAlign	–	–	$\checkmark$	–	86.46 $\pm$ 0.28	80.36 $\pm$ 0.50	43.27 $\pm$ 1.33	14.72 $\pm$ 1.59	56.20
AttnAlign	–	–	–	$\checkmark$	86.22 $\pm$ 0.29	78.68 $\pm$ 0.49	54.76 $\pm$ 1.14	12.91 $\pm$ 1.21	58.14
NAT+KD ( Charan et al., 2020 )	$\checkmark$	$\checkmark$	–	–	87.20 $\pm$ 0.17	82.69 $\pm$ 0.29	74.99 $\pm$ 0.54	62.45 $\pm$ 0.76	76.83

Table 1. Full component ablation for ViT-Tiny on CIFAR-100 under matched train–deploy noise. Accuracy under noisy inference is reported as mean

\pm

std over 30 seeds. Shaded rows denote the two KD-based ABNAT variants; ABNAT without

L_{task}

retains logit KD and our attention-alignment loss.

Figure 6. Cross-setting results for ViT-Small on CIFAR-100, ViT-Tiny on ImageNet-100, Swin-Tiny on CIFAR-100, Pythia-160M and Pythia-2.8B on WikiText-103, and CLIP ViT-B/32 on Flickr30K.

Figure 7. Train–deploy noise mismatch for ViT-Tiny on CIFAR-100 at

\sigma_{\text{train}}\!\in\!\{0.15,0.20\}

Figure 8. Layerwise attention drift of NAT, NAT+KD, and ABNAT at $\sigma\!=\!0.20$ .

Figure 9. Prediction confidence versus attention-route fidelity at

\sigma\!=\!0.20

on the common anchor subset.

Figure 10. Accuracy and training cost under selective block alignment.

$\lambda_{logits}$		$\lambda_{attn}$		$\tau_{\mathrm{KD}}$		$\tau_{attn}$
Value	Avg.	Value	Avg.	Value	Avg.	Value	Avg.
0.25	78.41	1	78.52	0.5	79.36	0.125	67.57
0.5	79.00	5	79.44	1	80.09	0.25	74.53
1	79.52	10	79.85	2	80.17	0.5	78.47
3	80.09	25	80.16	3	79.90	1	80.09
5	80.19	50	80.09	4	79.80	2	79.93

Table 2. Hyperparameter sensitivity for ViT-Tiny on CIFAR-100. Bold marks the best setting in each sweep; shading marks defaults.

Noise Robustness

Noise robustness studies how to preserve reliable predictions when data, labels, signals, or computing hardware are corrupted.

See all 93 papers

Trilinear Compute-in-Memory Architecture for Energy-Efficient Transformer Acceleration
Apr 8, 2026Md Zesun Ahmed Mia, Jiahui Duan, Kai Ni +1Compute-in-Memory Efficient Transformer Inference
The Pennsylvania State University, Department of Electrical Engineering, University Park, PA, USA · University of Notre Dame, Department of Electrical Engineering, Notre Dame, IN, USA
NoiseFormer -- Noise Diffused Symmetric Attention Transformer
Jan 10, 2026Phani Kumar, Nyshadham, Jyothendra Varma +1Transformer Attention Attention Mechanisms
Intel Corporation
Compute-in-Memory Attention: A Time-Domain Analog Softmax Circuit with RC-Tunable Temperature
Sep 2, 2026Ankur Singh, Ashish Gautam, Shruti R. Kulkarni +1Softmax Attention Analog Computing
Department of Electrical and Computer Engineering, University of Wisconsin–Madison, Madison, WI, USA. · ORNL Oak Ridge, Tennessee, USA.

Apr 8, 2026J/K move · Enter open · S save

Trilinear Compute-in-Memory Architecture for Energy-Efficient Transformer Acceleration

Md Zesun Ahmed Mia, Jiahui Duan, Kai Ni, Abhronil Sengupta

The Pennsylvania State University, Department of Electrical Engineering, University Park, PA, USA · University of Notre Dame, Department of Electrical Engineering, Notre Dame, IN, USA

Compute-in-Memory Efficient Transformer Inference Energy-Efficient ML Transformer Attention

Open paper AskPDF arXiv

Self-attention in Transformers generates dynamic operands that force conventional Compute-in-Memory (CIM) accelerators into costly non-volatile memory (NVM) reprogramming cycles, degrading throughput and stressing device endurance. Existing solutions either reduce but retain NVM writes through matrix decomposition or sparsity, or move attention computation to digital CMOS at the expense of NVM density. We present TrilinearCIM, a Double-Gate FeFET (DG-FeFET)-based architecture that uses back-gate modulation to realize a three-operand multiply-accumulate primitive for in-memory attention computation without dynamic ferroelectric reprogramming. Evaluated on BERT-base (GLUE) and ViT-base (ImageNet and CIFAR), TrilinearCIM outperforms conventional CIM on seven of nine GLUE tasks while achieving up to 46.6% energy reduction and 20.4% latency improvement over conventional FeFET CIM at 37.3% area overhead. To our knowledge, this is the first architecture to perform complete Transformer attention computation exclusively in NVM cores without runtime reprogramming.

Abstract

Noise Robustness

Explore similar work

Trilinear Compute-in-Memory Architecture for Energy-Efficient Transformer Acceleration

NoiseFormer -- Noise Diffused Symmetric Attention Transformer

Compute-in-Memory Attention: A Time-Domain Analog Softmax Circuit with RC-Tunable Temperature