ICLR2023 | 用于图像复原的基础二值卷积单元

CV君

发布于 2023-09-27 13:57:04

4090

发布于 2023-09-27 13:57:04

文章被收录于专栏：我爱计算机视觉我爱计算机视觉

本文分享 ICLR 2023 论文Basic Binary Convolution Unit For Binarized Image Restoration Network ，介绍用于图像复原的基础二值卷积单元。

具体信息如下：

论文链接：https://arxiv.org/pdf/2210.00405.pdf
项目链接：https://github.com/Zj-BinXia/BBCU

01. 二值网络对于复原任务的意义

虽然目前模型的量化以8bit 4bit比较多，但随着大模型的兴起，以及图像复原任务在移动端部署的重要需求，我们需要一种更加强大的模型压缩方案。二值网络（Binary Neural Network, BNN, 1bit网络）就是实现这一目标的不二之选。二值网络是部署模型到资源受限的设备上最有前途的神经网络压缩方法之一（配合特定的芯片以及框架，可以实现32x的内存压缩和高达64x的计算量减少）。事实上，在8bit 4bit压缩方案趋于成熟以后，各方也在积极地探索二值网络的落地。这篇论文全面的探索了二值网络对于复原网络的影响，经过大量实验以及分析以后，提出了一种用于图像复原的基础二值卷积单元（Basic binary convolurion unit, BBCU）以及二值化方案，相较于之前的二值复原网络在超分，去噪，去压缩等任务上均有了显著的性能提升（比如超分任务上维持相同性能的情况下，还能加速10倍）。

图1. BBCU在不同复原任务上高效地取得了SOTA的性能

02 二值网络简介

与全精度（32位）CNN模型相比，二值神经网络（BNN）可以实现高达32倍的内存节省。此外，由于激活也是二值化的，因此可以通过按位XNOR操作和位计数操作实现二值卷积运算。图2描述了一个简单的二值卷积示例。与需要浮点运算的全精度卷积不同，二值卷积依赖于位运算，可以获得高达64倍的计算节省。

图2. 二值卷积中的同或（XNOR）和位计数（bit counting）算子的示意图

03 摘要

图像复原（Image Restoration, IR）旨在从受各种退化因素破坏的低质量（Low Quality, LQ）图像中恢复高质量（High Quality, HQ）图像。典型的IR任务包括图像去噪、超分辨率（SR）和去除压缩伪影。

由于其不适定性和高实用价值，图像恢复是计算机视觉中一个活跃而具有挑战性的研究课题。最近，深度卷积神经网络（Convolutional Neural Network, CNN）通过学习从LQ到HQ图像块的映射,获得了出色的性能。

然而，大多数图像复原任务需要密集的像素预测，而基于CNN的模型的强大性能通常依赖于增加模型大小和计算复杂度，这需要大量的计算和内存资源。因此，大大降低CNN模型的计算和内存成本同时保持模型性能对于推广图像复原模型是非常重要的。

二值神经网络（Binary Neural Network, BNN，也称为1bit CNN）已被认为是部署模型到资源受限的设备上最有前途的神经网络压缩方法之一。在特殊设计的处理器上，BNN可以实现32x的内存压缩和高达64x的计算量减少。

目前，BNN的研究主要集中在高级视觉任务，特别是图像分类，但尚未完全探索低级视觉任务，例如图像超分辨率。考虑到BNN对于部署图像复原深度网络的重要性以及高级和低级视觉任务之间的差异，我们迫切需要探索BNN在低级视觉任务中的特性，并为后续研究和部署提供简单、强大、通用和可扩展的基线。

实验表明，我们的BBCU在多个复原任务上均获得了显著的效果提升。

04 用于图像复原的基础二值卷积单元设计

如图3所示，我们首先构建了BBCU-V1。具体来说，全精度卷积

X_{j}^{f} \otimes W_{j}^{f}

（其中

X_{j}^{f}

，

W_{j}^{f}

，和

\otimes

