首个在ImageNet上精度超过80%的二值神经网络BNext问世,-1与+1的五年辛路历程

两年前,当 MeliusNet 问世时,机器之心曾发表过一篇技术文章《第一次胜过 MobileNet 的二值神经网络,-1 与 + 1 的三年艰苦跋涉​》,回顾了 BNN 的发展历程。彼时,依靠早期 BNN 工作 XNOR-Net 起家的 XNOR.AI 被苹果公司收购,大家曾畅想过这种低功耗、高效能的二值神经网络技术会不会即将开启广阔的应用前景。

然而,过去的两年时间,我们很难从对技术严格保密的苹果公司获得关于 BNN 技术应用的更多信息,而不论是学界和业界也未出现其他特别亮眼的应用案例。另一方面,随着终端设备数量的暴涨,边缘 AI 应用和市场正在迅速增长: 预计到 2030 年将产生 500 到 1250 亿台边缘设备,边缘计算市场将暴涨到 600 亿美元规模。这其中有几个目前热门的应用领域:AIoT、元宇宙和机器人终端设备。相关业界正在加速发力技术落地,与此同时 AI 能力已经嵌入到以上领域的诸多核心技术环节中,如 AI 技术在三维重建、视频压缩以及机器人实时感知场景中的广泛应用。在这样的背景下,业界对基于边缘的高能效、低功耗 AI 技术、软件工具以及硬件加速的需求变的日益迫切。

目前,制约 BNN 应用的瓶颈主要有两方面:首先,无法有效缩小和传统 32-bit 深度学习模型的精度差距;第二则是缺乏在不同硬件上的高性能算法实现。机器学习论文上的加速比通常无法体现在你正在使用的 GPU 或 CPU 上。第二个原因的产生可能正是来自于第一个原因,BNN 无法达到令人满意的精度,因此无法吸引来自系统和硬件加速、优化领域的从业者的广泛关注。而机器学习算法社区通常无法自己开发高性能的硬件代码。因此,想要同时获得高精度和强加速效果,BNN 应用或加速器无疑需要来自这两个不同领域的开发者进行协作。

BNN 为什么对于计算和内存高效

举个例子,Meta 推荐系统模型 DLRM 使用 32-bit 浮点数来储存权重和激活参数,它的模型大小约为 2.2GB。而一个少量精度下降 (<4%) 的二值版本模型的大小仅为 67.5MB,由此可见二值神经网络可以节省 32 倍的模型大小。这对于储存受限的边缘设备来说具备了非常强的优势。此外,BNN 也非常适合应用在分布式 AI 场景中,例如联邦学习常常对模型权重进行传输和聚合操作,因此模型大小和带宽受限往往成为整个系统的瓶颈。常用的手段是通过增加模型本地优化迭代次数来换取权重传输频率的降低,用这种折中的方法来提升整体效率,如著名的谷歌 FedAvg 算法。但增大本地计算量的策略对端侧用户并不友好,这里,如果采用 BNN 就可以很轻易的降低几十倍的数据传输量。

BNN 的第二个显著的优势是计算方式极度高效。它只用 1bit 也就是两种 state 来表示变量。这意味着所有运算只靠位运算就能完成,借助与门、异或门等运算,可以替代传统乘加运算。位运算是电路中的基础单元,熟悉电路设计的同学应该明白,有效减小乘加计算单元的面积以及减少片外访存是降低功耗的最有效手段,而 BNN 从内存和计算两个方面都具备得天独厚的优势,WRPN[1]展示了在定制化 FPGA 和 ASIC 上,BNN 对比全精度可获得 1000 倍的功耗节省。更近期的工作 BoolNet [2]展示了一种 BNN 结构设计可以几乎不使用浮点运算并保持纯 binary 的信息流,它在 ASIC 仿真中获得极佳的功耗、精度权衡。

第一个精度上 80% 的 BNN 是什么样子的?

来自德国 Hasso Plattner 计算机系统工程研究院的 Nianhui Guo 和 Haojin Yang 等研究者提出了 BNext 模型,成为第一个在 ImageNet 数据集上 top1 分类准确率突破 80% 的 BNN:

首个在ImageNet上精度超过80%的二值神经网络BNext问世,-1与+1的五年辛路历程

