[论文解读] Fast Relative Entropy Coding with A* coding

Gergely Flamich, Stratis Markou+1 more|arXiv (Cornell University)|Jan 30, 2022

Generative Adversarial Networks and Image Synthesis被引用 6

一句话总结

本文提出 AS* 和 AD* 编码，这是一种基于 A* 采样的相对熵编码（REC）算法，可在密度比单峰的连续分布上实现 $ϵ(D_{∞}[Q\|P])$ 的期望运行时间，同时保持 $ϵ(D_{∕}[Q\|P])$ 的期望码长。该方法通过所提出的 IsoKL VAE（IKVAE）实现了高效、可微分的压缩，在 MNIST 上实现了接近最优的压缩效果，且无需量化。

ABSTRACT

Relative entropy coding (REC) algorithms encode a sample from a target distribution $Q$ using a proposal distribution $P$, such that the expected codelength is $\mathcal{O}(D_{KL}[Q \,||\, P])$. REC can be seamlessly integrated with existing learned compression models since, unlike entropy coding, it does not assume discrete $Q$ or $P$, and does not require quantisation. However, general REC algorithms require an intractable $Ω(e^{D_{KL}[Q \,||\, P]})$ runtime. We introduce AS* and AD* coding, two REC algorithms based on A* sampling. We prove that, for continuous distributions over $\mathbb{R}$, if the density ratio is unimodal, AS* has $\mathcal{O}(D_{\infty}[Q \,||\, P])$ expected runtime, where $D_{\infty}[Q \,||\, P]$ is the Rényi $\infty$-divergence. We provide experimental evidence that AD* also has $\mathcal{O}(D_{\infty}[Q \,||\, P])$ expected runtime. We prove that AS* and AD* achieve an expected codelength of $\mathcal{O}(D_{KL}[Q \,||\, P])$. Further, we introduce DAD*, an approximate algorithm based on AD* which retains its favourable runtime and has bias similar to that of alternative methods. Focusing on VAEs, we propose the IsoKL VAE (IKVAE), which can be used with DAD* to further improve compression efficiency. We evaluate A* coding with (IK)VAEs on MNIST, showing that it can losslessly compress images near the theoretically optimal limit.

研究动机与目标

• 解决一般相对熵编码（REC）算法运行时间过高的问题，其运行时间随 KL 散度呈指数增长。
• 克服在训练时使用连续近似而推理时使用离散量化所导致的量化失配问题。
• 设计具有可证明快速运行时间与低码长的 REC 算法，适用于连续分布，尤其适用于变分自编码器（VAEs）。
• 通过引入具有固定码长的分块编码，减少编码可变长度码所带来的码长开销。
• 提出 IsoKL VAE（IKVAE），通过 KL 散度与 Rényi ∞-散度参数化潜在分布，实现高效分块编码，支持 DAD*。

提出的方法

• 提出 AS* 和 AD* 编码，基于 A* 采样的 REC 算法，可高效地从提议分布 P 中采样，以编码来自目标分布 Q 的样本。
• 证明在单峰密度比假设下，AS* 的期望运行时间为 $ϵ(D_{\infty}[Q\|P])$，其中 $D_{\infty}[Q\|P]$ 为 Rényi ∞-散度。
• 设计 DAD*，即 AD* 的近似变体，通过固定每一块的码长，降低开销，同时保持低偏差与快速运行时间。
• 提出 IsoKL VAE（IKVAE），使用 $D_{\mathrm{KL}}[Q\|P]$ 与 $D_{\infty}[Q\|P]$ 参数化潜在分布，而非均值与方差，从而实现在块之间共享码长。
• 通过无约束重参数化方法，利用指数与双曲正切变换实现参数化，确保对均值与方差的约束得到满足。
• 使用朗伯 W 函数求解高斯分布的方差，并采用 Pade 逼近以在实际中实现数值稳定与高效计算。

实验结果

研究问题

• RQ1 能否利用基于 A* 的采样设计出对连续分布具有可证明快速运行时间的 REC 算法？
• RQ2 在单峰密度比假设下，AS* 是否能实现 $ϵ(D_{\infty}[Q\|P])$ 的运行时间上界？
• RQ3 尽管缺乏正式证明，AD* 是否仍能实现与 AS* 相当的运行时间缩放？
• RQ4 通过使用固定码长的近似 REC，能否在不增加偏差的前提下减少分块编码中的码长开销？
• RQ5 通过使用 $D_{\mathrm{KL}}[Q\|P]$ 与 $D_{\infty}[Q\|P]$ 参数化 VAE 的潜在分布，能否实现与 DAD* 配合的高效、低开销压缩？

主要发现

• 在单峰密度比假设下，AS* 实现了 $ϵ(D_{\mathrm{KL}}[Q\|P])$ 的期望码长与 $ϵ(D_{\infty}[Q\|P])$ 的期望运行时间。
• AD* 基于大量实验结果也表现出 $ϵ(D_{\infty}[Q\|P])$ 的期望运行时间，尽管未提供正式证明。
• DAD* 作为 AD* 的近似变体，保留了 AD* 的快速运行时间，同时通过使用固定码长对块进行编码，降低了码长开销。
• IsoKL VAE（IKVAE）通过使用 $D_{\mathrm{KL}}[Q\|P]$ 与 $D_{\infty}[Q\|P]$ 参数化潜在分布，实现了高效的分块编码，降低了每变量的码长成本。
• 在 MNIST 上，DAD* 与 IKVAE 的结合实现了接近理论最优极限的无损压缩，证明了其实际效率。
• 结合 Pade 逼近的朗伯 W 函数可实现高斯参数化中数值稳定且高效的方差计算。