【论文精读】Pre-Trained Image Processing Transformer

出处

摘要

随着现代硬件计算能力的强劲增长，在大规模数据集上学习的预训练深度学习模型(如BERT，GPT-3)已经显示出比传统方法更有效的效果。这一重大进展主要得益于transformer及其变体网络模型的表现能力。本文对底层计算机视觉任务(如去噪、超分辨率和去噪)进行了研究，提出了一种新的预训练模型，即图像处理变换器(IPT)。为了最大限度地挖掘转换器的能力，我们提出利用著名的ImageNet基准测试来生成大量损坏的图像对。IPT模型在这些具有多头和多尾的图像上进行训练。此外，为了更好地适应不同的图像处理任务，还引入了对比学习。因此，经过微调后，预先训练的模型可以有效地应用于期望的任务。由于只有一个预先训练的模型，IPT在各种低级基准上的表现优于当前最先进的方法。

论文贡献

提出了一种处理低水平视觉任务（超分辨率、图像去雾、图像去噪）的预训练模型。

网络结构

IPT的总体结构由四部分组成：头部用于从输入受损图像(如含噪声图像和低分辨率图像)中提取特征，编解码器转换器用于恢复输入数据中的丢失信息，尾部用于将特征映射到恢复图像。

1.头部（Heads）

为了调整不同的图像任务，采用多头结构来分别处理每个任务，每个头由三个卷积层组成。计算过程可表示为：
$$f_H=H^i(x)\text{\hspace{1em}(1)}$$
其中，$x$为输入的原彩色图像，$H^i$表示第$i$个任务的头部。

2.编码器（Transformer encoder）

在将特征输入到编码器之前，需要对特征图像$f_H\in {\mathbb{R}}^{\mathbb{C}\times \mathbb{H}\times \mathbb{W}}$分割成块 f p i ∈ R P 2 × C , i = { 1 , . . . , N } f_{p_i}\in\mathbb{R}^{P^2\times C},i=\{1,...,N\} fpi​​∈RP2×C,i={1,...,N}其中$N=\frac{HW}{P^2}$表示块的数量，$P$表示块的大小。为了保持每个块的位置信息，增加了可学习的位置编码$E_{p_i}\in {\mathbb{R}}^{P^2\times C}$。结合每一个特征块组成$E_{p_i}+f_{p_i}$的形式输入编码器中，编码器的输出$f_{E_i}\in {\mathbb{R}}^{P^2\times C}$与输入块$f_{p_i}$大小相同。计算方法可表示为：
$$y_0=[E_{p_1}+f_{p_1},E_{p_2}+f_{p_2},...,E_{p_N}+f_{p_N}]$$
$$q_i=k_i=v_i=LN(y_{i-1})\text{\hspace{1em}(2)}$$
$$y_i^{\prime }=MSA(q_i,k_i,v_i)+y_{i-1}$$
$$y_i==FFN(LN(y_i^{\prime }))+y_i^{\prime },i=1,...,l$$
$$[f_{E_1},f_{E_2},...,f_{E_N}]=y_l$$
其中$l$表示编码器中的层数，MSA表示多头自关注模块，FFN表示包含两个全连接层的前馈网络。

3.解码器（Transformer decoder）

解码器也采用相同的结构，由两个多头自关注（MSA）层和一个前馈网络（FFN）组成。这里使用了特定于任务的嵌入作为解码器的附加输入，这些特定的任务嵌入$E_t^i\in {\mathbb{R}}^{P^2\times C}$被学习来为不同的任务解码特征。解码器的公式表达为：
$$z_0=[f_{E_1},f_{E_2},...,f_{E_N}]$$
$$q_i=k_i=LN(z_{i-1})+E_t,v_i=LN(z_{i-1})$$
$$z_i^{\prime }=MSA(q_i,k_i,v_i)+z_{i-1}$$
$$q_i^{\prime }=LN(z_i^{\prime })+E_t,k_i^{\prime }=v_i^{\prime }=LN(z_0)\text{\hspace{1em}(3)}$$
$$z_i^{\prime \prime }=MSA(q_i^{\prime },k_i^{\prime },v_i^{\prime })+z_i^{\prime }$$
$$z_i=FFN(LN(z_i^{\prime \prime }))+z_i^{\prime \prime },i=1,...,l$$
$$[f_{D_1},f_{D_2},...,f_{D_N}]=y_l$$
其中$f_{D_i}$为解码器的输出，将解码的$N$个大小为$P^2\times C$的特征块，重塑为尺寸为$C\times H\times W$的特征图$f_D$

尾部（Tails）

尾部和头部性质相同，使用多个尾部来处理不同的任务。计算公式为：
$$f_T=T^i(f_D)\text{\hspace{1em}(4)}$$
其中$T^i(i=1,...,N_t)$表示第$i$个任务的头部，$N_t$表示任务的数量。输出$f_T$是由特定任务决定的结果图像尺寸。