[论文解读] NAX: Co-Designing Neural Network and Hardware Architecture for Memristive Xbar based Computing Systems

Shubham Negi, Indranil Chakraborty+2 more|arXiv (Cornell University)|Jun 23, 2021

Advanced Memory and Neural Computing参考文献 24被引用 7

一句话总结

NAX 是一种神经架构搜索框架，通过联合优化卷积核大小和交叉条形尺寸，协同设计深度神经网络与基于忆阻器交叉条形的存内计算（IMC）硬件，以在准确率与硬件效率之间取得平衡。与基线 ResNet 模型相比，其在 CIFAR-10 和 Tiny ImageNet 上分别实现了 17% 和 4% 的 EDAP 降低，以及 0.8% 和 0.2% 的准确率提升。

ABSTRACT

In-Memory Computing (IMC) hardware using Memristive Crossbar Arrays (MCAs) are gaining popularity to accelerate Deep Neural Networks (DNNs) since it alleviates the "memory wall" problem associated with von-Neumann architecture. The hardware efficiency (energy, latency and area) as well as application accuracy (considering device and circuit non-idealities) of DNNs mapped to such hardware are co-dependent on network parameters, such as kernel size, depth etc. and hardware architecture parameters such as crossbar size. However, co-optimization of both network and hardware parameters presents a challenging search space comprising of different kernel sizes mapped to varying crossbar sizes. To that effect, we propose NAX -- an efficient neural architecture search engine that co-designs neural network and IMC based hardware architecture. NAX explores the aforementioned search space to determine kernel and corresponding crossbar sizes for each DNN layer to achieve optimal tradeoffs between hardware efficiency and application accuracy. Our results from NAX show that the networks have heterogeneous crossbar sizes across different network layers, and achieves optimal hardware efficiency and accuracy considering the non-idealities in crossbars. On CIFAR-10 and Tiny ImageNet, our models achieve 0.8%, 0.2% higher accuracy, and 17%, 4% lower EDAP (energy-delay-area product) compared to a baseline ResNet-20 and ResNet-18 models, respectively.

研究动机与目标

• 解决基于忆阻器交叉条形的存内计算（IMC）系统中，深度神经网络（DNN）模型参数（如卷积核大小、深度）与硬件参数（如交叉条形尺寸）之间的相互依赖关系。
• 克服忆阻器交叉条形中器件与电路非理想性导致的硬件效率（能耗、延迟、面积）与模型准确率之间的权衡。
• 开发一种神经架构搜索（NAS）框架，协同优化神经网络架构与 IMC 硬件配置，以在非理想条件下实现最优性能。
• 证明在不同网络层中采用异构的交叉条形尺寸，可同时提升模拟 IMC 系统的硬件效率与推理准确率。

提出的方法

• 提出一种协同设计搜索空间，包含可变卷积核大小（如 3x3、5x5）映射到不同交叉条形尺寸（如 64x64、128x128）的 DNN 层。
• 在 NAS 过程中将忆阻器交叉条形中的非理想性（如寄生电阻、存取晶体管引起的非线性、I-V 非理想性）整合到准确率估计中。
• 采用类似 DARTS 的可微分架构搜索方法，构建包含多种卷积核与交叉条形尺寸选项的混合层超级网络。
• 应用两阶段搜索：i-search（忽略非理想性）与 ni-search（包含非理想性），以在真实硬件约束下评估准确率。
• 在架构参数更新过程中使用 GENIEx 推理来估计非理想性的影响，从而提升准确率感知的搜索能力。
• 使用 Adam 优化器配合学习率衰减策略训练架构参数，并在测试集上验证最终紧凑网络的性能。

实验结果

研究问题

• RQ1 交叉条形尺寸如何影响映射到忆阻器交叉条形 IMC 系统的 DNN 的硬件效率（EDAP）与准确率？
• RQ2 在考虑忆阻器器件与电路非理想性的情况下，协同优化 DNN 架构与交叉条形尺寸时，硬件效率与模型准确率之间的最优权衡是什么？
• RQ3 一种联合探索 DNN 与硬件参数的 NAS 框架，是否能在 EDAP 与准确率方面优于将 DNN 映射到固定硬件配置的传统方法？
• RQ4 较大交叉条形中的非理想性如何影响推理准确率？是否可通过按层智能选择交叉条形尺寸来缓解？
• RQ5 在搜索过程中考虑非理想性，能在多大程度上提升最终模型在真实 IMC 硬件上的准确率？

主要发现

• 与基线 ResNet-20 模型相比，NAX 在映射到忆阻器交叉条形硬件时，于 CIFAR-10 上实现了 17% 的 EDAP 降低与 0.8% 的准确率提升。
• 在 Tiny ImageNet 上，NAX 在相同硬件约束下，相比基线 ResNet-18 模型，实现了 4% 的 EDAP 降低与 0.2% 的准确率提升。
• 最优交叉条形尺寸因层而异：当输出通道数超过 128 时，使用更大的交叉条形（如 128x128）以避免资源利用率低下；而对较小操作则更偏好使用较小的交叉条形（如 64x64）。
• 考虑非理想性的 ni-search 变体在搜索过程中显著提升了模型准确率，尤其在 Tiny ImageNet 上，最佳模型的准确率差距达到 18.57%。
• ni-search 的搜索成本更高（在 Tiny ImageNet 上为 350 GPU 小时），高于 i-search（5 GPU 小时），但所得模型在准确率与硬件效率方面表现更优。
• 结果表明，按层采用异构交叉条形尺寸对于平衡硬件效率与准确率至关重要，其中 128x128 交叉条形适用于高通道数操作，而 64x64 适用于小规模操作。