Discrete Variational Autoencoding via Policy Search

📄 arXiv: 2509.24716v1 📥 PDF

作者: Michael Drolet, Firas Al-Hafez, Aditya Bhatt, Jan Peters, Oleg Arenz

分类: cs.LG, cs.AI, cs.RO

发布日期: 2025-09-29

💡 一句话要点

提出基于策略搜索的离散变分自编码器，提升高维图像重建质量。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 离散变分自编码器 策略搜索 自然梯度 图像重建 非参数编码器

📋 核心要点

  1. 现有离散VAE依赖近似重参数化或高方差无梯度方法，在高维图像重建任务中效果有限。
  2. 论文提出一种基于策略搜索的训练框架，利用非参数编码器的自然梯度更新参数化编码器，无需重参数化。
  3. 实验表明，该方法在ImageNet等数据集上优于现有方法，显著提升了图像重建质量，FID得分提升20%。

📝 摘要（中文）

离散变分自编码器(VAE)中的离散潜在瓶颈提供了高比特效率，并且可以使用自回归离散分布进行建模，从而能够利用transformers进行参数高效的多模态搜索。然而，离散随机变量不允许精确的可微参数化；因此，离散VAE通常依赖于近似方法，例如Gumbel-Softmax重参数化或straight-through梯度估计，或者采用高方差的无梯度方法，例如REINFORCE，这些方法在高维任务（如图像重建）上的成功有限。受到策略搜索中常用技术的启发，我们提出了一种离散VAE的训练框架，该框架利用非参数编码器的自然梯度来更新参数化编码器，而无需重参数化。我们的方法结合了自动步长调整和基于transformer的编码器，可以扩展到具有挑战性的数据集（如ImageNet），并且在从紧凑潜在空间重建高维数据方面优于近似重参数化方法和基于量化的离散自编码器，在ImageNet 256上的FID得分提高了20%。

🔬 方法详解

问题定义：论文旨在解决离散变分自编码器（Discrete VAEs）在高维数据（如图像）重建任务中，由于离散潜在变量不可微而导致的训练困难问题。现有方法，如Gumbel-Softmax重参数化和REINFORCE，要么引入近似误差，要么具有高方差，限制了其在高维数据上的表现。

核心思路：论文的核心思路是借鉴策略搜索的思想，将离散VAE的训练过程视为一个策略优化问题。通过利用非参数编码器的自然梯度来更新参数化编码器，避免了直接对离散变量进行梯度估计，从而绕过了不可微的问题。这种方法允许更稳定和高效的训练。

技术框架：整体框架包括一个参数化的编码器（通常是Transformer网络），一个非参数化的编码器，以及一个解码器。训练过程主要分为以下几个步骤：1) 使用非参数编码器生成离散潜在变量；2) 使用参数化编码器预测这些离散潜在变量的概率分布；3) 计算非参数编码器的自然梯度；4) 使用自然梯度更新参数化编码器的参数；5) 使用解码器重建输入数据。

关键创新：最重要的创新点在于使用非参数编码器的自然梯度来指导参数化编码器的训练。这避免了对离散变量进行近似梯度估计，从而减少了误差和方差。此外，结合自动步长调整，可以进一步提高训练的稳定性和效率。

关键设计：关键设计包括：1) 使用Transformer作为参数化编码器，以捕捉图像中的长程依赖关系；2) 使用KL散度作为损失函数，衡量参数化编码器预测的概率分布与非参数编码器生成的离散潜在变量之间的差异；3) 采用自动步长调整算法，动态调整参数化编码器的学习率，以提高训练的稳定性。

📊 实验亮点

实验结果表明，该方法在ImageNet 256数据集上取得了显著的性能提升，FID得分相比现有方法提高了20%。这表明该方法能够更有效地学习图像的潜在表示，并生成更高质量的图像。此外，该方法在训练稳定性和收敛速度方面也优于其他离散VAE方法。

🎯 应用场景

该研究成果可应用于图像压缩、图像生成、图像编辑等领域。通过学习紧凑的离散潜在表示，可以实现高效的图像存储和传输。此外，该方法还可以用于生成具有特定属性的图像，例如，通过控制离散潜在变量来改变图像的风格或内容。未来，该方法有望扩展到其他高维数据领域，如视频和音频。

📄 摘要（原文）

Discrete latent bottlenecks in variational autoencoders (VAEs) offer high bit efficiency and can be modeled with autoregressive discrete distributions, enabling parameter-efficient multimodal search with transformers. However, discrete random variables do not allow for exact differentiable parameterization; therefore, discrete VAEs typically rely on approximations, such as Gumbel-Softmax reparameterization or straight-through gradient estimates, or employ high-variance gradient-free methods such as REINFORCE that have had limited success on high-dimensional tasks such as image reconstruction. Inspired by popular techniques in policy search, we propose a training framework for discrete VAEs that leverages the natural gradient of a non-parametric encoder to update the parametric encoder without requiring reparameterization. Our method, combined with automatic step size adaptation and a transformer-based encoder, scales to challenging datasets such as ImageNet and outperforms both approximate reparameterization methods and quantization-based discrete autoencoders in reconstructing high-dimensional data from compact latent spaces, achieving a 20% improvement on FID Score for ImageNet 256.