[2112] On Efficient Transformer and Image Pre-training for Low-level Vision

paper
code

Abstract

• a highly efficient and generic transformer for low-level vision
  achieve SOTA with constrained parameters and computational complexity
• the first in-depth study on image pre-training for low-level vision
  how pre-training affect internal representations of models, how to conduct an effective pre-training

Comparison on PSNR (dB) performance of the proposed EDT and state-of-the-art methods in different low-level tasks.

Method

model architecture

The proposed encoder-decoder-based Transformer (EDT). EDT processes high-resolution (e.g., in deraining) and low-resolution (e.g., in SR, $s$ refers to the scale) inputs using different paths, modeling long-range interactions at a low resolution, for efficient computation.

consist of a lite conv-based encoder-decoder and a transformer body
limitation of transformer: high computational cost, so be difficult to handle HR inputs

Structures of the convolution blocks (CB) in the encoder and decoder. The dotted boxes and lines represent additional down-sampling and up-sampling operations when processing high-resolution inputs.

encoder

• sr maintain original input size
• dn, dr down-sample to $\frac{1}{4}$ input size with stride conv

note that stack of early convs is useful for stabling the optimization.

transformer
achieve larger receptive fields at low computational cost $⟸$ down-sampler in encoder

Structure of the transformer block, including layer normalizations (LN), a (shifted) crossed local multi-head attention module ((S)CL-MSA) and an anti-blocking feed-forward network (Anti-FFN).

decoder

• sr maintain original input size, an extra up-sampler before output
• dn, dr down-sample to $\frac{1}{4}$ input size with stride conv

skip connection
fast convergence during training

(shifted) cross local attention

Shifted Crossed Local Attention. The horizontal and vertical window sizes are set to $(h,w)$ and $(w,h)$.

given input features $X\in {\mathbb{R}}^{(H\times W)\times C}$, split into 2 parts in channel dimension
$$X=[X_1,X_2]$$

where, $X_1,X_2\in {\mathbb{R}}^{(H\times W)\times \frac{C}{2}}$
each half perform MSA in either horizontal or vertical window
$$\begin{array}{rl}{X}_1^{\prime }& =H\text{-}MSA(X_1)\\ X_2^{\prime }& =V\text{-}MSA(X_2)\end{array}$$

where, window size is $(w,h)$ or $(h,w),w>h$
projection layer finally fuse attention results
$$CL\text{-}MSA(X)=proj([X_1^{\prime },X_2^{\prime }])$$

window shifted by $(\lfloor \frac{h}{2}\rfloor ,\lfloor \frac{w}{2}\rfloor )$ or $(\lfloor \frac{w}{2}\rfloor ,\lfloor \frac{h}{2}\rfloor )$
$⟹$ crossed windows with shifts dramatically increase effective receptive field
complexity of crossed window attention for $H\times W$ input
$$\Omega ((S)CL\text{-}MSA)=4HW{C}^2+2hwHWC$$

anti-blocking FFN

$$Y=Linear(Act(ABFF(Act(Linear(X)))))$$

where, $ABFF(\cdot )$ implemented by depth-wise conv with 5x5 kernel
$⟹$ eliminate possible blocking effect caused by window partition

architecture variants

Configurations of four variants of EDT. The parameter numbers and FLOPs are counted in denoising at $192\times 192$ size.

uniformly set block number in each stage to 6, expansion ratio of FFN to 2, window size $(h,w)$ to (6, 24)

Pre-Training

pre-training on ImageNet

dataset ImageNet: 200K, 15.6%
degradation

• sr: $\times 2$, $\times 3$, $\times 4$ bilinear interpolation (BI)
• dn: Gaussian noise $\sigma =15,20,50$
• dr: light, heavy rain stroke

pre-training strategy

• train a single model for a specific task, eg. $\times 2$ sr
• train a single model for highly related tasks, eg. $\times 2$, $\times 3$, $\times 4$ sr
• train a single model for unrelated tasks, eg. $\times 2$, $\times 3$ sr and 15-level dn

pre-training for unrelated tasks is used in IPT.

centered kernel alignment (CKA)

study representation similarity hidden layers in network
given $m$ data points, output of 2 layers are $X\in {\mathbb{R}}^{m\times p_1}$ and $Y\in {\mathbb{R}}^{m\times p_2}$, with $p_1$ and $p_2$ output neurons respectively
use Gram matrices $K=X{X}^T$ and $L=Y{Y}^T$ to compute CKA
$$CKA(K,L)=\frac{HSIC(K,L)}{\sqrt{HSIC(K,K)HSIC(L,L)}}$$

where, HSIC is Hilbert-Schimit independence criterion

衡量特征之间的相似性，一般通过衡量特征中样本对之间的相似性，其指标是相似性度量矩阵 (representational similarity matrices)
比较常见的相似性度量矩阵，有以下几种：
a. 点积相似性度量 dot product based similarity
(s.1)
$$⟨vec(X{X}^T),vec(Y{Y}^T)⟩=tr(X{X}^{T}Y{Y}^T)={|Y^{T}X|}_F^2$$

其中，$X{X}^T$表示从特征X的视角来衡量样本之间的关系，$Y{Y}^T$表示从特征Y的视角来衡量样本之间的关系，两者的内积即为特征X和特征Y之间的相似性
b. Hilbert-Schimit独立性度量 Hilbert-Schimit independence criterion
对于均值为0的特征X和Y，(s.1)可变形为(s.2)
$$\frac{1}{(m-1{)}^2}tr(X{X}^{T}Y{Y}^T)={|cov(X^T,Y^T)|}_F^2$$

