[论文解读] Scale-Dropout: Estimating Uncertainty in Deep Neural Networks Using Stochastic Scale

Soyed Tuhin Ahmed, Kamal Danouchi+4 more|arXiv (Cornell University)|Nov 27, 2023

Adversarial Robustness in Machine Learning被引用 5

一句话总结

本文提出了一种新型正则化技术 Scale-Dropout，用于二值神经网络（BNNs），该技术通过随机丢弃整个缩放向量而非单个神经元，实现仅需一个随机单元的高效蒙特卡洛推理，且与模型大小无关。该方法在自旋电子学存算一体（CIM）架构上实现了高达100倍的能效提升，并在不确定性估计方面实现高达100%的分布外（OOD）检测能力，且预测准确率比当前最先进（SOTA）的BNN高出1.33%。

ABSTRACT

Uncertainty estimation in Neural Networks (NNs) is vital in improving reliability and confidence in predictions, particularly in safety-critical applications. Bayesian Neural Networks (BayNNs) with Dropout as an approximation offer a systematic approach to quantifying uncertainty, but they inherently suffer from high hardware overhead in terms of power, memory, and computation. Thus, the applicability of BayNNs to edge devices with limited resources or to high-performance applications is challenging. Some of the inherent costs of BayNNs can be reduced by accelerating them in hardware on a Computation-In-Memory (CIM) architecture with spintronic memories and binarizing their parameters. However, numerous stochastic units are required to implement conventional dropout-based BayNN. In this paper, we propose the Scale Dropout, a novel regularization technique for Binary Neural Networks (BNNs), and Monte Carlo-Scale Dropout (MC-Scale Dropout)-based BayNNs for efficient uncertainty estimation. Our approach requires only one stochastic unit for the entire model, irrespective of the model size, leading to a highly scalable Bayesian NN. Furthermore, we introduce a novel Spintronic memory-based CIM architecture for the proposed BayNN that achieves more than $100 imes$ energy savings compared to the state-of-the-art. We validated our method to show up to a $1\%$ improvement in predictive performance and superior uncertainty estimates compared to related works.

研究动机与目标

• 解决边缘设备中贝叶斯神经网络（BayNNs）因过多随机单元和内存访问导致的高硬件开销问题。
• 将BayNNs中每神经元的随机单元减少为每层的缩放向量，实现BNN中可扩展的不确定性估计。
• 设计一种基于自旋电子学存储器的存算一体（CIM）架构，通过最小化外围电路改动支持随机缩放，实现硬件高效性。
• 在不增加模型参数或计算复杂度的前提下，提升BNN中的不确定性量化与预测性能。
• 利用基于SOT-MRAM的CIM实现资源受限边缘设备上能效高效的不确定性感知BNN部署。

提出的方法

• 提出Scale-Dropout，一种正则化方法，以概率p随机丢弃BNN层的整个缩放向量，从而在训练过程中减少共适应现象。
• 提出蒙特卡洛缩放丢弃（MC-Scale Dropout），一种贝叶斯推理方法，利用随机缩放向量通过蒙特卡洛采样估计预测不确定性。
• 设计基于SOT-MRAM的存算一体（CIM）架构，通过外围电路注入随机性，避免对存储阵列的修改。
• 利用SOT-MRAM的高电阻动态范围（高达数MΩ）支持大规模交叉条阵列中可扩展、能效高效的贝叶斯推理。
• 将缩放向量作为BNN中的可学习参数，以减少量化误差，并在推理过程中实现稳定的随机缩放。
• 在基于交叉条的CIM架构上实现该方法，其中矩阵-向量乘法在内存中完成，最大限度减少数据移动与能耗。

Figure 1: Several nodes (neurons) a) at training time that are scaled with a probability of $p$ and dropped (ignored) with a probability of $1-p$ , b) At test time, if point estimate prediction is preferred, all the nodes are always scaled. However, for Bayesian inference, all nodes behave similarly

实验结果

研究问题

• RQ1 是否可以用每个模型仅一个随机单元替代贝叶斯BNN中每神经元的丢弃单元，同时保持不确定性估计质量？
• RQ2 与传统MC-Dropout相比，对整个缩放向量进行随机缩放在BNN中的不确定性量化与预测性能方面有何影响？
• RQ3 基于自旋电子学存储器的CIM架构，仅通过最小外围修改，能在多大程度上实现贝叶斯BNN的能效推理？
• RQ4 器件非均匀性与存储精度对所提出的Scale-Dropout方法的鲁棒性与可扩展性有何影响？
• RQ5 与当前最先进BNN和BayNN相比，该方法是否能在分布外（OOD）检测与预测准确率方面实现更优表现？

主要发现

• 所提出的MC-Scale Dropout方法在CIFAR-10上实现了高达100%的分布外（OOD）检测率，显著优于先前方法。
• 在VGG结构上，即使在亮度退化导致准确率下降6.89%的情况下，该方法仍能保持29.53%的最差情况OOD检测率。
• 与当前最先进二值贝叶斯神经网络相比，该方法在预测性能上最高提升1.33%。
• 该方法将传统BNN的准确率最高提升0.26%，表明通过不确定性感知训练可实现更好的泛化能力。
• 基于SOT-MRAM的CIM架构相比当前最先进方案实现了超过100倍的能效提升，仅需外围电路修改。
• 实证分析表明，使用1000次蒙特卡洛采样后，各类别的后验分布趋近于高斯分布，验证了贝叶斯近似的合理性。

Figure 2: Spin Scale-Dropout Module based on SOT MTJ.