基于Transformer的熵模型在图像压缩领域的突破性研究
基于Transformer的熵模型在图像压缩领域的突破性研究
本文介绍了一篇发表在ICLR 2022上的论文《ENTROFORMER: A TRANSFORMER-BASED ENTROPY MODEL FOR LEARNED IMAGE COMPRESSION》，该论文提出了一种基于Transformer的熵模型Entroformer，用于改进图像压缩技术。研究团队通过引入多头注意力机制、菱形相对位置编码和并行双向上下文模型等创新技术，实现了在低比特率下的高性能压缩效果。
问题定义
在图像压缩的变换编码方法中，编码器使用参数分析变换将输入图像转化为潜在表示y，然后量化形成y ̂。因为y ̂是离散值，它可以使用熵编码技术，如算术编码进行无损压缩，并作为比特序列传输。另一方面，解码器从压缩信号中恢复y ̂，并对其进行参数合成变换，恢复重构图像x ̂。在本文中认为变换g_a和g_s是一般的参数化函数，如人工神经网络(ann)，而不是传统压缩方法中的线性变换。然后，参数θ_g和φ_g封装了神经元的权重等。因此问题定义如下：
方法
基于Transformer的熵模型
本文提出的基于Transformer的熵模型如图所示，图中对于超先验，我们堆叠N个Transformer编码器层以产生超潜势z，该z被量化并提供给N个Transformer编码器层以生成超先验。生成层次代表离子时，通过降采样和高档模块，改变了特征的分辨率。对于自回归先验，我们堆叠2N个Transformer解码器层来生成自回归特征。生成高斯一个参数，一个线性层被附加到超先验特征和上下文特征的组合上。
位置编码
为了探索位置在图像压缩中的影响，我们首先设计了一个基于Transformer的熵模型，没有任何位置编码。在训练过程中，我们随机戴上一个面具来集中注意力。在测试过程中，我们评估了每个位置i都采用相应的掩码，除位置i外，掩码设置为1。下图绘制了将每个位置的比特率与所有位置的上下文结果进行比较。这一结果突显了利率的影响根据上下文的位置。这一观察结果提供了全面的理解，并提供了经验指导用于新的位置编码设计。
为了获取潜在因素的空间信息，一种常见的方法是使用基于相对关系的偏置注意权重。基于上图中的结果，我们提出了一个扩展到这种关系感知的自注意，以考虑元素位置对图像压缩性能的影响。具体如下式所示：
自注意力Top-k机制
Transformer中常见的关注是对比例点积的关注。它计算具有所有键的查询的点积，将每个键除以√dk，并应用一个softmax函数得到值的权重。因此本文提出了基于Top-k的自注意力机制，具体如下面公式所示：
将查询、键和值打包成矩阵Q、K和V。然后计算注意力矩阵eij = (QK^T + P)/√dk,其中P是相对位置编码。接着使用Top-k选择操作fk(·)从每一行中选择k个最大的元素,最后通过softmax得到加权的值V'。
这种Top-k自注意力机制有两个好处：1)缓解了序列长度不匹配的问题；2)有助于去除无关的上下文信息，提高了模型的收敛速度。
实验
整体效果
本文通过计算速率失真（RD）性能来评估Transformer熵模型的影响。下图显示了公开可获得的柯达数据集的RD曲线，采用峰值信噪比（PSNR）作为图像质量度量。如左侧所示，我们结合超先验模块和上下文模块的性能优于CNNs方法，提高了5.2%，比BPG提高了20.5%。如右图所示，评估了两个变熵模型，即仅超优先级模型和仅上下文模以隔离Entroformer架构的影响。此外，具有双向上下文的并行Entroformer是也很有竞争力，比之前的基线更好。
位置编码影响
本文进行了组件分析，以研究不同位置的影响本文的Entroformer中的编码。本文构建了具有不同位置编码的基于变换器的上下文模型。A.基于CNN的上下文模型被实现用于比较。如表所示，应用绝对位置时在Entroformer中编码或非位置编码，本文可以实现比基线有限的bpp节省。如果本文使用相对位置编码，其性能优于绝对位置编码（节省4.0%bpp，节省0.9%bpp节省）。此外，将位置编码从1D扩展到2D至关重要（节省4.9%bpp，而节省4.0%bpp节省）。最后，结合菱形边界，本文可以比其他方法节省更多的比特率，特别是，与CNN相比，节省了5.8%的bpp。
本文评估了改变裁剪距离h的影响。下图通过显示在单个速率点上的比特率的相对减少来比较h的数量。基线为由一个基于转换器的上下文模型实现，而不需要进行位置编码。当h = 3时，它达到了最好的结果。
Top-k中k的影响
Top-k方案的影响如图所示。参数k指定了自我关注所涉及的注意力数量。本文绘制了柯达测试集上不同k的完整学习bpp曲线，而不仅仅是单个RD点。给出了原始自我注意机制的附加曲线（虚线）以供比较。
值得注意的是，当k≤64时，压缩性能有所提高，这远小于训练和测试的序列长度576和1536。此外Top-k方案也会影响Entroformer的收敛速度。有趣的是，当k大于64时，本文看到结果完全不同，top-k方案与原始自我注意没有区别。一个假设是存在在密集的自我关注中压缩大量无关信息用于图像压缩。这些观察结果反映了去除不相关的标记有利于Entroformer训练的收敛。
总结
本文提出了一种基于Transformer的熵模型Entroformer，用于改进图像压缩技术。研究团队通过引入多头注意力机制、菱形相对位置编码和并行双向上下文模型等创新技术，实现了在低比特率下的高性能压缩效果。实验结果表明，Entroformer在低比特率下优于目前最先进的基于CNN的方法以及标准的BPG编解码器，分别提高了5.2%和20.5%。这是首次成功将基于Transformer的方法应用于图像压缩领域，开创了新的研究方向。