分别表示全精度激活、权重、和卷积）可以被二值卷积

X_{j}^{b} \otimes W_{j}^{b}

所近似。对于二值卷积，权重和激活都被二值化为-1和+1。高效的按位XNOR和bitcount操作可以取代计算量大的浮点矩阵乘法，定义如下：

其中，

X_{j}^{f} \in \mathbb{R}^{C\times H\times W}

和

W_{j}^{f} \in \mathbb{R}^{C_{out}\times C_{in} \times K_{h} \times K_{w}}

分别是 j 层中的全精度激活和卷积权重。同样，

X_{j}^{b} \in \mathbb{R}^{C\times H\times W}

和

W_{j}^{b} \in \mathbb{R}^{C_{out}\times C_{in} \times K_{h} \times K_{w}}

分别是 j 层中的二值化激活和卷积权重。

x_{i,j}^{f} ，w_{i,j}^{f} ，x_{i,j}^{b} 和 w_{i,j}^{b}

分别是

X_{j}^{f} ，W_{j}^{f} ，X_{j}^{b} 和 W_{j}^{b}

中第 i 个通道的元素。

\alpha _{i,j}

是可学习的调整系数，用于控制

x_{j}^{f}

中第 i 个通道的符号函数的阈值。值得注意的是，权重二值化方法继承自 XNOR-Net，其中

\frac{\parallel W_{j}^{f} \parallel _{1} }{n}

是绝对权重值的平均值，作为缩放因子以最小化二值卷积权重和全精度卷积权重之间的差异。

接下来，我们使用RPReLU作为激活函数，其定义如下：

其中，

Y_{j} \in \mathbb{R}^{C\times H\times W}

是RPReLU函数

f(\cdot )

在第 j 层的输入特征图。

y_{i,j}

是

Y_{j}

的第 i 个通道的元素。

\gamma _{i,j}

和

\zeta _{i,j }

是学习得到的偏移量，用于移动分布。

\beta _{i,j }

是一个可学习的系数，用于控制负值部分的斜率，作用于

Y_{i}

的第 i 个通道。

与常见的由两个全精度卷积组成的残差块不同，我们发现对于二值卷积，残差连接是必不可少的，它可以补充二值化所造成的信息丢失。因此，我们为每个二值卷积设置了一个残差连接。因此，BBCU-V1可以用以下数学公式表示：

其中，

k _{j} ，\tau _{j} \in \mathbb{R} ^{C}

是第 j 个BatchNorm ( BN) 层的可学习参数。

在BNN中，用于二值化的符号函数的导数是一个脉冲函数，无法用于训练。因此，我们采用近似导数函数作为符号函数的导数。它可以用数学公式表示为：

然而，EDSR证明了BN会改变图像的分布，这对于超分辨率中精确的像素预测是不利的。那么，我们是否可以直接删除BBCU-V1中的BN以获得BBCU-V2（图3（b））呢？在BBCU-V1（图3（a））中，可以看到二值卷积的bitcount操作倾向于输出大的值域（从-15到15）。相比之下，残差连接保留了全精度信息，其值范围是从约-1到1的小值域，这是从图像复原网络的最开始输入图像（值域0到1）传递来的。添加BN后，BN的可学习参数可以缩小二值卷积的值域，并使二值卷积的值域接近残差连接的值域，以避免全精度信息被覆盖。该过程可以表示为：

其中

k _{j} ，\tau _{j} \in \mathbb{R} ^{C}

是第 j 层BN的可学习参数。因此，与BBCU-V1相比，BBCU-V2简单地移除了BN并且性能大幅下降。

经过上述探索，我们知道BN对于BBCU非常重要，因为它可以平衡二值卷积和残差连接的值范围。然而，BN会改变图像分布，从而限制了图像的复原效果。在BBCU-V3中，我们提出了残差对齐（Residual Alignment, RA）方案，通过将输入图像的值域乘以一个放大因子

k (k>1)

来代替BN，如图3（c）和图3（b）所示：

