Bi-MTDP：通过二值网络加速多任务密集预测，又快又提点 | CVPR 2024

论文提出二值化多任务密集预测器 Bi-MTDP，通过二值神经网络（BNNs）显著加速多任务密集预测模型，同时保持甚至提高模型性能。为了避免信息严重退化而导致二值化带来性能下降，论文引入了深度信息瓶颈层，在前向传播时强制要求下游任务表示满足高斯分布；此外，还引入知识蒸馏机制来纠正反向传播中信息流方向

来源：晓飞的算法工程笔记 公众号

论文: Efficient Multitask Dense Predictor via Binarization

• 论文地址：https://arxiv.org/abs/2405.14136
• 论文代码：https://github.com/42Shawn/BiMTDP

Introduction

---

  在计算机视觉领域，有一个日益增长的趋势，即以多任务学习的方式处理密集预测任务，例如语义分割、单眼深度估计和人体解析。由于多任务模型中通过交叉连接结构进行信息补充机制的收益，一系列密集预测任务的整体性能得到了极大提升（见图1）。然而，处理多个复杂像素级任务的当前最先进（SoTA）多任务密集预测模型具有相当大的计算需求，这种高计算要求限制了它们在资源受限环境中（如自动驾驶、机器人技术和虚拟现实等领域）的应用。论文旨在优化这些庞大的SoTA模型以适用于边缘设备，在速度和性能之间取得平衡。

  已经不是研究探索神经网络压缩策略，包括剪枝、网络量化和知识蒸馏。值得注意的是，网络二值化作为一种量化形式，将权重和激活最小化为 $\pm 1$ ，从而能够用二值神经网络（BNNs）中更高效的XNOR-Bitcount操作替代全精度网络中计算密集型的内积运算。理论上，二值化可以将存储成本降低32倍，并将推断速度提高64倍，使得BNNs适用于边缘设备。

  虽然二值神经网络（BNNs）在图像分类方面取得了令人印象深刻的成果，几乎达到了全精度ResNet级别的准确性，但其应用主要局限于小规模模型，无法用于其他计算密集型的计算机视觉任务。将BNNs扩展到更大的模型应该是下一步的研究，但这种扩展受到过拟合和信息降级等问题的阻碍。在全精度模型中有效的技术，如标签平滑、dropout和mixup，在BNNs上效果较差。此外，最先进的多任务密集预测任务通常需要深层复杂模型并配备多模态融合结构，加剧了有效实施二值化所面临的挑战。

  在应用二值化进行多任务密集预测任务时，主要障碍是在深度模型中信息流的显著降级，导致性能下降。为了解决这个问题，论文提出了一个Binary Multitask Dense Predictor（Bi-MTDP）基线，通过二值化模块构建了一个多任务密集预测器。基于对这个基线的彻底评估，论文得出结论：二值化操作破坏了多任务模型中的信息流，因此与其全精度对应物相比，下游任务的表示不够丰富。为解决这一问题，论文通过两个方向上额外的信息流校准机制更新Bi-MTDP

1. 在正向传播中实施变分信息瓶颈机制，使嵌入遵循具有稀疏性质的高斯分布以过滤掉与任务无关的因素。
2. 在反向传播中利用现有全精度模型通过基于特征的知识蒸馏来校准二值网络梯度。

  Bi-MTDP的好处可以从两个正交的角度进行分析：

1. 从网络二值化的角度来看，通过将二值化与多任务密集预测框架相结合，Bi-MTDP的成功证明了它可以有效地补充信息，从而提高了各个二值模型的性能。
2. 从多任务密集预测任务的角度来看，加速这些繁琐的模型有利于在其中设计更有效和高效的交叉模块，如图1所示。由于现有的密集预测模型在建模交叉模块方面存在严重限制，因为它们过度依赖卷积操作，所以对于多任务密集预测来说，通过交叉机制学习涵盖多任务上下文各个范围的交互和推理是至关重要的。

