Restructuring Vector Quantization with the Rotation Trick

Vector Quantized Variational AutoEncoders（VQ-VAEs）旨在将连续的输入压缩为离散的潜在空间，并以最小失真重构它。它们通过维护一组向量（通常称为码本）来运作，并将每个编码器输出量化为码本中最近的向量。然而，由于向量量化是不可微分的，因此梯度流绕过向量量化层而不是直接通过它进行直通逼近。这种逼近可能是不可取的，因为向量量化操作的所有信息都丢失了。在这项工作中，我们提出了一种通过VQ-VAE的向量量化层传播梯度的方法。我们通过旋转和重新缩放的线性变换，将每个编码器输出平滑地转换为其相应的码本向量，该变换在反向传播过程中被视为常数。因此，编码器输出和码本向量之间的相对大小和角度被编码到梯度中，随着梯度通过向量量化层和返回编码器进行传播。在11种不同的VQ-VAE训练范例中，我们发现这种重构改善了重构指标、码本利用率和量化误差。我们的代码可在https://github.com/cfifty/rotation_trick上找到。

论文旨在解决Vector Quantized Variational AutoEncoders (VQ-VAEs)中向量量化层的反向传播问题，以便更好地利用编码器输出和码本向量之间的信息。

论文提出了一种通过旋转和缩放线性变换平滑地将编码器输出转换为其对应的码本向量的方法，以便在反向传播期间将梯度传播到向量量化层，并在维持压缩的同时实现更好的重构。

论文在11种不同的VQ-VAE训练范式中测试了这种方法，并发现这种重构方式可以提高重构度量、码本利用率和量化误差。代码已开源。

在向量量化和自编码器方面，已有许多相关的研究，例如《Neural Discrete Representation Learning》和《Deep Clustering for Unsupervised Learning of Visual Features》等。