利用旋转技巧重构向量量化
机器学习
2025-03-18 v2 计算机视觉与模式识别
摘要
向量量化变分自编码器旨在将连续输入压缩到离散潜空间，并以最小失真进行重构。其工作原理是维护一组向量（通常称为码本），并将每个编码器输出量化到码本中最近的向量。然而，由于向量量化是不可微的，编码器的梯度会绕过向量量化层，而不是通过直通估计穿过它。这种近似可能是不理想的，因为来自向量量化操作的所有信息都会丢失。在这项工作中，我们提出了一种通过 VQ-VAE 的向量量化层传播梯度的方法。我们通过一个旋转和缩放线性变换，将每个编码器输出平滑地转换为其对应的码本向量，该变换在反向传播过程中被视为常数。因此，编码器输出与码本向量之间的相对幅度和角度被编码到梯度中，随着梯度穿过向量量化层并传回编码器。在 11 种不同的 VQ-VAE 训练范式上，我们发现这种重构改善了重建指标、码本利用率和量化误差。我们的代码可在 https://github.com/cfifty/rotation_trick 获取。
引用
@article{arxiv.2410.06424,
title = {Restructuring Vector Quantization with the Rotation Trick},
author = {Christopher Fifty and Ronald G. Junkins and Dennis Duan and Aniketh Iyengar and Jerry W. Liu and Ehsan Amid and Sebastian Thrun and Christopher Ré},
journal= {arXiv preprint arXiv:2410.06424},
year = {2025}
}