Multitask Network Binarization

---

Binary Neural Network

  首先，简要回顾二值神经网络（BNNs）的一般思想，详细阐述二值化的加速机制和信息流退化问题。定义一个具有 $K$ 个层的全精度（FP）神经网络， $f(\mathbf{x})=({\mathbf{W}}^K\times \sigma \times {\mathbf{W}}^{K-1}\cdots \sigma \times {\mathbf{W}}^1)(\mathbf{x})$ ，其中 $\mathbf{x}$ 是输入样本， ${\mathbf{W}}^k:{\mathbb{R}}^{d_{k-1}}⟼{\mathbb{R}}^{d_k}(k=1,...,K)$ 代表连接第 $(k-1)$ 层和第 $k$ 层的权重矩阵， $d_{k-1}$ 和 $d_k$ 分别表示第 $k$ 层网络输入和输出的大小，函数 $\sigma (\cdot )$ 对特征图执行逐元素激活。

  BNNs在前向操作和反向梯度逼近方面与FP神经网络有所不同。具体而言，在前向传播过程中，BNN保持FP潜在权重 ${\mathbf{W}}_F$ 用于梯度更新，第 $k$ 个权重矩阵 ${\mathbf{W}}_F^k$ 被二值化为 $\pm 1$ ，通过二值化函数 $\text{sign}(\cdot )$ 得到二值权重矩阵 ${\mathbf{W}}_B^k$ ，即 ${\mathbf{W}}_B^k=\text{sign}({\mathbf{W}}_F^k)$ 。然后第 $k$ 层的中间激活图（全精度）由 ${\mathbf{A}}_F^k={\mathbf{W}}_B^k{\mathbf{A}}_B^{k-1}$ 产生。相同的量化方法用于将全精度激活图二值化为 ${\mathbf{A}}_B^k=\mathrm{sign}\left({\mathbf{A}}_F^k\right)$ ，并且通过对这个过程进行 $L$ 次迭代执行整个二值化的前向传递，如图2所示。对于BNNs,权重和激活都是1-bit，因此网络在内存成本方面加速了32倍。重要的是，在推断阶段，BNN加速了64倍, 因为FP网络中的FP乘法被BNN中的Xnor-Bitcount所取代。

  在反向传播中，主要挑战是普遍存在的sign函数在理论上是不可微的，因此通过梯度传播极大地破坏了信息流。为了解决这个问题，利用直通估计器（STE）来数值近似整个BNN的导数，即

KaTeX parse error: Undefined control sequence: \label at position 239: …array} \right. \̲l̲a̲b̲e̲l̲{eq:critic_new}…

  值得注意的是，在实践中没有采用上述的传统逼近方法，而是利用普遍存在的Bi-Real和IR-Net逐渐逼近符号函数，这已被证明是更好的估计方法。尽管已经有许多方法来消除二值化引起的信息流恶化，但由于权重、激活和梯度的严重精度损失，恶化仍然是不可避免的。因此，二值化破坏了复杂计算机视觉模型的性能。

