Extending Straight-Through Estimation for Robust Neural Networks on Analog CIM Hardware

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Feng, Yuannuo, Zhou, Wenyong, Lyu, Yuexi, Zhang, Yixiang, Liu, Zhengwu, Wong, Ngai, Kang, Wang
Format: Preprint
Published: 2025
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866908491780718592
author Feng, Yuannuo
Zhou, Wenyong
Lyu, Yuexi
Zhang, Yixiang
Liu, Zhengwu
Wong, Ngai
Kang, Wang
author_facet Feng, Yuannuo
Zhou, Wenyong
Lyu, Yuexi
Zhang, Yixiang
Liu, Zhengwu
Wong, Ngai
Kang, Wang
contents Analog Compute-In-Memory (CIM) architectures promise significant energy efficiency gains for neural network inference, but suffer from complex hardware-induced noise that poses major challenges for deployment. While noise-aware training methods have been proposed to address this issue, they typically rely on idealized and differentiable noise models that fail to capture the full complexity of analog CIM hardware variations. Motivated by the Straight-Through Estimator (STE) framework in quantization, we decouple forward noise simulation from backward gradient computation, enabling noise-aware training with more accurate but computationally intractable noise modeling in analog CIM systems. We provide theoretical analysis demonstrating that our approach preserves essential gradient directional information while maintaining computational tractability and optimization stability. Extensive experiments show that our extended STE framework achieves up to 5.3% accuracy improvement on image classification, 0.72 perplexity reduction on text generation, 2.2$\times$ speedup in training time, and 37.9% lower peak memory usage compared to standard noise-aware training methods.
format Preprint
id arxiv_https___arxiv_org_abs_2508_11940
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle Extending Straight-Through Estimation for Robust Neural Networks on Analog CIM Hardware
Feng, Yuannuo
Zhou, Wenyong
Lyu, Yuexi
Zhang, Yixiang
Liu, Zhengwu
Wong, Ngai
Kang, Wang
Machine Learning
Artificial Intelligence
Hardware Architecture
Analog Compute-In-Memory (CIM) architectures promise significant energy efficiency gains for neural network inference, but suffer from complex hardware-induced noise that poses major challenges for deployment. While noise-aware training methods have been proposed to address this issue, they typically rely on idealized and differentiable noise models that fail to capture the full complexity of analog CIM hardware variations. Motivated by the Straight-Through Estimator (STE) framework in quantization, we decouple forward noise simulation from backward gradient computation, enabling noise-aware training with more accurate but computationally intractable noise modeling in analog CIM systems. We provide theoretical analysis demonstrating that our approach preserves essential gradient directional information while maintaining computational tractability and optimization stability. Extensive experiments show that our extended STE framework achieves up to 5.3% accuracy improvement on image classification, 0.72 perplexity reduction on text generation, 2.2$\times$ speedup in training time, and 37.9% lower peak memory usage compared to standard noise-aware training methods.
title Extending Straight-Through Estimation for Robust Neural Networks on Analog CIM Hardware
topic Machine Learning
Artificial Intelligence
Hardware Architecture
url https://arxiv.org/abs/2508.11940