图 1 基于 ImageNet 的 SOTA BNN 性能对比

首个在ImageNet上精度超过80%的二值神经网络BNext问世,-1与+1的五年辛路历程

论文地址:https://arxiv.org/pdf/2211.12933.pdf

作者首先基于 Loss Landscape 可视化的形式深入对比了当前主流 BNN 模型同 32-bit 模型在优化友好度方面的巨大差异 (图 2),提出 BNN 的粗糙 Loss Landscape 是阻碍当前研究社区进一步探索 BNN 性能边界的的主要原因之一。

基于这一假设,作者尝试利用新颖的结构设计对 BNN 模型优化友好度进行提升,通过构造具备更平滑 Loss Landscape 的二值神经网络架构以降低对高精度 BNN 模型的优化难度。具体而言,作者强调模型二值化大幅度限制了可用于前向传播的特征模式,促使二值卷积仅能在有限的特征空间进行信息提取与处理,而这种受限前馈传播模式带来的优化困难可以通过两个层面的结构设计得到有效缓解:(1) 构造灵活的邻接卷积特征校准模块以提高模型对二值表征的适应性;(2) 探索高效的旁路结构以缓解前馈传播中由于特征二值化带来的信息瓶颈问题。

首个在ImageNet上精度超过80%的二值神经网络BNext问世,-1与+1的五年辛路历程

图 2 针对流行 BNN 架构的 Loss Landscape 可视化对比(2D 等高线视角)

基于上述分析,作者提出了 BNext,首个在 ImageNe 图像分类任务达到 > 80% 准确率的二值神经网络架构,具体的网络架构设计如图 4 所示。作者首先设计了基于 Info-Recoupling (Info-RCP) 模块的基础二值处理单元。针对邻接卷积间的信息瓶颈问题,通过引入额外的 Batch Normalization 层与 PReLU 层完成对二值卷积输出分布的初步校准设计。接着作者构造了基于逆向残差结构与 Squeeze-And-Expand 分支结构的二次动态分布校准设计。如图 3 所示,相比传统的 Real2Binary 校准结构,额外的逆向残差结构充分考虑了二值单元输入与输出间的特征差距,避免了完全基于输入信息的次优分布校准。这种双阶段的动态分布校准可以有效降低后续邻接二值卷积层的特征提取难度。

首个在ImageNet上精度超过80%的二值神经网络BNext问世,-1与+1的五年辛路历程

图 3 卷积模块设计对比图

其次,作者提出结合 Element-wise Attention (ELM-Attention) 的增强二值 Basic Block 模块。作者通过堆叠多个 Info-RCP 模块完成 Basic Block 的基础搭建,并对每个 Info-RCP 模块引入额外的 Batch Normalization 和连续的残差连接以进一步缓解不同 Info-RCP 模块间的信息瓶颈问题。基于旁路结构对二值模型优化影响的分析, 作者提出使用 Element-wise 矩阵乘法分支对每个 Basic Block 的首个 3×3 Info-RCP 模块输出进行分布校准。额外的空域注意力加权机制可以帮助 Basic Block 以更灵活的机制进行前向信息融合与分发,改善模型 Loss Landscape 平滑度。如图 2.e 和图 2.f 所示,所提出的模块设计可以显著改善模型 Loss Landscape 平滑度。

首个在ImageNet上精度超过80%的二值神经网络BNext问世,-1与+1的五年辛路历程

图 4 BNext 架构设计。”Processor 代表 Info-RCP 模块,”BN “代表 Batch Normalization 层,”C “表示模型基本宽度,”N “和”M ” 表示模型不同阶段的深度比例参数。

首个在ImageNet上精度超过80%的二值神经网络BNext问世,-1与+1的五年辛路历程

Table 1 BNext 系列。“Q”表示输入层、SEbranch以及输出层量化设置。

作者将以上结构设计同流行的MoboleNetv1基准模型结合,并通过改变模型深度与宽度的比例系数构建了四种不同复杂度的 BNext 模型系列 (Table 1):BNex-Tiny,BNext-Small,BNext-Middle,BNext-Large。