Multitask Dense Predictor

  在为密集预测任务应用二值化技术后，二值化模型的性能是不可接受的，如图1和表3中的二进制单任务结果所示。因为这些SoTA密集预测模型的架构相对较重且较深（例如，HRNet-48或ResNet-101具有特定任务头），而不是用于分类的模型（例如，ResNet-18具有全连接层作为分类头）。此外，在二进制模型中通过反向传播传递信息，特别是在深度模型中，效率极低。

  密集预测任务可以相互补充信息，例如，法线和深度可以直接从彼此派生出来，这可以被建模为彼此的正则化。密集预测任务之间的相关性值得利用以提高模型的整体性能，例如在深度学习时代之前，一些开创性的工作利用带有深度信息的RGB-D图像来预测场景语义以改进预测质量。在深度学习时代，最近引起关注的基于注意力机制的多任务学习方法中，明确和隐含地从其他任务中蒸馏信息作为补充信号，从而提高目标任务的性能。简而言之，上述方法通过将现有骨干网络与后续解码过程结合起来实现初始任务预测，并如图3(左)所示。

  具体地，骨干网络提取共享特征后通过一组特定任务头进行处理，产生了一系列对于 $T$ 个任务的初始预测，即 { Y i k }   ( k = 1 , ⋯   , T ) \{Y_i^k\}~(k = 1,\cdots, T) {Yik​} (k=1,⋯,T) （骨干网络和特定任务头被称为网络的前端）。将 $Y_i^t$ 转换并二值化成1-bit特征图形式，得到相应场景的一组二值特征图，即 { F B , i t }   ( t = 1 , ⋯   , T ) \{F_{B,i}^t\}~(t = 1,\cdots,T) {FB,it​} (t=1,⋯,T) ，这比骨干网络的共享二值特征更加关注任务。在最终进行任务预测之前，通过二值化注意力机制将这些来自特定任务的表示融合为多模态蒸馏信息。论文的方法也是可扩展到不同任务，尤其是可能有一些任务只在网络前端（初始预测）中被使用。这些初始任务也被称为辅助任务，因为它们充当代理以改善最终任务的性能，如图3所示。

• Multi-Modal Distillation (MMD) via Binarized Attention Mechanism

  多模态蒸馏模块是多任务密集预测模型的关键。具体而言，利用注意力机制来引导不同任务的不同模态生成的二值特征图之间的信息传递。由于传递的信息流并不总是有帮助的，注意力可以作为一个门函数来控制信息的流动。换句话说，使网络能够自动学习集中或忽略来自其他二值特征的信息。基于二值化操作，可以通过二值化注意力公式化MMD。在传递信息给第 $k$ 个任务时，首先获得一个二值化注意力图 ${\mathbf{A}}_{B,i}^k$ ，即

$$\begin{array}{c}{\mathbf{A}}_{B,i}^k⟵\text{bool}({\mathbf{W}}_B^k\otimes {\mathbf{F}}_{B,i}^k)\end{array}$$

  其中， ${\mathbf{W}}_B^k$ 是二值卷积层的参数， ${\mathbf{F}}_{B,i}^k$ 是初始预测的二值特征图， $\otimes$ 表示卷积操作。然后，通过以下方式传递信息：

$$\begin{array}{c}{\mathbf{F}}_{B,i}^{o,k}⟵\text{sign}\left[{\mathbf{F}}_{B,i}^k+\sum _{t=1,t\ne k}^T{\mathbf{A}}_{B,i}^k\odot ({\mathbf{W}}_{B,t}\otimes {\mathbf{F}}_{B,i}^t)\right]\end{array}$$

  其中， $\odot$ 表示逐元素相乘。蒸馏过程如图3左所示，然后在图3右侧的头部中，输出的二值特征图 ${\mathbf{F}}_{B,i}^{o,k}$ 被用于相应的第 $t$ 个任务。通过使用特定任务的蒸馏激活，网络可以为每个任务保留更多信息，这对于BNNs尤其有益，因为恶化的信息流主要导致性能下降。

  另一方面，多任务密集预测模型在性能方面受益于网络二值化。尽管这些多任务模型已经取得了令人期待的性能，但它们仍然受到基于卷积的蒸馏模块的限制，这些模块在多尺度上被大量使用，模拟了相对局部感知领域中的关键空间和任务相关交互。从理论上讲，在不同网络节点中增加更多的蒸馏模块可以提高模型性能，但由于计算限制，不能无限制地向现有模型中添加蒸馏模块。幸运的是，通过二进制网络节省的计算成本，可以在模型中实现额外的蒸馏模块。

