STanH：用参数化的量化模块实现变比特率的图像压缩

在端到端学习图像压缩中，编码器和解码器联合训练以最小化 R + λD 成本函数，其中 λ 控制量化潜在表示的速率和图像质量之间的权衡。但是，必须为每个 λ 训练具有数百万个参数的不同编码器-解码器对，因此需要切换编码器并在用户设备上为每个目标速率存储多个编码器和解码器。本文提出利用一种可微量化器，该量化器围绕双曲正切的参数和设计，称为 STanH，可放宽逐步量化函数。STanH 实现为具有可学习量化参数的可微激活层，可插入预训练的固定速率模型并进行细化以实现不同的目标比特率。实验结果表明，本文的方法能够实现与最先进技术相当的可变速率编码效率，但在部署简易性、训练时间和存储成本方面却有显著的节省。

来源：IEEE Transactions on Image Processing, 2024
题目：STanH : Parametric Quantization for Variable Rate Learned Image Compression
作者：Alberto Presta, Enzo Tartaglione, Attilio Fiandrotti, Marco Grangetto
原文链接：https://arxiv.org/abs/2410.00557
内容整理：刘潮磊

引言

研究问题：

用参数化的量化模块实现变比特率的图像压缩。

SOTA 工作与所属团队：

Asymmetric gained deep image compression with continuous rate adaptation 华为，CVPR 2021
Selective compression learning of latent representations for variable-rate image compression 韩国科学技术院，NeurIPS 2022
EVC : Towards Real-Time Neural Image Compression with Mask Decay 南京大学，ICLR 2023

动机：

同一个模型针对R+λD中不同的λ重新训练需要很长时间
存储具有不同λ的模型也要消耗很多资源

贡献：

提出了StanH参数化量化模块，只需要针对不同λ训练不同的StanH
可以插入到任何一个预训练好的编码器中
StanH参数量很小，只有几百个参数

方法

总体框架

STanH：用参数化的量化模块实现变比特率的图像压缩 | TIP 2024 — 图1

基本的编解码框架不改动
将量化模块用本文提出的StanH模块替代

StanH的设计

目标：

为每个λ找到合适的量化

要求：

可微
可以控制量化的宽度和高度

StanH设计：

L：量化等级
β：斜率
w：每一级量化的高度
b：每一级量化的宽度

StanH函数图像：

训练过程：

训练流程：

针对某一个λ训练一个多个anchor model
固定模型主要框架，只训练StanH模块

StanH的使用：

训练：w、b、β参数都在β参数主要作用是使其可微
测试：β趋于无穷大，StanH是一个由w、b控制的阶梯函数

参数w、b、β的更新：

w、b：由整个模型的loss去backward更新
β：由提出的StanH模块与实际测试采用的量化模块间的差距更新

码率控制

对于一个新的λ，有两种方式进行训练：

直接训练：针对新的λ训练一个新的StanH模块
插值：将两个StanH模块的参数(w、b)直接插值

插值方式：

实验设计与验证

实验设置

实验设置：

compressai三种编码器：Zou22, Xie21, Cheng20
量化等级：Cheng20, Xie21——L=60(y,z) Zou22——L=120(y)/40(z)

数据集：

Kodak, Clic Professional, Tecnik

实验结果

不同anchor数量的设置(使用Zou22进行实验)

不同anchor数量的设置(使用Zou22进行实验)

anchor数量设置为3在不同数据集上测试：

训练参数量和时长：

与其他变比特率编码器对比：

插值的效果：

结论

本文提出了一种新方法，通过利用 STanH（一种近似量化的参数模块）将固定速率 LIC 模型转换为可变速率。根据定义，如果 STanH 的参数 β 在训练时得到适当退火，则 STanH 会收敛到逐步量化器。本文通过仅端到端训练几个锚模型来实现可变速率，然后仅针对其他 RD 权衡将 STanH 层细化为不同的派生模型。一旦锚模型经过训练，细化其他派生模型的训练和存储成本就可以忽略不计，通过计算现有 STanH 的加权平均值，实际上可以实现细粒度和连续速率控制。总之，本文表明其的方法在固定速率和可变速率 LIC 模型方面都取得了可比的结果；此外，由于其简单性，它与参考架构完全无关。