论文地址：
https://arxiv.org/abs/2405.14136
论文代码：
https://github.com/42Shawn/BiMTDP
Introduction
在计算机视觉领域，有一个日益增长的趋势，即以多任务学习的方式处理密集预测任务，例如语义分割、单眼深度估计和人体解析。由于多任务模型中通过交叉连接结构进行信息补充机制的收益，一系列密集预测任务的整体性能得到了极大提升（见图
1
）。然而，处理多个复杂像素级任务的当前最先进（
SoTA
）多任务密集预测模型具有相当大的计算需求，这种高计算要求限制了它们在资源受限环境中（如自动驾驶、机器人技术和虚拟现实等领域）的应用。论文旨在优化这些庞大的
SoTA
模型以适用于边缘设备，在速度和性能之间取得平衡。
已经不是研究探索神经网络压缩策略，包括剪枝、网络量化和知识蒸馏。值得注意的是，网络二值化作为一种量化形式，将权重和激活最小化为
，从而能够用二值神经网络（
BNNs
）中更高效的
XNOR-Bitcount
操作替代全精度网络中计算密集型的内积运算。理论上，二值化可以将存储成本降低
32
倍，并将推断速度提高
64
倍，使得
BNNs
适用于边缘设备。
虽然二值神经网络（
BNNs
）在图像分类方面取得了令人印象深刻的成果，几乎达到了全精度
ResNet
级别的准确性，但其应用主要局限于小规模模型，无法用于其他计算密集型的计算机视觉任务。将
BNNs
扩展到更大的模型应该是下一步的研究，但这种扩展受到过拟合和信息降级等问题的阻碍。在全精度模型中有效的技术，如标签平滑、
dropout
和
mixup
，在
BNNs
上效果较差。此外，最先进的多任务密集预测任务通常需要深层复杂模型并配备多模态融合结构，加剧了有效实施二值化所面临的挑战。
在应用二值化进行多任务密集预测任务时，主要障碍是在深度模型中信息流的显著降级，导致性能下降。为了解决这个问题，论文提出了一个
Binary Multitask Dense Predictor
（
Bi-MTDP
）基线，通过二值化模块构建了一个多任务密集预测器。基于对这个基线的彻底评估，论文得出结论：二值化操作破坏了多任务模型中的信息流，因此与其全精度对应物相比，下游任务的表示不够丰富。为解决这一问题，论文通过两个方向上额外的信息流校准机制更新
Bi-MTDP
在正向传播中实施变分信息瓶颈机制，使嵌入遵循具有稀疏性质的高斯分布以过滤掉与任务无关的因素。
在反向传播中利用现有全精度模型通过基于特征的知识蒸馏来校准二值网络梯度。
Bi-MTDP
的好处可以从两个正交的角度进行分析：
从网络二值化的角度来看，通过将二值化与多任务密集预测框架相结合，
Bi-MTDP
的成功证明了它可以有效地补充信息，从而提高了各个二值模型的性能。
从多任务密集预测任务的角度来看，加速这些繁琐的模型有利于在其中设计更有效和高效的交叉模块，如图
1
所示。由于现有的密集预测模型在建模交叉模块方面存在严重限制，因为它们过度依赖卷积操作，所以对于多任务密集预测来说，通过交叉机制学习涵盖多任务上下文各个范围的交互和推理是至关重要的。
Multitask Network Binarization
Binary Neural Network
首先, 简要回顾二值神经网络 ( BNNs ) 的一般思想, 详细阐述二值化的加速机制和信息流退化问题。定义一个具有
个层的全精度（FP）神经网络，
, 其中
是输入样本,
代表连接第
层和第
层的权重矩阵，
和
分别表示第
层网络输入和输出的大小, 函数
对特征图执行逐元素激活。
BNNS 在前向操作和反向梯度逼近方面与 FP 神经网络有所不同。具体而言, 在前向传播过程中, BNN 保持 FP 潜在权重
用于梯度更新, 第
个权重矩阵
被二值化为
, 通过二值化函数
得到二值权重矩阵
, 即
。然后第
层的中间激活图（全精度）由
产生。相同的量化方法用于将全精度激活图二值化为
，并且通过对这个过程进行
次迭代执行整个二值化的前向传递，如图 2 所示。对于 BNNs, 权重和激活都是 1-bit, 因此网络在内存成本方面加速了 32 倍。重要的是, 在推断阶段, BNN 加速了 64 倍, 因为 FP 网络中的 FP 乘法被 BNN 中的 Xnor-Bitcoun
所取代。
在反向传播中，主要挑战是普遍存在的sign函数在理论上是不可微的，因此通过梯度传播极大地破坏了信息流。为了解决这个问题，利用直通估计器（STE）来数值近似整个BNN的导数，即
值得注意的是，在实践中没有采用上述的传统逼近方法，而是利用普遍存在的
Bi-Real
和
IR-Net
逐渐逼近符号函数，这已被证明是更好的估计方法。尽管已经有许多方法来消除二值化引起的信息流恶化，但由于权重、激活和梯度的严重精度损失，恶化仍然是不可避免的。因此，二值化破坏了复杂计算机视觉模型的性能。
Multitask Dense Predictor
在为密集预测任务应用二值化技术后，二值化模型的性能是不可接受的，如图
1
和表
3
中的二进制单任务结果所示。因为这些
SoTA
密集预测模型的架构相对较重且较深（例如，
HRNet-48
或
ResNet-101
具有特定任务头），而不是用于分类的模型（例如，
ResNet-18
具有全连接层作为分类头）。此外，在二进制模型中通过反向传播传递信息，特别是在深度模型中，效率极低。
具体地, 骨干网络提取共享特征后通过一组特定任务头进行处理, 产生了一系列对于
个任务的初始预测, 即
（骨干网络和特定任务头被称为网络的前端）。将
转换并二值化成 1-bit 特征图形式, 得到相应场景的一组二值特征图, 即
, 这比骨干网络的共享二值特征更加关注任务。在最终进行任务预测之前, 通过二值化注意力机制将这些来自特定任务的表示融合为多模态蒸馏信息。论文的方法也是可扩展到不同任务, 尤其是可能有一些任务只在网络前端（初始预测）中被使用。这些初始任务也被称为辅助任务, 因为它们充当代理以改善最终任务的性能, 如图 3 所示。
Multi-Modal Distillation (MMD) via Binarized Attention Mechanism
多模态蒸馏模块是多任务密集预测模型的关键。具体而言，利用注意力机制来引导不同任务的不同模态生成的二值特征图之间的信息传递。由于传递的信息流并不总是有帮助的，注意力可以作为一个门函数来控制信息的流动。换句话说，使网络能够自动学习集中或忽略来自其他二值特征的信息。基于二值化操作，可以通过二值化注意力公式化
MMD
。在传递信息给第
个任务时, 首先获得一个二值化注意力图
，即
其中,
是二值卷积层的参数,
是初始预测的二值特征图,
表示卷积操作。然后, 通过以下方式传递信息:
其中,
表示逐元素相乘。蒸馏过程如图 3 左所示, 然后在图 3 右侧的头部中, 输出的二值特征图
被用于相应的第
个任务。通过使用特定任务的蒸㽞激活, 网络可以为每个任务保留更多信息, 这对于 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
的分类目标函数可以被表述为信息损失项，如下所示：
其中,
是输入的二进制骨干表示,
是潜在表示变量,
是一个多变量高斯分布,
是一个标准正态分布。一般来说, 后者是一个正则化项, 控制过滤掉输入信息的程度。
Feature-based Knowledge Distillation for Guiding the Direction of Information Flow
蒸馏是一种常见且重要的优化方法, 用于缓解量化模型在超低比特宽设置下性能下降的问题, 可以灵活地应用于任何架构, 利用全精度教师模型的知识。通常做法是逐层地从全精度教师向量量化对应的激活进行蒸馏，即
和
，其中
表示网络层数）。使用均方误差（MSE）作为距离函数来衡量学生和教师之间相应特征的差异。知识蒸馏损失可以写成如下形式:
Counter-Intuitive Results of Bi-MTDP-A
直观地说，在
FP
网络上实施二值化必然会导致表示退化，因为
sign
函数的梯度无法完美估计。因此，二值化模型不可能优于其全精度对应模型。然而，
Bi-MTDP
的变体
Bi-MTDP-C
（即仅二值化多模态蒸馏的全精度骨干模型）的性能优于其全
FP
版本。具体来说，仅对多模态蒸馏进行二值化处理就能同时将模型加速
∼39%
，并将分割
mIoU
提高
∼4%
。这一结果表明，论文的方法并不是在模型性能和效率之间进行天真的权衡，而是提升多任务密集预测器的有力工具。
这一令人兴奋的成果甚至有点违反直觉，推测原因是：i) 在
MMD
上进行二值化可以过滤掉与任务无关的信息；
ii
) 网络内的信息流更加有效。为了验证这一推测，从两个方面进行了一系列实验，即
Bi-MTDP-C
的表示能力和网络内信息流的补充。