• Binary Baseline for Multitask Dense Prediction, Bi-MTDP

  为了在多任务学习框架下使用BNNs获得密集预测，将MTI-Net二值化为二进制基线。具体来说，完整精度的MTI-Net中的主要模块，包括骨干网络、头部和多模态蒸馏模块，都被替换为二进制模块（权重和激活都是1-bit），将这个基线称为Bi-MTDP。

Information Flow Supplementation

  尽管为多任务密集预测构建了一个完全二值化的基线Bi-MTDP，并使用常见技术对管道进行训练，但性能仍然是一个主要关注点。无关因素在前向传播时过度拟合，而在反向传播时优化方向严重不匹配，导致信息严重退化。为了解决这些问题，论文进一步提出了Bi-MTDP的变体，Bi-MTDP-F。具体来说，在共享的二进制骨干网络的输出后引入了一个变分信息瓶颈（VIB）层，精确地强制特征提取器保留输入数据的最小充分信息。此外，采用基于特征的知识蒸馏来引导优化方向。

• Variational Information Bottleneck for Filter-Out Nuisance Factors

  通过共享的骨干网络获取输入图像的初始二进制表示，需要训练一系列有针对性的头部来对其进行分割，一个直接的策略是将这些表示送到后续的MMD模块中。然而，二值化表示缺乏同质化，导致了模型过拟合问题。因此，需要对二值化表示进行正则化，而正则化不应该污染表示中的信息流。幸运的是，信息瓶颈（IB）原理直接涉及与最佳假设相关的压缩，即数据不匹配和模型复杂性应同时被最小化。

  由于VIB能够有效地捕捉相关部分并过滤掉输入中的无关部分，论文在骨干网络后设计了一个基于VIB的新颖层。具体来说，它明确要求特征提取器保留输入数据的最小充分信息。换句话说，它可以帮助确保信息流灵活地学习针对性任务的清晰表示。基于VIB的分类目标函数可以被表述为信息损失项，如下所示：

$$\begin{array}{c}{\mathcal{L}}_{vib}=\text{KL}\left[p(\mathbf{Z}\mid {\mathbf{A}}_{\mathbf{B}}),r(\mathbf{Z})\right],\end{array}$$

  其中， ${\mathbf{A}}_{\mathbf{B}}$ 是输入的二进制骨干表示， $\mathbf{B}$ 是潜在表示变量， $p(\mathbf{Z}\mid {\mathbf{A}}_{\mathbf{B}})$ 是一个多变量高斯分布， $r(\mathbf{Z})$ 是一个标准正态分布。一般来说，后者是一个正则化项，控制过滤掉输入信息的程度。

• Feature-based Knowledge Distillation for Guiding the Direction of Information Flow

  蒸馏是一种常见且重要的优化方法，用于缓解量化模型在超低比特宽设置下性能下降的问题，可以灵活地应用于任何架构，利用全精度教师模型的知识。通常做法是逐层地从全精度教师向量量化对应的激活进行蒸馏，即 ${\mathbf{F}}_{B,l}$ 和 ${\mathbf{F}}_{FP,l}$ （ $l=1,\cdots ,L$ ，其中 $L$ 表示网络层数）。使用均方误差（MSE）作为距离函数来衡量学生和教师之间相应特征的差异。知识蒸馏损失可以写成如下形式：

$$\begin{array}{c}{\mathcal{L}}_{kd}=\sum _{l=1}^L\text{MSE}({\mathbf{F}}_{B,l},{\mathbf{F}}_{FP,l}).\end{array}$$

Counter-Intuitive Results of Bi-MTDP-A

  直观地说，在FP网络上实施二值化必然会导致表示退化，因为sign函数的梯度无法完美估计。因此，二值化模型不可能优于其全精度对应模型。然而，Bi-MTDP的变体Bi-MTDP-C（即仅二值化多模态蒸馏的全精度骨干模型）的性能优于其全FP版本。具体来说，仅对多模态蒸馏进行二值化处理就能同时将模型加速∼39%，并将分割 mIoU提高∼4%。这一结果表明，论文的方法并不是在模型性能和效率之间进行天真的权衡，而是提升多任务密集预测器的有力工具。

  这一令人兴奋的成果甚至有点违反直觉，推测原因是：i) 在MMD上进行二值化可以过滤掉与任务无关的信息；ii) 网络内的信息流更加有效。为了验证这一推测，从两个方面进行了一系列实验，即Bi-MTDP-C的表示能力和网络内信息流的补充。

