改进梯度估算方法,让神经网络在模拟存算硬件上更稳定高效。
Extending Straight-Through Estimation for Robust Neural Networks on Analog CIM Hardware
- 分离前向噪声模拟与反向梯度计算,实现更精确的硬件噪声建模。
- 图像分类准确率提升5.3%,文本生成困惑度降低0.72,训练速度加快2.2倍。
- 适合部署在模拟存算硬件上的鲁棒神经网络训练,尤其关注能效与稳定性。
模拟存算(Analog Compute-In-Memory, CIM)架构在神经网络推理中具有显著能效优势,但受复杂硬件噪声影响,部署困难。现有噪声感知训练方法通常依赖理想化且可微的噪声模型,无法充分捕捉模拟CIM硬件的多样性。受量化中的直通估计器(Straight-Through Estimator, STE)启发,本文将前向噪声模拟与反向梯度计算解耦,支持在模拟CIM系统中使用更精确但计算不可行的噪声模型进行训练。理论分析表明,该方法保持了关键梯度方向信息,同时确保计算可行性与优化稳定性。大量实验显示,所提扩展STE框架在图像分类任务上实现最高5.3%的准确率提升,在文本生成任务中实现0.72的困惑度下降,训练速度提升2.2倍,峰值内存占用降低37.9%,优于标准噪声感知训练方法。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。