【论文笔记】RAQ-VAE: Rate-Adaptive Vector-Quantized Variational Autoencoder
目录
- 笔记作者
- 1.引言
- 2.背景
- 3.Rate-Adaptive VQ-VAE
- 3.1 基于模型的 Rate-Adaptive VQ-VAE
- 3.2数据驱动的自适应码率 VQ-VAE
- 4.实验
- 4.1 视觉任务的主要结果
- 4.2 详细分析
- 5 .结论
笔记作者
吴思雨,女,西安工程大学电子信息学院,2023级研究生,张宏伟人工智能课题组
研究方向:人工智能与机器视觉
电子邮件:2879944563@qq.com
原文链接:https://arxiv.org/html/2405.14222v1
1.引言
向量量化(VQ)是机器学习中用于学习离散表示的基本技术,在多个任务中取得了显著成果。VQ-VAE(向量量化变分自编码器)在保留变分自编码器(VAE)的编码器-解码器结构的基础上,提出了用于处理离散潜在表示的方法。通过VQ-VAE学习离散潜在变量模型在计算机视觉、音频、语音以及其他模态中显示出了很好的效果,因为数据本身就是离散的,这些离散表示自然适用于学习复杂的推理和预测。最近,通过深度生成模型的应用,如生成对抗网络(GANs)和去噪扩散概率模型(DDPMs),VQ-based的离散表示学习也取得了进展。
随着VQ-VAE被集成到多种深度生成模型框架中,其在各种任务中的实用性和适用性日益显现。然而,尽管取得了成功,基于代码本的量化过程的可扩展性仍然是一个重大挑战。随着大规模数据集的普及和对实时处理能力的需求,传统的VQ-VAE架构面临着在处理动态压缩需求时的计算复杂度限制,包括在用户希望增加或减少计算负载时需要重新训练模型。因此,解决VQ过程的可扩展性已成为实现VQ-VAE潜力的关键挑战,并使其能够与最先进的生成模型进行集成。
为了解决这些问题,Li等人提出了一种方法,通过应用双曲嵌入(hyperbolic embeddings)来增强代码本向量,结合共现信息并通过Hilbert曲线重新排列改进的代码本,从而无需重新训练公开可用的VQ模型。为了实现更全面的代码本表示,使用多代码本一直是一个挑战,旨在为不同任务实现更丰富的表示。Malka等人基于渐进学习设计并学习了一个嵌套代码本,以支持不同量化级别。Guo等人提出了一种框架,通过重构教师标签生成作为知识蒸馏中的编解码问题,利用多代码本向量量化预测通过学生模型的嵌入生成的代码本索引。最近,Huijben等人集中研究了基于残差量化的无监督代码本生成,研究了向量量化器本身。然而,通过多代码本或残差量化来解决这些问题通常需要对现有的VQ-VAE结构进行重大更改,或在量化特征图的分辨率上做出妥协。
为此,我们提出了一种速率自适应VQ-VAE(RAQ-VAE)框架,它允许使用单个VQ-VAE模型在不同速率下进行离散表示。首先,我们提出了基于模型的RAQ-VAE,它可以使用现有的VQ-VAE模型,通过一个差分k-means聚类(DKM)算法及其反向功能化,获得速率自适应的VQ,而无需任何额外的参数和重新训练。接着,我们提出了基于数据驱动的RAQ-VAE,它结合Seq2Seq模型与VQ-VAE架构来生成速率自适应的代码本。数据驱动的RAQ-VAE能够通过Seq2Seq模型实现任何所需速率的离散表示,并且在一定程度上超越了独立训练的传统VQ-VAE模型。我们的框架解决了在大规模计算机视觉任务中,尤其是在需要高容量表示的任务中,训练不同压缩率的多个VQ-VAE模型的问题。此外,它可以无缝集成到各种VQ应用中,而无需对现有的VQ-VAE结构进行重大修改。
本文贡献总结如下:
我们引入了RAQ-VAE框架,提供了两种VQ代码本表示方法:基于模型的RAQ-VAE,利用现有训练好的VQ-VAE模型;以及基于数据驱动的RAQ-VAE,结合Seq2Seq模型与VQ-VAE架构。
我们提出了基于模型的RAQ-VAE,通过动态代码本聚类方法来适应VQ-VAE模型的代码本,使量化器能够在不重新训练的情况下调整速率。
我们提出了基于数据驱动的RAQ-VAE,通过Seq2Seq模型生成速率自适应的代码本。该方法使用单一代码本和训练方法“交叉强制”训练递归网络生成不同速率的代码本。
我们的实验表明,单个RAQ-VAE模型的性能与多个传统VQ-VAE模型(按固定速率训练)相当,甚至超越了它们,并且使用相同的编码器-解码器架构。
2.背景
VQ-VAE 通过将连续的潜在变量离散化为其在代码本中的最近邻向量,有效地表示了跨越数据空间多个维度的有意义特征。在 VQ-VAE 中,离散表示的学习是通过将编码后的潜在变量量化为其在可训练代码本中的最近邻向量,并从离散的潜在变量解码输入数据来实现的。对于数据集
D
D
D中的离散表示
x
∈
D
x \in D
x∈D,使用一个包含
K
K
K个可学习代码向量
{
e
i
}
i
=
1
K
\{ e_i \}_{i=1}^{K}
{ei}i=1K的代码本
e
e
e。量化后的离散潜在变量
z
q
(
x
∣
e
)
z_q(x | e)
zq(x∣e)被解码以重建数据
x
x
x。量化器
Q
Q
Q被建模为确定性的分类后验,它通过一个确定性的编码器
f
ϕ
f_\phi
fϕ将数据
x
x
x的连续潜在表示
f
ϕ
(
x
)
f_\phi(x)
fϕ(x)映射到
z
q
(
x
∣
e
)
z_q(x|e)
zq(x∣e),通过在
D
D
D维代码本
e
=
{
e
i
}
i
=
1
K
e = \{ e_i \}_{i=1}^{K}
e={ei}i=1K中找到最近邻来实现。
量化表示被固定为索引
i
i
i的代码向量
e
i
e_i
ei的大小
K
K
K的代码本的
l
o
g
2
K
log_2 K
log2K位。在确定性解码器
f
θ
f_\theta
fθ中,量化的离散潜在变量
z
q
(
x
∣
e
)
z_q(x | e)
zq(x∣e) 被解码为数据
x
x
x,即
x
^
=
f
θ
(
z
q
(
x
∣
e
)
∣
e
)
\hat{x} = f_\theta(z_q(x | e) | e)
x^=fθ(zq(x∣e)∣e)。在训练过程中,编码器
f
ϕ
f_\phi
fϕ、解码器
f
θ
f_\theta
fθ和代码本
e
e
e被联合优化,以最小化损失函数
L
V
Q
(
ϕ
,
θ
,
e
;
x
)
=
L_{VQ}(\phi, \theta, e; x)=
LVQ(ϕ,θ,e;x)=
其中
sg
[
⋅
]
\text{sg}[·]
sg[⋅]是停止梯度操作符。
L
recon
L_{\text{recon}}
Lrecon是输入数据
x
x
x和解码器重建输出
x
^
\hat{x}
x^之间的重建损失。两个损失函数
L
embed
L_{\text{embed}}
Lembed和
L
commit
L_{\text{commit}}
Lcommit仅作用于代码本变量和编码器权重,并使用一个加权超参数
β
\beta
β,以防止从一个代码向量到另一个代码向量的波动。由于量化过程是不可微的,代码本损失通常通过直通梯度估计器(straight-through gradient estimator) 进行近似,例如
∂
L
∂
f
ϕ
(
x
)
≈
∂
L
∂
z
q
(
x
)
\frac{\partial L}{\partial f_\phi(x)} \approx \frac{\partial L}{\partial z_q(x)}
∂fϕ(x)∂L≈∂zq(x)∂L。
传统的变分自编码器(VAE)和 VQ-VAE 的目标函数都由重建误差和潜在正则化项之和组成。为了提高性能和收敛速度,通常会应用指数移动平均(EMA)更新方法来优化代码本。
Seq2Seq
序列到序列(Seq2Seq)模型广泛应用于序列预测任务,如语言建模和机器翻译。该模型采用初始的 LSTM(称为编码器)按顺序处理输入序列,并生成一个固定维度的向量表示,称为上下文向量。然后,输出序列由另一个 LSTM(即解码器)生成。特别地,解码器是基于输入序列进行条件化的,这使得它在架构中作为一个独立的组件。在训练过程中,Seq2Seq 模型通常使用教师强制(teacher forcing),即在每个时间步提供目标序列给解码器,而不是让解码器使用自己上一步的输出作为输入。这种方法通过在训练过程中提供正确的上下文,帮助模型更快地收敛。
3.Rate-Adaptive VQ-VAE
尽管 VQ-VAE 已成功应用于多个领域,但它仍然面临可扩展性问题。特别是,常见的固定速率 VQ-VAE 模型在处理不同数据集时需要修改代码本的大小 K K K。此外,调整计算负载需要重新训练模型,这带来了额外的挑战。为了克服这些限制,我们提出了两个 Rate-Adaptive VQ-VAE(RAQ-VAE)框架,它们可以通过增加或减少代码本大小 K K K来调整 VQ-VAE 的速率。RAQ-VAE 框架的大纲如图 1 所示。RAQ-VAE 基于一个代码本映射 Ψ : ( R D ) × K → ( R D ) × K ~ \Psi : (\mathbb{R}^D) \times K \rightarrow (\mathbb{R}^D) \times \tilde{K} Ψ:(RD)×K→(RD)×K~,对于任何整数 ( K ~ ∈ N \tilde{K} \in \mathbb{N} K~∈N),该值可以比原始代码本大小 K K K小(即 K ~ < K \tilde{K} < K K~<K)或大(即 K ~ > K \tilde{K}>K K~>K)。我们设计了两种映射方式:(i)基于模型的 RAQ-VAE;(ii)基于数据的 RAQ-VAE。基于模型的 RAQ-VAE(第 3.1 节)通过可微分的 k-means 聚类(DKM)[3] 算法获得速率自适应 VQ,无需任何额外的参数。此外,基于数据的 RAQ-VAE(第 3.2 节)是一种离线训练的 RAQ-VAE 方法,它采用了代码本生成的序列到序列(Seq2Seq) 模型。
3.1 基于模型的 Rate-Adaptive VQ-VAE
之前的尝试提出了通过在模型训练过程中周期性地对代码本进行聚类来增强代码本学习。相比之下,我们提出了一种基于模型的速率自适应 VQ-VAE,它在模型训练后执行在线代码本聚类。通过加载一个使用原始代码本
e
e
e训练的 VQ-VAE 模型,并动态地将代码本聚类到适应后的代码本
e
~
\tilde{e}
e~,这使得向量量化器能够适应整体模型中的细微模式,从而提供灵活性和可扩展性(见图 1a)。
代码本聚类
为了实现适应后的代码本大小
K
~
\tilde{K}
K~(即 (
∣
e
~
∣
|\tilde{e}|
∣e~∣)所需的速率,我们从原始代码本
e
e
e推导出聚类后的代码本
e
~
\tilde{e}
e~。代码本聚类的详细公式在补充材料 A.2 中提供。为了确保聚类过程能有效地融入到训练好的 VQ-VAE 模型中,我们采用了可微分的 k-means 聚类(DKM)算法。该算法最初是为了 DNN 模型压缩而提出的,使用基于注意力的权重聚类方法。我们将 DKM 算法适应于 VQ 代码本聚类,重点关注聚类代码本和 VQ-VAE 模型架构的微调。此外,我们利用 DKM 进行代码本增量化(逆函数化过程)来处理需要增加代码本大小的场景。
减少码率 (
K
~
<
K
\tilde{K} < K
K~<K)
在码率降低任务中,可微分 k-means 聚类 (DKM)可以在
K
~
\tilde{K}
K~个簇上执行迭代的可微分码本聚类。设C代表聚类中心,e代表原始码本。针对 VQ 码本的 DKM 算法操作如下:
初始化质心
C
=
{
c
j
}
j
=
1
K
~
C = \{c_j\}_{j=1}^{\tilde{K}}
C={cj}j=1K~,通过随机选择
K
~
\tilde{K}
K~ 个码本向量从
e
e
e初始化,或使用k-means++进行初始化;在所有后续迭代中,使用上一批次的最终
C
C
C作为起点。
计算原始码本向量
e
i
e_i
ei 与初始化质心
c
j
c_j
cj之间的欧几里得距离,以欧几里得距离作为度量标准,计算距离矩阵
d
i
j
=
−
f
(
e
i
,
c
j
)
d_{ij} = -f(e_i, c_j)
dij=−f(ei,cj),其中
D
D
D是所有
d
i
j
d_{ij}
dij组成的矩阵。
计算注意力矩阵
A
A
A,每一行
A
A
A 按如下方式计算,
a
i
,
j
=
exp
(
d
i
,
j
/
τ
)
∑
k
exp
(
d
i
,
k
/
τ
)
a_{i,j} = \frac{\exp(d_{i,j} / \tau)}{\sum_k \exp(d_{i,k} / \tau)}
ai,j=∑kexp(di,k/τ)exp(di,j/τ),其中
τ
\tau
τ是softmax 温度参数,表示
e
i
e_i
ei对
c
j
c_j
cj的注意力分布。
计算新的聚类中心候选
C
~
\tilde{C}
C~并更新
C
C
C,计算新的质心候选
c
~
j
\tilde{c}_j
c~j,
c
~
j
=
∑
i
a
i
,
j
e
i
∑
i
a
i
,
j
\tilde{c}_j = \frac{\sum_i a_{i,j} e_i}{\sum_i a_{i,j}}
c~j=∑iai,j∑iai,jei。
重复此过程,直到收敛,当
∣
C
−
C
~
∣
≤
ϵ
|C - \tilde{C}| \leq \epsilon
∣C−C~∣≤ϵ时,表明 DKM已收敛;或者达到最大迭代次数时,停止训练;计算
A
C
A_C
AC以获得新的压缩码本
e
~
\tilde{e}
e~。
上述迭代过程可总结如下:
在文献 [3] 中,作者实现了用于软权重聚类分配的 DKM,并可以通过施加约束来增强聚类的硬度。在 softmax 操作中,温度参数
τ
τ
τ 可用于控制聚类的硬度水平。在 DKM 过程结束时,我们使用最后的注意力矩阵
A
A
A,将每个码本向量映射到最近的聚类中心,从而完成码本的聚类。
增加码本大小(
K
~
>
K
\tilde{K}>K
K~>K)
虽然k-means 聚类在压缩码本向量方面表现良好,但它在算法上存在一定的限制,无法有效增加额外的码本。为了解决这一问题,我们提出了逆函数 DKM(IKM),这是一种用于增加码本大小的技术。
该迭代方法旨在逼近已生成码本的后验分布,并使用最大均值差异(MMD)来比较基础码本与聚类生成码本之间的分布差异。MMD 是一种 基于核函数的统计检验方法,可用于测量两个分布之间的相似性。
假设原始码本向量
e
e
e具有
K
K
K个向量,并已在基础VQ-VAE中训练完成,使用IKM 算法生成码本𝑒̃的过程如下:
初始化一个D维的自适应码本向量
e
~
=
{
e
~
i
}
i
=
1
K
~
\tilde{e} = \{\tilde{e}_i\}_{i=1}^{\tilde{K}}
e~={e~i}i=1K~,并从标准正态分布中采样:
e
~
∼
N
(
0
,
D
−
1
2
I
)
\tilde{e} \sim \mathcal{N} \left(0, D^{-\frac{1}{2}} I \right)
e~∼N(0,D−21I)。
使用 DKM 过程对
e
~
\tilde{e}
e~进行聚类(参考方程3):
g
DKM
(
e
~
)
=
arg
min
g
DKM
(
e
~
)
L
DKM
(
e
~
;
g
DKM
(
e
~
)
)
g_{\text{DKM}}(\tilde{e}) = \arg \min_{g_{\text{DKM}}(\tilde{e})} L_{\text{DKM}}(\tilde{e}; g_{\text{DKM}}(\tilde{e}))
gDKM(e~)=argmingDKM(e~)LDKM(e~;gDKM(e~))
计算 MMD,衡量原始码本
e
e
e与 DKM 生成的
g
DKM
(
e
~
)
g_{\text{DKM}}(\tilde{e})
gDKM(e~)之间的分布差异。
优化
e
~
\tilde{e}
e~以最小化 MMD 目标函数:
L
IKM
(
e
;
e
~
)
=
M
M
D
(
e
,
g
DKM
(
e
~
)
)
+
λ
∥
e
~
∥
2
2
L_{\text{IKM}}(e; \tilde{e}) = MMD(e, g_{\text{DKM}}(\tilde{e})) + \lambda \|\tilde{e}\|^2_2
LIKM(e;e~)=MMD(e,gDKM(e~))+λ∥e~∥22, 其中,
λ
\lambda
λ 是正则化参数,用于控制
L
2
L_2
L2正则项的强度。
由于 DKM 不会阻塞梯度流,我们可以使用随机梯度下降(SGD)来更新码本
e
~
\tilde{e}
e~,其更新方式为:
e
~
=
e
−
η
∇
L
IKM
(
e
,
e
~
)
\tilde{e} = e - \eta \nabla L_{\text{IKM}}(e, \tilde{e})
e~=e−η∇LIKM(e,e~),其中,
η
\eta
η表示学习率。
借助DKM 和 IKM,生成的码本
e
~
\tilde{e}
e~可用于以不同的码率对编码向量进行量化,即
z
q
(
x
∣
e
~
z_q(x|\tilde{e}
zq(x∣e~),而无需向已训练的VQ-VAE添加额外的模型参数。由于DKM 不会阻塞梯度流,因此即使在离线训练和在线训练过程中,也可以轻松更改码本的聚类分配。
在离线训练(Offline Training)期间,可以选择最适合 VQ 任务损失的聚类并进行采用。
尽管本研究不专注于使用 DKM 进行多码本(Multi-Codebook)建模(主要目的是在训练后利用自适应码本),但仍然可以在训练过程中调整
K
~
\tilde{K}
K~以动态优化码本大小,并通过分层优化(Hierarchical Optimization)使VQ-VAE适应多码本聚类(Multi-Codebook Clustering)。
3.2数据驱动的自适应码率 VQ-VAE
Seq2Seq 模型 在机器翻译任务中得到了广泛应用,因为不同语言的句子长度可能会有较大差异。受此启发,我们提出了一种基于Seq2Seq 的方法,在 VQ-VAE 框架内 生成自适应码本(Rate-Adaptive Codebook)。本节介绍 数据驱动的RAQ-VAE,该方法结合了 学习向量量化(Learning Vector Quantization)层和Seq2Seq模型。
概述 如图1b所示,数据驱动的RAQ-VAE是利用确定性编码器-解码器对、可训练的原始码本e和Seq 2Seq模型来构造的。通过Seq2Seq模型从原始码本
e
e
e生成自适应码本
e
e
e。数据驱动的 RAQ-VAE 通过
e
e
e和
e
~
\tilde{e}
e~以层次化方式对数据
x
x
x的连续潜在表示
f
ϕ
(
x
)
f_{\phi}(x)
fϕ(x)进行量化,分别得到
z
q
(
x
∣
e
)
z_q(x|e)
zq(x∣e)和
z
q
(
x
∣
e
~
)
z_q(x|\tilde{e})
zq(x∣e~)。基于传统的 VQ-VAE 结构,数据驱动的 RAQ-VAE 在训练码本
e
e
e 及其生成过程
G
ψ
G_{\psi}
Gψ的同时,学习编码器-解码器对。
码本编码
速率自适应码本生成过程
G
ψ
G_{\psi}
Gψ利用 Seq2Seq 模型中的 LSTM 单元,从原始码本
e
e
e动态生成适应性码本
e
~
\tilde{e}
e~ 。 首先,初始化目标码本大小
e
K
e_K
eK。在训练过程中,数据驱动的 RAQ-VAE 可以在任意码本大小
e
K
e_K
eK下进行训练。而在测试阶段,Seq2Seq 模型按照用户指定的速率生成适应性码本
e
~
\tilde{e}
e~ 。该初始化为算法 1 中的编码和解码步骤奠定了基础。 在具体实现中,原始码本
e
e
e 的每个向量
e
i
e_i
ei 依次由一组 LSTM 单元进行编码。LSTM 的隐藏状态和记忆单元状态
(
h
,
c
)
(h, c)
(h,c) 用于捕捉每个基础码本向量的上下文信息,从而构建更加适应不同速率需求的量化表征。
训练过程
为了训练数据驱动的RAQ-VAE,我们联合优化基本的VQ-VAE和RAQ-VAE目标,以学习原始码本
e
e
e 和速率自适应码本生成过程
G
ψ
G_{\psi}
Gψ的良好表示。我们将约束优化问题
L
RAQ
L_{\text{RAQ}}
LRAQ公式化,以联合更新
G
ψ
G_{\psi}
Gψ和
f
ϕ
f_{\phi}
fϕ、
f
θ
f_{\theta}
fθ、以及
e
e
e 作为
L
VQ
(
ϕ
,
θ
,
e
;
x
)
≥
L
RAQ
(
ϕ
,
θ
,
ψ
,
e
;
x
)
=
L_{\text{VQ}} (\phi, \theta, e; x) \geq L_{\text{RAQ}} (\phi, \theta, \psi, e; x)=
LVQ(ϕ,θ,e;x)≥LRAQ(ϕ,θ,ψ,e;x)=
其中,
s
g
[
⋅
]
sg[·]
sg[⋅]是停止梯度操作符。数据驱动的 RAQ-VAE 联合最小化
L
VQ
L_{\text{VQ}}
LVQ(方程 2)和
L
RAQ
L_{\text{RAQ}}
LRAQ (方程 5)。反向传播
L
VQ
L_{\text{VQ}}
LVQ 会引发与基础 VQ-VAE 相同的梯度流。此外,反向传播
L
RAQ
L_{\text{RAQ}}
LRAQ 会引发 Seq2Seq 模型的梯度流,从而实现有效的码本生成。所提数据驱动 RAQ-VAE 的整体训练过程总结在算法 2 中。在训练过程中,Seq2Seq 模型动态生成码本,并在每次训练迭代中适应不同的速率。
4.实验
为了展示所提出的 RAQ-VAE 的进步,我们将传统的 VQ-VAE 和二级层次结构的 VQ-VAE(VQ-VAE-2)作为基线。我们在视觉数据集上进行经验评估:CIFAR10(32×32) 和 CelebA(64×64)进行定量评估,ImageNet(256×256) 进行定性评估。我们设计 RAQ-VAE 以适应传统的 VQ-VAE 及其改进的模型结构 ,以在单一模型中实现多个速率。
我们对所有方法使用相同的架构和参数,将默认的码字(离散潜变量)嵌入维度
D
D
D 设置为 64 用于 CIFAR10 和 CelebA,设置为 128 用于 ImageNet。码本大小范围为:CIFAR10 从 16 到 1024,CelebA 从 32 到 2048,ImageNet 从 128 到 4096,传统的 VQ-VAE 模型在2的幂大小上训练,RAQ-VAE 模型则设置为该范围的中间值,适用于基于模型和数据驱动的方法。
评估指标
使用峰值信噪比(PSNR)、结构相似性指数(SSIM)、rFID和码本困惑度对我们的方法进行了定量评估。PSNR 测量信号的最大可能功率与影响数据保真度的噪声功率之间的比率。SSIM 评估两幅图像之间的结构相似性 [52]。rFID 评估在测试数据上的重建图像质量,并与人类感知判断高度相关。码本困惑度定义为:
e
−
∑
i
p
e
i
log
p
e
i
e^{-\sum_{i} p_{e_i} \log p_{e_i}}
e−∑ipeilogpei,其中
p
e
i
=
N
e
i
∑
j
N
e
j
p_{e_i} = \frac{N_{e_i}}{\sum_j N_{e_j}}
pei=∑jNejNei ,
N
e
i
N_{e_i}
Nei 表示与码本
e
i
e_i
ei 的潜在表示相对应的编码数量,当困惑度值达到码本大小
K
~
\tilde{K}
K~时,表示均匀的先验分布,意味着所有码本被平等使用。
4.1 视觉任务的主要结果
定量评估
我们将 RAQ-VAE 模型与传统的 VQ-VAE [48] 进行了图像重建性能的经验比较。我们训练并评估了每个不同码本大小
K
K
K的 VQ-VAE 作为定量基准,然后通过调整速率(通过调整
K
~
\tilde{K}
K~在单个基于模型和数据驱动的 RAQ-VAE 模型上验证了 RAQ-VAE。图 2 显示了在 CIFAR10 和 CelebA 数据集上评估的结果。在相同的压缩率和网络架构下,所有提出的 RAQ-VAE 模型的性能接近于多个 VQ-VAE 模型的性能。当增加速率(码本大小)时,数据驱动的 RAQ-VAE 在 PSNR 和 SSIM 指标上略低,但在 rFID 分数上显著更好,rFID 评估数据集级别的感知图像质量。特别地,传统 VQ-VAE 模型在 CelebA 上的困惑度得分较低,但提出的基于数据的 RAQ-VAE 在困惑度和 rFID 上表现更好,尤其是在高每像素比特(bpp)时。基于模型的 RAQ-VAE 整体表现较差,但在降低速率的任务中,在CIFAR10 上间歇性地取得了更可靠的结果。我们提出的方法具有高度的可移植性,考虑到每个单独的 VQ-VAE 模型所投入的资源,RAQ-VAE 只需一个模型即可覆盖多个固定速率的 VQ-VAE 模型。
定性评估
在定性评估中,我们对比了在不同速率(0.4375 bpp 到 0.75 bpp)下训练的 VQ-VAE 和单个数据驱动的 RAQ-VAE 在 ImageNet (256 × 256) 上的表现。在图 3 中,第一行展示的 VQ-VAE 是针对每个速率单独训练的,结果表明,随着速率的降低,图像质量下降,这与定量评估中观察到的结果一致。
如果我们随机选择在
K
=
4096
K = 4096
K=4096训练的 VQ-VAE 码本(第二行)来降低速率,相较于原始图像,颜色会发生变化,并且当 bpp 低于 0.5 时,图像的形状基本保持不变,但颜色偏向红色系。然而,数据驱动的 RAQ-VAE(第三行),在低 bpp
K
=
512
K = 512
K=512的基础码本(0.5625 bpp)上训练,仅需一个模型即可较好地保留输入图像的高级语义特征和颜色。特别是,它比随机选择码本的方法更好地恢复了图像的细节(如猫的胡须)。
尽管在最低 bpp 下,与原始图像相比,整体图像质量和色调有所下降,但如果离散表示与最先进的 PixelCNN 或 PixelSNAIL 作为先验估计器相结合,并输入 VQ 解码器,则可以增强图像生成质量,这与大多数 VQ-VAE 研究 [的结论一致。
4.2 详细分析
码本可用性
参考先前的研究,我们观察到,随着码本大小的增加,数据驱动的 RAQ-VAE 的码本困惑度(perplexity)也随之增加,从而提升了重建性能。在大多数 VQ-VAE 框架中,当码本困惑度接近码本大小时,即在码本大小有限的情况下能够有效利用可用资源时,被认为是最优的。如定量评估结果所示(见图 2),在较高的 bits per pixel (bpp) 设置下,数据驱动的 RAQ-VAE 在码本困惑度方面优于传统 VQ-VAE。这一改进突显了 Seq2Seq 模型在生成码本方面的有效性,使解码器能够始终高效地利用码本。数据驱动的 RAQ-VAE 具备维持较高码本困惑度的能力,这不仅确保了更好的特征表示和重建质量,也证明了其在处理大规模码本时的稳健性。
适用性
为了展示我们方法的更广泛适用性,我们将其扩展至两级层次化 VQ-VAE(VQ-VAE-2)。VQ-VAE-2 在原始 VQ-VAE 框架的基础上引入了层次结构,使其具备更强的表示能力和重建能力。图 4 展示了以 VQ-VAE-2 作为基线模型的重建性能。从结果可以看出,在 CelebA 数据集的多个码率下,数据驱动的 RAQ-VAE 模型均明显优于原始 VQ-VAE-2。
特别地,在相同的原始码本大小下,VQ-VAE-2 的参数量约为 VQ-VAE 的 2 倍,但由于其层次结构,它能够捕获数据中更细微的特征,从而提升整体重建质量。当我们将速率自适应量化(RAQ)方法应用于 VQ-VAE-2 后,发现它不仅保持了层次化模型的高性能,同时还提供了无需重新训练整个模型即可适应不同码率的灵活性。这表明,我们的方法不仅适用于 VQ-VAE,还能够有效扩展至更高级的版本(如 VQ-VAE-2),展现了其在各种 VQ-VAE 框架中的通用性和稳健性。RAQ-VAE 框架能够无缝集成并扩展至 VQ-VAE-2 等模型,这进一步证明了其在数据重建和生成任务中的广泛应用潜力。
5 .结论
本文提出了速率自适应 VQ-VAE(RAQ-VAE)框架,旨在解决传统 VQ-VAE 在可扩展性方面的局限性,使其能够动态调整码本大小。实验结果表明,RAQ-VAE 在多个码率下均能实现优异的重建性能,而无需重新训练。这一速率自适应特性为需要动态压缩级别的应用(如可变码率的图像/视频压缩和实时端到端通信系统 )提供了更大的灵活性。 凭借其已验证的通用性,RAQ-VAE 具有推动机器学习在理论和实践两方面进一步发展的潜力。