由于相对粗糙的 Loss Landscape,当前二值模型优化普遍依赖于 knowledge distillation 等方法提供的更精细监督信息,以此摆脱广泛存在的次优收敛。BNext作者首次考虑了优化过程中teache模型与二值student模型预测分布巨大差距可能带来的影响,,指出单纯基于模型精度进行的 teacher 选择会带来反直觉的 student 过拟合结果。为解决这一问题,作者提出了 knowledge-complexity (KC) 作为新的 teacher-selection 度量标准,同时考虑 teacher 模型的输出软标签有效性与 teacher 模型参数复杂度之间的关联性。

首个在ImageNet上精度超过80%的二值神经网络BNext问世,-1与+1的五年辛路历程

如图 5 所示,基于 knowledge complexity,作者对流行全精度模型系列如 ResNet、EfficientNet、ConvNext 进行了复杂度度量与排序,结合 BNext-T 作为 student 模型初步验证了该度量标准的有效性,并基于排序结果用于后续实验中的 knowledge distillation 模型选择。

首个在ImageNet上精度超过80%的二值神经网络BNext问世,-1与+1的五年辛路历程

图 5 反直觉的过拟合效应和不同教师选择下的知识复杂性影响 

在此基础上,论文作者进一步考虑了强 teacher 优化过程中由于早期预测分布差距造成的优化问题,提出 Diversified Consecutive KD。如下所示,作者通过强弱 teachers 组合的知识集成方法对优化过程中的目标函数进行调制。在此基础上,进一步引入 knowledge-boosting 策略,利用多个预定义候选 teachers 在训练过程中对弱 teacher 进行均匀切换,将组合知识复杂度按照从弱到强的顺序进行课程式引导,降低预测分布差异性带来的优化干扰。

首个在ImageNet上精度超过80%的二值神经网络BNext问世,-1与+1的五年辛路历程

在优化技巧方面,BNext 作者充分考虑了现代高精度模型优化中数据增强可能带来的增益,并提供了首个针对现有流行数据增强策略在二值模型优化中可能带来影响的分析结果,实验结果表明,现有数据增强方法并不完全适用于二值模型优化,这为后续研究中特定于二值模型优化的数据增强策略设计提供了思路。

基于所提出架构设计与优化方法,作者在大规模图像分类任务 ImageNet-1k 进行方法验证。实验结果如图 6 所示。

首个在ImageNet上精度超过80%的二值神经网络BNext问世,-1与+1的五年辛路历程

图 6 基于 ImageNet-1k 的 SOTA BNN 方法比较。

相比于现有方法,BNext-L 在 ImageNet-1k 上首次将二值模型的性能边界推动至 80.57%,对大多数现有方法实现了 10%+ 的精度超越。相比于来自 Google 的 PokeBNN, BNext-M 在相近参数量前提下要高出 0.7%,作者同时强调,PokeBNN 的优化依赖于更高的计算资源,如高达 8192 的 Bacth Size 以及 720 个 Epoch 的 TPU 计算优化,而 BNext-L 仅仅以常规 Batch Size 512 迭代了 512 个 Epoch,这反映了 BNext 结构设计与优化方法的有效性。在基于相同基准模型的比较中,BNext-T 与 BNext-18 都有着大幅度的精度提升。在同全精度模型如 RegNetY-4G (80.0%)等的对比中,BNext-L 在展现相匹配的视觉表征学习能力同时,仅仅使用了有限的参数空间与计算复杂度,这为在边缘端部署基于二值模型特征提取器的下游视觉任务模型提供了丰富想象空间。

What next? 

BNext 作者在论文中提到的,他们和合作者们正积极在 GPU 硬件上实现并验证这个高精度 BNN 架构的运行效率,未来计划扩展到其他更广泛的硬件平台上。然而在编者看来,让社区对 BNN 重拾信心,被更多系统和硬件领域的极客关注到,也许这个工作更重要的意义在于重塑了 BNN 应用潜力的想象空间。从长远来看,随着越来越多的应用从以云为中心的计算范式向去中心化的边缘计算迁移,未来海量的边缘设备需要更加高效的 AI 技术、软件框架和硬件计算平台。而目前最主流的 AI 模型和计算架构都不是为边缘场景设计、优化的。因此,在找到边缘 AI 的答案之前,相信 BNN 始终都会是一个充满技术挑战又蕴涵巨大潜力的重要选项。

© 版权声明

相关文章