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
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 sensl,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)
σ=0.05
σ=0.10
σ=0.15
σ=0.20
Avg.
Original+Noise
–
–
–
–
86.28 ± 0.37
64.34 ± 1.37
10.06 ± 1.51
2.09 ± 0.61
40.69
NAT ( Rasch et al., 2023 )
✓
–
–
–
84.98 ± 0.35
79.23 ± 0.51
71.52 ± 0.67
58.86 ± 0.66
73.65
KD
–
✓
–
–
87.21 ± 0.15
82.59 ± 0.32
75.00 ± 0.53
63.29 ± 0.74
77.02
FeatureAlign
–
–
✓
–
86.46 ± 0.28
80.36 ± 0.50
43.27 ± 1.33
14.72 ± 1.59
56.20
AttnAlign
–
–
–
✓
86.22 ± 0.29
78.68 ± 0.49
54.76 ± 1.14
12.91 ± 1.21
58.14
NAT+KD ( Charan et al., 2020 )
✓
✓
–
–
87.20 ± 0.17
82.69 ± 0.29
74.99 ± 0.54
62.45 ± 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 ± std over 30 seeds. Shaded rows denote the two KD-based ABNAT variants; ABNAT without Ltask 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 σtrain∈{0.15,0.20} .
Figure 8. Layerwise attention drift of NAT, NAT+KD, and ABNAT at σ=0.20 .
Figure 9. Prediction confidence versus attention-route fidelity at σ=0.20 on the common anchor subset.
Figure 10. Accuracy and training cost under selective block alignment.
λlogits
λattn
τKD
τ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.
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.
Md Zesun Ahmed Mia, Jiahui Duan, Kai Ni +1
The Pennsylvania State University, Department of Electrical Engineering, University Park, PA, USA · University of Notre Dame, Department of Electrical Engineering, Notre Dame, IN, USA
Softmax is a key operation in Transformer attention, but its exponentiation and normalization add significant overhead in compute-in-memory (CIM) accelerators, especially when analog attention scores must first be converted to the digital domain. This work presents a tunable-temperature analog softmax circuit in GlobalFoundries 22-nm fully depleted silicon-on-insulator (FDSOI) technology that operates directly on CIM-generated score voltages without intermediate analog-to-digital conversion. Each input score is converted into a time-domain event using a shared falling ramp. The corresponding comparator transition samples an RC-decaying reference to generate an exponential weight, which is then processed by an in-circuit normalization stage. In contrast to analog softmax circuits that rely on transistor weak-inversion behavior for exponentiation, the proposed architecture controls the softmax response through the ramp slope and RC time constant, enabling programmable effective temperature. The 128-element architecture is evaluated using transistor-level and post-layout extracted simulations, including multi-level input vectors, capacitance variation and mismatch, process and temperature variation, monte carlo analysis, and shared-interconnect parasitics. The complete 128-element implementation occupies 9453.42~μm2 including the shared global ramp circuitry, while each replicated softmax element occupies 70.2~μm2. The circuit achieves a 242.97-ns evaluation latency at 13.44mW total power, corresponding to 25.5pJ per output element. The simultaneous 128-element evaluation achieves an RMSE of 24.46~mV relative to the ideal softmax response. The extracted circuit characteristics are further incorporated into a MemTorch-based hardware-aware Transformer model, where the proposed softmax achieves a validation loss within 2.5% of the ideal-softmax baseline.
Ankur Singh, Ashish Gautam, Shruti R. Kulkarni +1
Department of Electrical and Computer Engineering, University of Wisconsin–Madison, Madison, WI, USA. · ORNL Oak Ridge, Tennessee, USA.
Analog compute-in-memory (CIM) enables energy-efficient neural network inference, but device variation and read noise can severely degrade low-bit quantized models. Existing CIM-oriented quantization methods mainly minimize ideal quantization error, ignoring the hardware noise floor and thus causing inefficient precision allocation. We propose NANQ, a noise-aware mixed-precision non-uniform quantization framework for analog CIM. NANQ models magnitude-dependent weight noise from measured responses of an eFlash CIM array and converts the noise profile into an adaptive quantization density, assigning finer resolution to low-noise regions while avoiding ineffective precision in noise-dominated regions. It further assigns layer-wise bit-widths by identifying each layer's precision saturation point under hardware noise using a unified threshold. On-chip experiments on an eFlash CIM SoC show that, under 2-bit weight-magnitude quantization, NANQ improves vision-model accuracy by 8.05 percentage points and reduces language-model PPL by 54.7% on average over PowerQuant. Mixed-precision NANQ captures most of the gains obtainable from additional quantization resources with only 3.2-3.8 equivalent bits.
Yizhe Chen, Wenshuai Yao, Saiya Wang +6
School of Integrated Circuit Science and Engineering, Beihang University, Beijing, China · School of Integrated Circuits, Peking University, Beijing, China · Department of Electrical and Computer Engineering, The University of Hong Kong, Hong Kong SAR, China