• Qualitative Study of Learned Features with Bi-MTDP

  为了研究Bi-MTDP-C及其FP对应模型的表示能力，使用t-SNE算法在2-D空间中可视化了二值化多模态蒸馏（MMD）模块后面的特征图，以及使用Grad-Cam算法确定网络认为重要的区域，结果如图5所示。显然，二值化模型Bi-MTDP-C能够通过二值化注意力模块过滤掉无关信息（见图5（a）），从而有助于学习更具有区分性的特征（见图5（b）），导致更高的定量结果。总体而言，分割生成的空间特征图更好。增强的代表能力可以促进更高的定量结果。

• Analysis of Information Flow Supplementation within Network via Centered Kernel Alignment

  分析神经网络层内的分布信息流是具有挑战性的，因为层的输出分布在大量神经元之间。中心化核对齐（CKA）可以解决这些挑战，通过定量比较网络内部或跨网络的激活。具体而言，对于一个由 $m$ 个样本输入的网络，CKA算法将 $\mathbf{X}\in {\mathbb{R}}^{m\times p_1}$ 和 $\mathbf{Y}\in {\mathbb{R}}^{m\times p_2}$ 作为输入，它们是两个层的输出激活（分别具有 $p_1$ 和 $p_2$ 个神经元）。令 $\mathbf{K}\triangleq \mathbf{X}{\mathbf{X}}^{\top }$ 和 $\mathbf{L}\triangleq \mathbf{Y}{\mathbf{Y}}^{\top }$ 表示两个层的格拉姆矩阵，CKA计算如下：

$$\begin{array}{c}\mathrm{CKA}(\mathbf{K},\mathbf{L})=\frac{\mathrm{HSIC}(\mathbf{K},\mathbf{L})}{\sqrt{\mathrm{HSIC}(\mathbf{K},\mathbf{K})\mathrm{HSIC}(\mathbf{L},\mathbf{L})}},\end{array}$$

  在这里，HSIC是希尔伯特-施密特独立准则。给定中心矩阵 $\mathbf{H}={\mathbf{I}}_n-\frac{1}{n}11^{\top }$ 和中心化的格拉姆矩阵 ${\mathbf{K}}^{\prime }=\mathbf{HKH}$ 以及 ${\mathbf{L}}^{\prime }=\mathbf{HLH}$ ， $\mathrm{HSIC}=\frac{\mathrm{vec}({\mathbf{K}}^{\prime })\mathrm{vec}({\mathbf{L}}^{\prime })}{(m-1{)}^2}$ ，表示这些中心化格拉姆矩阵之间的相似性。重要的是，CKA对于表示的正交变换（包括神经元的排列）是不变的，而归一化项确保了对各向同性缩放的不变性。这些特性使得能够对神经网络隐藏表示进行有意义的比较和分析。

  因此，引入CKA来研究多任务密集预测模型中的信息流。在热图中，点越浅，两个相应层之间的相似度越高。两个层的输出表示之间的相似得分越高，意味着这两个层共享更多信息。结果如图5（c）所示，我们可以看到Bi-MTDP-C前层和后层之间的相似得分要比MTI-Net中的相似得分高得多。这表明Bi-MTDP-C能够补充网络内部的信息流，从而提高模型性能。

Experiments

---

 
 
 

如果本文对你有帮助，麻烦点个赞或在看呗～
更多内容请关注 微信公众号【晓飞的算法工程笔记】