[2021-ICCV] MUSIQ: Multi-scale Image Quality Transformer 论文简析

论文：https://arxiv.org/abs/2108.05997

代码：https://github.com/google-research/google-research/tree/master/musiq

概述

当前SOTA的IQA（图像质量评估）模型都是基于CNN的，基于CNN的模型通常受到在一个批次内，图像尺寸必须固定的限制，所以其输入图像通常会进行缩放或者裁剪，当然这会导致图像质量的下降。为了解决这个问题，本文设计了一个多尺度的图像质量transformer（multi-scale image quality transformer）来处理不同尺寸、不同长宽比的原分辨率图像。通过多尺度的图像表示，本文的模型可以捕捉到不同粒度（granularity）的图像质量。另外，本文提出了一种新型的基于哈希的（hash-based）二维空间嵌入方法和一种尺度嵌入，来作为多尺度表示中的位置嵌入。

上图中右侧(b)图是传统的CNN模型的做法，必须要缩放或者裁剪来固定输入图像尺寸，而这无疑会影响原图的图像质量；而左侧(a)图则是本文的多尺度图像质量transformer，基于图像块可以以多尺度的形式直接处理原图。

另外，由于MUSIQ只改变输入编码，因此它可以适应任何transformer变体，也就是说，本文提出的创新点是在编码阶段的处理方法，得到输入编码之后可以放到Swin、CvT之类的新型transformer里都是可以的。

本文的创新点总结如下：

1. 本文提出一种基于图像块的多尺度图像质量transformer（multi-scale image quality transformer MUSIQ），可以处理不同尺寸、不同长宽比的全尺寸输入图像，并且可以提取多尺度的特征。
2. 本文提出了一种新型的基于哈希的（hash-based）二维空间嵌入方法和一种尺度嵌入，来帮助transformer捕获空间间和尺度间的信息。
3. 本文提出的MUSIQ在四个大规模IQA数据集上取得SOTA性能。

方法

模型框架

MUSIQ的整体结构如上图所示，首先得到输入图像的多尺度表示，包括原图和固定长宽比缩放（ARP(aspect ratio preserved) resized）的变体。不同尺度的图像被分成固定大小的图像块，然后被送入到模型中，由于图像块是来自不同的空间分辨率的图像，我们需要高效地将这些多种长宽比、多种尺度的输入编码为一个token序列，捕获像素、空间和尺度信息。

为此，本文设计了三个编码模块：

1. 图像块编码模块
2. 基于哈希的空间嵌入模块
3. 可学习的尺度嵌入

分别对来自多尺度图像的图像块本身、每个图像块的二维空间位置、不同的尺度进行编码。

在将多尺度的图像输入编码为一个token序列之后，我们先准备一个额外的可学习的分类头classification head（CLS）。transformer encoder输出中的CLS token将作为最终的图像表示。然后在最后加一个全连接层来预测图像质量分。由于MUSIQ只改变输入编码，因此它可以适应任何transformer变体，也就是说，本文提出的创新点是在编码阶段的处理方法，得到输入编码之后可以放到Swin、CvT之类的新型transformer里都是可以的。

多尺度图像块嵌入

ARP resize = aspect ratio preserved resize 即固定长宽比缩放，后面不在赘述，直接简称ARP resize。

为了同时捕获局部信息和全局信息（各种多尺度方法的老说辞了^^），本文提出对图像的多尺度表示进行建模。

记全尺寸原图的高、宽、通道数分别为 $H,W,C$， 使用高斯核对全尺寸原图进行ARP resize（保持长宽比的缩放）之后的多尺度图像的高、宽、通道数分别为 $h_k,w_k,C$ ，其中 $k=1,\dots ,K$ ，$K$是每个输入的ARP resize变体的个数。为了将多尺度输入对齐，从而有一致的全局视角，将每个多尺度变体的长边固定为 $L_k$ ，从而：
$${\alpha }_k=L_k/max(H,W),h_k={\alpha }_{k}H,w_k={\alpha }_{k}W$$
${\alpha }_k$ 即为每个尺度的缩放因子。

从每个多尺度图像中切分出尺寸为 $P$ 的正方形图像块。对于高、宽不是 $P$ 的整数倍的图像，用零填充（这里笔者有个问题：既然都padding 0了，那不就相当于也将输入的尺寸固定了吗，既然可以padding，那什么网络都能处理原尺寸图像啊，没搞懂这样设计还有什么意义）。每个图像块被图像块编码器模块patch encoder module编码为 $D$ 维的嵌入，$D$​ 即为transformer中的隐层token尺寸。本文的patch encoder module使用了5层的ResNet而非线性映射。将patch encoder module输出的图像块嵌入拼接起来就得到输入图像的多尺度映射序列，来自原尺寸图像和ARP resize的多尺度图像的图像块个数就分别为：$N=HW/P^2$，$n_k=h_k{w}_k/P^2$ 。

对于输入图像尺寸不同导致的 $N,n_k$ 不同，从而序列长度不同的问题。本文采用了NLP中常用的pad+mask的方式来得到固定长度的输入，从而进行训练。前面提到过ARP resize图像的长边固定在 $L_k$​ ，因此有 $n_k<=L_k^2/P^2=m_k$ ，所以直接pad到 $m_k$ 即可。

基于哈希的二维空间嵌入