从式子可以看出，由于残差连接的值来自放大值域后的输入图像，因此

X_{j}^{f}

的值域也会被放大，我们将其定义为

kX_{j}^{f}

。与此同时，二值卷积的值几乎不受RA的影响，因为

X_{j}^{b}

过滤了

kX_{j}^{f}

的幅度信息。与BN不同，我们所提出的RA策略反而使残差连接的值域接近二值卷积（-60到60之间）。

此外，我们发现使用RA去除BN后还具有两个主要优点：（1）与全精度复原网络类似，去除BN后，二值化的IR网络性能也更好。（2） BBCU的结构变得更加简单、高效、强大。

基于上述发现，我们意识到BBCU-V3（图3（c）中的激活函数会缩小残差连接的负值范围，这意味着它会丢失负的全精度信息。因此，我们进一步开发了BBCU-V4，通过将激活函数移入残差连接中，避免丢失负的全精度信息。实验结果（请见实验章节）表明我们的设计是有效的。因此，我们将BBCU-V4作为BBCU的最终设计。

图3. BBCU 改进过程中网络结构和激活分布的展示图

05 用于复原网络

如图4（a）所示，现有的图像恢复（IR）网络可以分为四个部分：头部

、主体

、上采样

和尾部

\tau

。如果IR网络不需要增加分辨率，则可以删除上采样

部分。具体而言，给定LQ输入

I_{LQ}

，IR网络恢复HQ输出

\hat{I} _{HQ}

的过程可以表示为：

先前的二值化超分辨率工作（如BAM,BTM}主要集中在对主体

部分进行二值化。然而，上采样

部分占总计算量的 52.3% ,是十分有必要进行二值化的。此外，如何对网络的头部

和尾部

\tau

进行二值化也是值得探索的。

图4 全精度图像复原网络和二值化图像复原网络的结构示意图

如图4（b）所示，我们进一步设计了适用于这四个部分的BBCU的不同变体。

对于头部

部分，其输入是

I_{LQ} \in \mathbb{R}^{3\times H\times W}

，而输出是具有

通道的二值卷积输出。由于通道数不同，我们不能直接将

I_{LQ}

加到二值卷积输出上。为了解决这个问题，我们提出了BBCU-H，通过重复

I_{LQ}

以得到

通道。

对于身体

部分，由于输入和输出通道相同，我们提出了BBCU-B，通过直接将输入激活添加到二值卷积输出中。

对于上采样

部分，我们通过重复输入激活的通道以与二值卷积相加，提出了BBCU-U。

对于尾部

\tau

部分，我们通过采用

I_{LQ}

作为残差连接，提出了BBCU-T。为了更好地评估二值化每个部分对计算和参数的好处，我们定义了两个指标：

其中，

PSNR^b

、

OPs^b

和

Parms^b

分别表示在对网络的某一个部分进行二值化后网络的性能、计算量和参数量。同样地，

PSNR^f

、

OPs^f

和

Parms^f

用于衡量全精度网络的性能、计算量和参数量。

06 消融实验

为了验证二值图像复原网络中BBCU的有效性 ，我们将SRResNet的主体部分分别用四种BBCU的变体（图3）进行二值化。结果如表1所示。

与BBCU-V1相比，BBCU-V2在Urban100和Manga109上分别降低了0.14dB和0.24dB。这是因为BBCU-V2仅仅去除了BatchNorm（BN），使得二值卷积的值域比残差连接大得多，并覆盖了全精度信息（图3）。
BBCU-V3在BBCU-V2的基础上添加了残差对齐方案，解决了二值卷积和残差连接之间的值域不平衡问题，并去除了BN。由于BN对图像复原网络不利，因此BBCU-V3在Set5、Urban100和Manga109上分别超过BBCU-V1 0.17dB、0.16dB和0.29dB。
BBCU-V4将激活函数移入残差连接中，保留了残差连接中的全精度负值信息（图3）。因此，BBCU-V4的表现优于BBCU-V3。