PokeBNN：二值神经网络的质量改进和成本优化

4556065707580 124750PokeBNN：轻量级准确性的二值追求0Yichi Zhang �0康奈尔大学0yz2499@cornell.edu0Zhiru Zhang康奈尔大学0zhiruz@cornell.edu0Łukasz Lew �0Google研究0lew@google.com0摘要0优化Top-1ImageNet会促使出现巨大的网络，这在推理环境中可能是不切实际的。二值神经网络（BNNs）有潜力显著降低计算强度，但现有模型的质量较低。为了克服这个缺陷，我们提出了PokeConv，一种二值卷积块，通过添加多个残差路径和调整激活函数等技术来改善BNNs的质量。我们将其应用于ResNet-50并优化ResNet的初始卷积层，这是难以二值化的。我们将结果网络系列命名为PokeBNN1。选择这些技术是为了在top-1准确率和网络成本两方面都获得有利的改进。为了能够联合优化成本和准确率，我们定义了算术计算工作量（ACE），这是一种针对量化和二值化网络的硬件和能量启发式成本指标。我们还确定了一个优化未充分探索的超参数，用于控制二值化梯度近似。我们在top-1准确率以及常用的CPU64成本、ACE成本和网络大小指标上建立了新的最先进（SOTA）。BNNs中的ReActNet-Adam[33]在70.5%的top-1准确率和7.9的ACE上实现了SOTA。PokeBNN的一个小变体在2.6的ACE上实现了70.5%的top-1准确率，成本降低了3倍以上；一个较大的PokeBNN在7.8的ACE上实现了75.6%的top-1准确率，准确率提高了5%以上而不增加成本。PokeBNN在JAX/Flax [6, 18]中的实现和复现说明已经开源。01. 引言0对帕累托优化的需求。深度学习研究在很大程度上受到基准测试和指标的驱动。0� 在Google期间进行的工作，贡献相等。1 Poke /'p6kI/的发音类似于pocket。PokeConv，PokeBNN和Pokemon分别是PocketConvolution，Pocket Binary Neural Network和PocketMonster的缩写。2源代码和复现说明可在AQT存储库中找到：github.com/google/aqt。0ACE（百万）0Top-1（%）0XNOR++0BiReal180BiReal340IR180IR340SQBWN180PCNN0MeliusNet290Melius0MeliusNet590RealtoBin0ReActNet-Adam0ReActNet02.0x01.75x01.5x01.4x01.25x01.0x00.75x00.5x（PokeBNN）0图1. 使用top-1和ACE（第3节）比较BNNs。0过去，每个基准测试中的单个指标，例如ImageNet上的top-1准确率，就足够了。但是现在，我们需要考虑各种模型架构、大小和计算成本。这促使我们优化一个质量指标（如top-1准确率）和另一个成本指标（如FLOPS、延迟或能耗）的帕累托前沿。优化指标的选择。由于目前行业是大规模推理的主要用户，成本指标应与每次推理的美元成本相关。随着机器学习硬件的成熟，越来越明显的是能耗是与推理成本成正比的关键指标，尤其是在数据中心。在第3节中，我们定义了一个称为算术计算工作量（ACE）的新的代理指标，旨在估计抽象出具体机器学习硬件的推理成本。量化和二值化的影响。在int8或float16等数值格式中，较不重要的位对网络的输出影响不如较重要的位大。然而，处理它们消耗的能量是相同的。这可能不是最优的。每减半一次量化位数（例如从16位到8位到4位到2位到1位），可能的推理成本降低（或等效的性能改进）至少为2倍（例如，NVIDIAAmpere在int1上比int8快8倍[38]），根据ACE指标的估计，最多可达到4倍。相比之下，这比将GPU或TPU升级一代或两代所带来的改进要大得多。124760二值化通过用逻辑XNOR和位计数操作替换浮点乘法运算将这种优势推向极致。如果二值神经网络（BNNs）能够达到高质量，它们很可能在数据中心和边缘推理中占据重要地位。BNN优化很困难。现代BNN的先驱曾与其浮点对应物相比，top-1准确率差距超过20%[23]。直到最近，BNNs才在质量上与流行的ResNet-18模型相媲美[33,34]。其中一个原因是BNNs往往具有混乱、不连续的损失曲面，这使得它们的优化具有挑战性[31,33]。事实上，为了使二值化起作用，与标准DNN实践相比，必须改变许多事物。BNNs需要多阶段训练、梯度的近似以及各种架构调整，以避免二值化信息瓶颈。我们的主要贡献如下：•我们提出PokeConv，一种二值卷积块，可以显著提高BNN的准确性。我们用PokeConv替换了ResNet[17]中的大部分卷积。0•我们提出PokeInit块来替换ResNet的初始卷积层，该层很难进行二值化。PokeInit显著降低了网络的成本。PokeInit和PokeConv构成了PokeBNN系列的基础。0•我们优化了BNN中一个未充分研究的剪裁边界超参数，该超参数控制二值化梯度的近似。第6节的消融实验证明，通过这个参数，我们在top-1准确率上获得了超过3%的提升。0•我们提出并定义了一种新颖的硬件和能量启发式成本度量标准ACE，它基于硬件上的推理成本，同时对现有硬件平台不加偏见。ACE改进了能源高效神经网络研究和ML硬件研究的一致性。我们使用ACE来量化PokeBNN的推理成本。0• 我们通过实验证明，在ImageNet[43]上，PokeBNN与CPU64、ACE和网络大小等成本度量一起建立了top-1的Pareto-SOTA。我们在相同的ACE成本下，比SOTAReActNet-Adam提高了5.1%的top-1准确率（图1）。02. 相关工作0有大量积极的研究致力于BNN的训练和加速。我们只回顾了对本文中的网络设计产生高影响的过去努力的一个子集。[49]中可以找到一份全面的调查报告。BNN的可行性。开创性的工作[9, 23,27]证明了BNN的可行性。他们建立了具有二值化权重和激活的神经网络的训练框架，并展示了有希望的结果。0在MINIST和CIFAR-10等小数据集上，这些初步的ImageNet结果显示AlexNet[29]的top-1准确率从62.5%下降到36.1%，GoogleNet[45]的top-1准确率从68.9%下降到47.1%。多阶段训练。多阶段训练[8, 33, 34,37]是一种关键的有效技术，其中首先训练一个未量化的模型，然后再启用二值化。一些方法采用三阶段训练——从未量化版本到仅二值化激活，再到二值化权重和激活[37]。知识蒸馏是另一种常用的技术，用于提高BNNs的准确性[8,33, 34,37]。BNN架构。另一条全面的工作线探索架构变化，以追求更好的模型质量。其中许多目标是产生可忽略的计算和参数开销。例如，对二值化张量进行通道级实值重新缩放可以有效减轻量化损失[2, 7,42]。将二值化卷积层的未量化输入激活与其输出通过快捷方式连接可以增强梯度流动和模型表示能力[35]。Squeeze-and-excitation（SE）[22]是另一种计算廉价的技术，可以改善包括BNNs在内的小型卷积模型的质量[37]。FracBNN[50]在BNN中包含额外的BatchNorm层[24]以加快收敛速度。[8]中的作者首先表明，在每个卷积层之后使用PReLU函数[16]可以提高二值模型质量。沿着这条线，最近有报道称，在PReLU函数中引入可学习的偏置可以进一步提高模型准确性[33,34]。随着发展，当前的BNNs最终超过了70%的top-1准确率。03. 算术计算工作量0在本节中，我们阐述和定义了ACE，它旨在反映使用CMOS方法实现的理想化ML硬件上的神经网络推理成本。ACE度量定义如下：0ACE = 0i ∈ I,j ∈ J n i,j ∙ i ∙ j (1)0其中n i,j是a i -bit数和b j-bit数之间的乘累加（MAC）操作的数量，可以从模型结构自动推导出来。I和J是在给定神经网络推理中使用的所有位宽的集合，通常I = J = {1, 2, 4, 8,16}。能源使用与计算的总成本高度相关。推理可能发生在数据中心或边缘设备上，并且可以由CPU、GPU或TPU提供服务。对于边缘设备，电池使用是主要关注点，这使得能源使用成为许多ML应用的关键瓶颈。在数据中心的情况下，令人惊讶的是，能源也是主要的成本驱动因素。为了在数据中心中运行推理，需要支付以下费用：124770硬件、电力和电源供应以及其他基础设施成本。一张GPU卡可能需要1000美元，并且在3-5年内使用，消耗400W。以65%的利用率和每千瓦时15美分的电价计算，三年的电费将达到0.4千瓦*24小时*365*3*0.65*0.15美元/千瓦时=约1000美元。有趣的是，数据中心的电源供应成本（冷却、变压器、电池、备用发电机）据报道是电费的两倍以上（至少在Google数据中心的情况下）[25]。此外，ML芯片成本与其TDP的相关性超过90%。总体而言，运行推理的成本主要由能源消耗驱动。大部分计算能源使用在算术运算能源上。与经典CPU相反，运行推理的ML硬件在实际算术运算（如乘法、加法、其他函数）上消耗了很高比例的能源。例如，在TPUs的情况下，计算控制的成本在巨大的SIMD大小（16K到64K）上摊销[25,26]。这通常使用系统阵列[30]实现。相比之下，CPU的典型SIMD大小为4到32（例如，SSE，AVX）。我们在论文的完整版本的附录中讨论了其他非算术能源消耗。算术运算能源与活动位加法器的数量成正比。要将两个无符号整数a <2^I，b <2^J相乘，首先使用逻辑AND操作计算I∙J位的值，并将它们分组求和：�00 ≤ i x.ch: r = avg_ch(r, x.ch) if r.shape != x.shape:0r = avg_pool_3x3(r, x.ch, st = 2)0return x + r0使用适合二值化的非线性函数。尽管二值化函数本身就是非线性的，但插入额外的非线性函数可以进一步提高BNN模型的质量，正如先前的研究[8,34]所报道的那样，特别是如果该函数学习如何移动和重塑激活分布[34]。我们建议在残差相加之后添加一个非线性函数——Dynamic PReLU（DPReLU）[36]，定义如下：0DPReLU(x) := η(x - α) - βx - α > 0 γ(x - α) - βotherwise (4)0这里α、β、γ和η都是逐通道可学习的参数。它们的初始值分别为0、0、0.25和1.0。除了激活移位外，DPReLU在两个线性片段上都具有可学习的斜率。与ReActNet[34]提出的RPReLU相比，它引入了额外的重塑灵活性。Squeeze-and-excitation（SE）[22]有助于缩放。SE是一种计算成本低廉的技术，可以提高模型质量。它允许网络在给定输入上融入全局知识。BNN（如real-to-binary[37]）也使用了SE的不同变体。在我们的设计中，我们应用了一个SE块，就像在MobileNetV3[20]中使用的那样。图3显示了其图示和伪代码。0空间均值04位线性0ReLU04位线性0硬Sigmoid01 x 1 x Cin01 x 1 x (Cin/8)01 x 1 x (Cin/8)01 x 1 x Cout0H x W x Cout0H x W x Cin0def SE_4b(x, ch)0s = SpatialMean(x) s =Linear_4b(s, x.ch / 8) s = ReLU(s) s= Linear_4b(s, ch) s = ReLU6(s +3) / 6 return s0图3. Squeeze-and-Excitation：量化为4位。0重要的是，ACE指标报告SEdense层会产生非可忽略的成本。因此，我们建议将SE块量化为4位，并设置隐藏投影0将长度缩小为输入的1/8。实验证明，这种修改不会损失准确性。额外的归一化层。我们在输出（即残差分割之前）上放置了一个额外的BatchNorm层[24]。FracBNN[50]建议这种修改以加快收敛速度。这个额外的层对于PokeBNN来说更加重要。首先，它的偏置项学习如何适当地移动分布，以平衡激活信号到下一个PokeConv层。其次，它对分布的调整允许对第一个SEdense层进行激进的4位量化。此外，它对二值卷积周围的快捷方式进行归一化，并促进梯度流动。限制。我们使用默认的零填充Conv层，这会为一小部分像素引入第三个值。我们相信大多数先前的工作都受到了相同的限制。FBNA[15]提出并评估了交替的1和-1填充，可以解决这个限制。04.3. PokeInit和投影层优化0将常规的Convs替换为PokeConvs后，我们发现ResNet-50中还有两个ACE代价较高的组件：（1）输入的7×7Conv层；（2）下采样块中用于快捷方式的1×1投影Conv层。这些层通常在二值化中被排除[32]。将1×1投影Conv层替换为ReshapeAdd。1×1投影Conv层将产生3.6亿个MACs。我们建议完全删除这些下采样投影层，并用为快捷方式定义的ReshapeAdd函数替换它们。我们使用平铺而不是零填充进行通道扩展，即对于具有n个通道的输入张量x，tile(x)i = x(i modn)。PokeInit。仅未量化的输入层就需要1.18亿个MACs。大量MACs的主要来源是（1）7×7的卷积核大小和（2）大的输出空间分辨率112×112。为了优化ACE成本，我们将步幅为2的最大池化与第一个步幅为2的卷积融合，得到步幅为4。这将使输出空间分辨率减小4倍。然后我们将卷积核大小从7×7减小到4×4。进一步减小卷积核将导致信息丢失，因为会有像素不与卷积核进行卷积。为了增加输入块的感受野，我们在其后跟随一个3×3的深度卷积层。我们将这样的输入层组合称为PokeInit。其伪代码如图4所示。为了进一步降低成本，我们将PokeInit量化为8位。这种优化将输入层的成本从1.18亿个浮点MACs降低到660万个int8 MACs。1248008位Conv4x40BatchNorm0DPReLU08位DWConv3x30BatchNorm0DPReLU0def PokeInit(x):0x = Conv4x4_8b(x, ch=32, st=4) x =BatchNorm(x) x = DPReLU(x) x =Conv3x3_depthwise_8b(x, ch=64) x =BatchNorm(x) x = DPReLU(x) return x0图4. PokeInit：量化为8位。04.4. 模型组装0我们使用ResNet-50模板将PokeConv和PokeInit组装在一起，其中线性分类器头部为8位，如伪代码所示。如前几节所讨论的，我们现在将PokeConv和PokeInit应用于基本的ResNet-50，并移除其中的1×1投影Conv层。除此之外，我们还将分类器量化为8位。最终网络的伪代码如下：0CH = 64 * M0def PokeBNN50(x):0x = PokeInit(x) for i inrange(16):0st = 2 if i in (3, 7, 13) else 1 if i < 3: ch = CH elif: i < 7: ch = CH* 2 elif: i < 13: ch = CH * 4 r = x x = PokeConv(x, None,Conv1x1_1b, ch, st=1) x = PokeConv(x, None, Conv3x3_1b, ch,st=st) x = PokeConv(x, r, Conv1x1_1b, 4 * ch, st=1)0x = SpatialMean(x) x =Linear(x) return x0需要注意的是，所有的Conv和线性层都不使用偏置，除非其后面没有BatchNorm层。我们按照ResNet[17]的配置来初始化BatchNorm。05. 实验0在本节中，我们在分辨率为224×224的ILSVRC12ImageNet[43]分类数据集上对PokeBNN进行了实证评估。我们只应用了随机裁剪和翻转作为数据增强。05.1. 训练设置0我们在64个TPU-v3芯片上进行实验，批量大小为8192。我们使用Adam优化器[28]（β1=0.9，β2=0.99），采用线性学习率衰减。初始学习率为6.4e-4。训练过程中的权重衰减设置为5e-5。BatchNorm的动量设置为0.9。为了估计非二值量化层激活的剪裁边界B，我们按照[1]中的方法进行估计，并0使用指数移动平均（α =0.9）来计算批次中的最大值。我们总共训练了750个epoch，并采用了两阶段训练。我们发现，第一个半非量化阶段只需要50个epoch。在第50个epoch时，4位和8位激活被量化。所有权重（8位、4位和二值化）在第50个epoch时被量化。在训练过程中，1位激活始终保持为1位。我们使用知识蒸馏的设置来训练PokeBNN，这需要计算KL散度损失。修改的方法很简单，只需用教师模型的预测替换独热编码的真实标签。我们使用一个8位的ResNet-50作为教师模型。我们还尝试了从视觉变换器[12]中进行蒸馏，但令人惊讶的是结果相似。为了测量准确性，在学习率衰减为零后，我们继续训练几个epoch。我们观察到由于批量归一化统计数据的进一步更新，top-1会出现震荡。我们发现训练和评估的top-1完全不相关。因此，按照取最大top-1的做法是不公平的。本文中的所有top-1数字都是在学习率已经为零的几个epoch上的平均top-1。平均准确性与最大准确性之间的差异约为0.5%至1%。05.2. 评估结果0为了进行公平比较，我们将PokeConv块中的输出通道数量进行缩放，以改变模型大小（例如，PokeBNN-2x表示将输出通道数量翻倍）。所有结果都在表2中收集。在不同随机权重下运行5次PokeBNN-1x时，top-1的标准差为0.034%。重要的是，对于之前的工作，我们假设所有的FP32操作都可以用BF16替代而不会损失精度。PokeBNN不使用FP32。基于我们分析的数据，我们比较了PokeBNN与文献中基线方法在准确性与成本指标之间的权衡，并绘制了准确性与成本指标的Pareto曲线。我们有几个关键观察结果：PokeBNN在ACE指标下建立了BNN的SOTAPareto前沿，如图1所示。PokeBNN的准确性随着模型大小的增加而显著提高。尽管从ResNet-50进行了二值化，但是在PokeBNN中将通道数量扩大2倍，top-1准确率达到了77.2%，略高于4位ResNet-50（表2），但效率提高了4.6倍以上。文献中的大多数BNN模型在ImageNet上的top-1准确率低于65%。ReActNet[34]和ReActNet-Adam[33]首次通过利用MobileNet架构[21]达到了接近70%的ResNet-18级准确性。在与当前SOTAReActNet-Adam相同的ACE预算下，我们的PokeBNN-1.4x达到了75.6%的top-1准确率，高出5%以上。一个小变体PokeBNN-0.75x有250556065707580XNORXNOR++BiReal18BiReal34IR18IR34SQBWN18PCNNBDense37CIBCNN18CIBCNN34MobiNetBinaryMobileNetMeliusNet29MeliusNet42MeliusNet59RealtoBinSABNN18SABNN34ReActNet-AdamReActNet2.0x1.75x1.5x1.4x1.25x1.0x0.75x0.5x (PokeBNN)0150556065707580XNOR++BiReal18BiReal34IR18IR34SQBWN18PCNNBDense37CIBCNN18CIBCNN34MobiNetMeliusNet29MeliusNet42MeliusNet59RealtoBinSABNN18SABNN34 ReActNet-AdamReActNet2.0x1.75x1.5x1.4x1.25x1.0x0.75x0.5x (PokeBNN)124810与ReActNet-Adam具有相同的top-1相比，ACE减少了3倍以上。与BNN文献中准确率最高的MeliusNet-59[4]变体相比（表2），大型变体PokeBNN-2x的准确率提高了6％，同时ACE效率提高了5.3倍。此外，我们在ResNet-18架构上测试了PokeConv，并观察到PokeBNN-0.5x优于它。PokeBNN还在常用的CPU64指标下建立了SOTAPareto前沿。我们使用文献中广泛采用的CPU64指标绘制Pareto曲线。如图5所示，权衡趋势与提出的ACE指标大致相同。值得注意的是，一些BNNs（例如MeliusNet-29）在CPU64指标下显示出比图1中不利的权衡。这是因为ACE捕捉到了二进制操作在能源使用方面比浮点操作便宜64倍以上的事实。0CPU64成本（百万）0Top-1（％）0图5. 使用top-1和CPU64成本进行BNN比较。0PokeBNN在尺寸-准确性权衡上优于先前的BNN。模型大小是指示内存需求的另一个重要维度。因此，我们在图6中绘制了模型大小与top-1准确性的关系，结果显示PokeBNN也处于SOTA Pareto前沿，与先前的方法相比。06. 剔除研究0在本节中，我们对我们提出的技术进行了详细的剔除研究。我们测量了每个单独技术对PokeBNN-1.0x的影响。剪切边界剔除。作为超参数，剪切边界B在低位宽量化[1,10]中起着重要作用，但在过去的BNN研究中很少被探索。在BNN中，虽然有一些工作手动设置了二进制激活的边界B= 1.3 [4,5]，但对其进行研究的工作很少，在大多数情况下，默认情况下B = 1 [23, 34, 35]。0模型大小（MB）0Top-1（％）0图6. 使用top-1和模型大小进行BNN比较。0在我们的实验中，我们发现B对BNN的准确性有显着影响。我们在一组值从1.0到6.0的范围内扫描B以进行二值化激活。每个PokeConv具有相同的边界。结果如表3所示。在B =3和最常用的值B =1之间存在3.3％的准确性差距。我们假设较大的剪切阈值可以改善损失曲面的Lipschitz常数，并减少死神经元的数量（即梯度为零的神经元）。这与先前观察到的[1,10]剪切边界对于超低位宽量化很重要的观察一致。PokeConv剔除。我们逐个删除PokeConv中的每个组件并研究其影响。结果如表4所示。删除PokeConv中的DPReLU导致最大的准确性下降，甚至大于用原始的二值化ResNet块替换PokeConv。我们假设这是因为该更改消除了两个快捷方式上的非线性，这阻碍了模型的学习。我们还将DPReLU替换为RPReLU[34]和ReLU。这分别导致了0.2％和0.6％的top-1降级。由于0.2％高于3个标准偏差（3 ×0.034％），因此DPReLU确实比其他两个候选者更改善了模型质量。其他组件（即4位SE，额外的快捷方式和BatchNorms）对模型的top-1至少有3％的影响。鉴于它们在ACE方面的开销可以忽略不计，它们是有利的设计选择。我们还尝试添加回1×1投影Conv层。top-1结果为73.5％，仅高出0.1％。因此，完全删除这些层是合理的。PokeInit剔除。在PokeBNN-1.0x中，我们将PokeInit替换为ResNet的原始7×7输入Conv层，然后是最大池化。top-1为73.5％，仅高出0.1％。鉴于PokeInit将输入层的操作数量减少了18倍，这是一个有利的权衡。我们还尝试删除PokeInit中的3×3深度卷积层。这导致73.1％的top-1。PokeBNN-1.75x0010.211.111037.111.9174.416.376.8PokeBNN-1.5x009.78.28111.78.9128.512.475.9PokeBNN-1.4x009.57.17037.27.8111.610.975.6PokeBNN-1.25x009.25.75635.86.389.69.075.0PokeBNN-1.0x008.73.63609.54.257.76.273.4PokeBNN-0.75x008.22.02032.72.632.93.870.5Top-1 (%)70.171.472.973.473.372.872.4Top-1 (%)70.660.468.170.261.9124820表2. 最终结果和与先前技术的比较 -在计算FP32操作的ACE时，我们假设它们可以转换为BF16而不会损失准确性。“-”表示数据不可用。PokeBNN-1.0x在5个不同种子上的top-1的标准差为0.034%。BF16PokeBNN是一个变种，其中所有卷积和全连接层都是BF16。底部四行显示了上下文的基础模型，所有其他模型都是二值化的。0模型MAC操作（10^6）ACE（10^9）CPU64（10^6）大小（MB）Top-1（%）FP32BF16INT8INT4二进制0GoogleNet-BNN [23] - - - - - - - - 47.10MobiNet [39] - - - - - - 52.0 4.6 54.40BinaryMobileNet [40] - - - - - - 154.0 - 60.90MeliusNet-59 [4] 245 - - - 18300 81.0 530.9 17.4 71.00Real-to-Binary Net [37] 156.4 - - - 1676 41.7 182.6 5.1 65.40QuickNetSmall [3] - - - - - - - 4.0 59.40QuickNetLarge [3] - - - - - - - 5.4 66.90ReActNet-A [34] 11.9 0 0 0 4816.9 7.9 87.2 7.4 69.40ReActNet-Adam [33] 11.9 0 0 0 4816.9 7.9 87.2 7.4 70.50INT4 ResNet-50 [1] 0 0 120.1 3969.1 0 71.2 263.1 13