传统的固定长度的位置嵌入无法适应可变分辨率的输入，并且也无法对齐来自不同尺度但空间位置接近的图像块。

本文认为一个有效地适合MUSIQ的位置嵌入应当满足以下条件：

1. 可以在不同长宽比、不同分辨率下有效地为图像块的空间信息进行编码；
2. 不同尺度下空间位置接近的图像块应当有相近的空间嵌入
3. 便于实现，不会干扰到transformer attention

据此，本文提出了一种基于哈希的二维空间嵌入（HSE），记某个图像块的位置在第 $i$ 行，第 $j$ 列，被哈希到 $G\times G$ 的网格中的相应的元素。该网格中的每一个元素是一个 $D$ 维的嵌入向量。即有一个可学习的矩阵 $T\in {\mathbb{R}}^{G\times G\times D}$ ，输入尺寸为 $H,W$ ，对于位置在 $(i,j)$ 的图像块，其空间嵌入被定义为 $T$ 中的 $(t_i,t_j)$ 位置的元素：
$$t_i=\frac{i\times G}{H/P},t_j=\frac{j\times G}{W/P}$$
将 $D$ 维的空间嵌入 $T_{t_i,t_j}$ 逐元素地加到patch embedding上。为了快速查找，将$(t_i，t_j)$四舍五入到最接近的整数。

为了在不同尺度之间对齐图像块，来自不同尺度的图像块都映射到一个同样的表格 $T$。这样，在空间上位置接近但是来自不同的尺度的图像块也会被映射到$T$中接近的嵌入上。因为$i$和$H$以及$j$和$W$与尺寸调整因子$\alpha$成比例变化。

$T$ 的尺寸 $G$ 存在一个trade-off，$G$​ 过小会导致过多的哈希碰撞，从而使得模型无法分辨空间位置接近的图像块；过大则会导致浪费内存并且需要更多的分辨率来进行训练。

尺度嵌入

由于本文对所用的图像复用一个相同的哈希矩阵，HSE是无法分别来自不同尺度的图像块的，因此本文引入一个额外的尺度嵌入SCE来帮助模型分辨来自不同尺度的图像块。

本文将SCE定义为一个可学习的嵌入 $Q\in {\mathbb{R}}^{(K+1)\times D}$ ，因为输入有 $K$ 个尺度的变体。$Q_0\in {\mathbb{R}}^D$ 逐元素地加到所有的 $D$ 维的原分辨率图像的pathch embedding上，$Q_k\in {\mathbb{R}}^D$ 分别逐元素地加到所有 $k$ 尺度的patch embedding上。

预训练和微调

本文预训练还是在ImageNet上做的，预训练阶段会使用各种数据增广的方法来提升性能。

微调则是在图像质量和美学质量数据集上进行，在微调阶段，保持原尺寸图像作为输入，数据增广只采用对图像质量无影响的水平翻转。

实验

吊打友商部分的实验就不在这里说了，具体指标大家可以到原文中去看。我们主要看一下和方法本身有关的消融实验和可视化实验。

ARP的重要性

本实验旨在说明固定长宽比缩放(ARP)的重要性，上面几个CNN和ViT的方法是直接介绍384/224的正方形square resize的输入，对于本文多尺度方法分别做了正方形square resize的的多尺度输入，和保持原图长宽比的多尺度输入。实验结果成功证实了保持长宽比在图像质量评估中的重要作用。

除此之外，作者还用折线图配合一张图像破坏长宽比的形式直观地展现了ARP的作用。图中蓝线，即带ARP的本文方法对图像长宽比的变化非常敏感，可以敏锐地察觉到图像长宽比变化对图像质量的影响。而其他在训练时接收的事固定正方形输入的模型则感知不到这种长宽比的变化带来的影响。

多尺度结合的全尺寸输入的有效性

本实验展示的是多尺度、全尺寸输入的有效性。实验结果如下表，可以看到多尺度输入性能优于单一尺度输入 ，原尺寸输入（full）性能优于固定尺寸输入，并且多尺度结合训练的性能也比多个尺度分别训练再结合

下面的注意力可视化实验也说明了多尺度学习的作用（注意一下三列分辨率是不同的，这里为了展示缩放到同样大小来适应表格）。可以看到在高分辨率的图像中，模型更加关注细节信息；而在低分辨率的图像中，模型更加关注全局信息。

下面是基于哈希的二维空间嵌入的每个网格位置 $(i,j)$ 与其他位置处嵌入的余弦相似度。空间位置嵌入是为了反应二维图像的空间信息，通俗点说就是行列信息，即某个图像块在二维图像的哪一行，哪一列。可以看到除了和自己的相似度最高之外，和同行同列的空间嵌入的相似度也较高。这时符合我们的预期的，说明二维空间嵌入可以准确地表征图像块的行列信息。

基于哈希的空间嵌入和尺度嵌入的有效性

正文最后一个消融实验对比了有无本文的基于哈希的空间嵌入和有无尺度嵌入的情况下的模型的性能表现。可以看到，这两个嵌入是必需的。

还有一个实验对比了不同的图像块编码模块，还是本文的方法较优。

本文的概述、方法和实验就简单地介绍到这里，已经囊括方法思路和大部分细节及正文实验，此外还有一些附录实验，有兴趣的话请移步到原文中查看。