(s.2)成立，有如下简单推导
$$\begin{array}{rl}\frac{1}{(m-1{)}^2}tr(X{X}^{T}Y{Y}^T)& =\frac{1}{(m-1{)}^2}tr((Y^{T}X{)}^T{Y}^{T}X)\\ & =\frac{1}{(m-1{)}^2}\sum _{i=1}^p\sum _{j=1}^p(Y^T{X}_{i,j})(X^T{Y}_{i,j})\\ & =\frac{1}{(m-1{)}^2}\sum _{i=1}^p\sum _{j=1}^p(Y^T{X}_{i,j}{)}^2\\ & =\frac{1}{(m-1{)}^2}cov(X^T,Y^T{)}_{i,j}^2\\ & ={|cov(X^T,Y^T)|}_F^2\end{array}$$

Hilbert-Schimit独立性度量就是把(s.2)中的协方差的Frobenius范数的平方替换成Hilbert-Schimit范数的平方
定义两个核函数$K_{i,j}=k(x_i,x_j),L_{(}i,j)=l(y_i,y_j)$，Hilbert-Schimit独立性度量可以写成(s.3)
$$HSIC=\frac{1}{(m-1{)}^2}⟨vec(HKH),vec(HLH)⟩=\frac{1}{(m-1{)}^2}tr(KHLH)$$

其中，$HKH$和$HLH$是中心化的Gram矩阵，$H_n=I_n-\frac{1}{n}11^T$
(s.2)和(s.3)的联系是：当使用线性核函数$k(x,y)=l(x,y)=x^{T}y$时，Hilbert-Schimit独立性度量退化为点积相似性度量
ref: zhihu

representation structure of EDT

Sub-figures (a)-© show CKA similarities between all pairs of layers in x2 EDT-S SR model, x2 EDT-B SR model, level-15 EDT-B denoising model with single-task pre-training, and the corresponding similarities between with and without pre-training are shown in (e)-(g). Sub-figure (d) shows the cross-model comparison between EDT-B SR and EDT-B denoising models and (h) shows the ratios of layer similarity larger than 0.6, where “s” means the similarity between the current layer in SR and any layer in denoising.

key findings

• sr models show clear stages in internal representations and proportion of each stage vary with model size, while dn models present a relatively uniform structure
• layers in dn models show more similarity to lower layers in sr models, containing more local information
• single-task pre-training mainly affect higher layers in sr models, but have limited impact on dn models

single and multi-task pre-training

Sub-figures (a)-(d) show CKA similarities of x2 SR models, without pre-training as well as with pre-training on a single task (x2), highly related tasks (x2, x3, x4 SR) and unrelated tasks (x2, x3 SR, level-15 denoising). Sub-figures (e)-(h) show the corresponding attention head mean distances of Transformer blocks. Note that we do not plot shifted local windows in (e)-(h) so that the last blue dotted line (“—”) has no matching point. The red boxes indicate the same attention modules.

key findings

• representations of sr models contain more local information in early layers while more global information in higher layers
• 与单任务预训练相比，相关多任务预训练(fig.g)将更多的第3个块的全局表征转化成了局部表征，增加了第2块的范围；不相关多任务预训练(fig.h)的这种现象的表现偏弱
• all 3 pre-training methods can greatly improve the performance by introducing different degrees of local information, treated as a kind of inductive bias, to the intermediate layers of model, among which multi-related-task pre-training performs best

PSNR(dB) improvements of single-task, multi-related-task and multi-unrelated-task pre-training in $\times 2$ SR.

multi-related-task pre-training is more effective and efficient, providing initialization for multiple tasks.
multi-task pretraining is more effective and data-efficient, for its setting enable transformer body to see more samples in an iteration.

PSNR(dB) improvements of different data scales during single-task pre-training for EDT-L in $\times 2$ SR.

incremental PSNR improvements on multiple SR benchmarks by increasing the data scale
noted that double pre-training iterations (from 500K to 1M) for the data scale of 400K so that data can be fully functional, and longer pre-training period largely increases training burden.

PSNR(dB) results of different pre-training (single-task) data scales in $\times 2$ SR. “EDT-B$†$” refers to the base model with single-task pre-training and “EDT-B$\ast$” represents the base model with multi-related-task pre-training. The best results are in bold.

Experiment

pre-training

• input
  • sr, in transformer: $48\times 48$-size
  • dn, dr: $192\times 192$-size
• optimizer Adam: $\beta 1=0.9,\beta 2=0.99$, batch size=8, 500K iterations
• learning rate initial 2e-4, halved at 250K, 400K, 450K, 475K-th iteration

fine-tuning

• extra 500K, 800K, 200K iterations for sr, dn, dr
• learning rate constant 1e-5

super resolution

Quantitative comparison for classical SR on PSNR(dB)/SSIM on the Y channel from the YCbCr space. “$‡$” means the $\times 3$ and $\times 4$ models of SwinIR are pre-trained on the $\times 2$ setup and training patch size is $64\times 64$ (ours is $48\times 48$). “$†$” indicates methods with a pre-training. Best and second best results are in red and blue colors.

Quantitative comparison for lightweight SR on PSNR(dB)/SSIM on the Y channel. “$‡$” means the $\times 3$ and $\times 4$ models of SwinIR are pre-trained on the $\times 2$ setup and the training patch size is $64\times 64$ (ours is $48\times 48$ and without pre-training).

denoising

Quantitative comparison for color image denoising on PSNR(dB) on RGB channels. “$‡$” means the $\sigma =25/50$ models of SwinIR are pre-trained on the $\sigma =15$ level. “$†$” indicates methods with a pre-training. “$\ast$” means our model without down-sampling and without pre-training.

deraining

Quantitative comparison for image deraining on PSNR(dB)/SSIM on the Y channel. “$†$” indicates methods with a pre-training.

ablation study

window size

Ablation study of window size on PSNR(dB) in $\times 2$ SR. Best results are in bold.