Discrete Variational Autoencoding via Policy Search
作者: Michael Drolet, Firas Al-Hafez, Aditya Bhatt, Jan Peters, Oleg Arenz
分类: cs.LG, cs.AI, cs.RO
发布日期: 2025-09-29 (更新: 2026-01-28)
💡 一句话要点
提出基于策略搜索的离散变分自编码器，用于高效高维数据重建
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 离散变分自编码器 策略搜索 自然梯度 非参数编码器 图像重建 深度学习 生成模型
📋 核心要点
- 传统离散VAE依赖近似重参数化或高方差无梯度方法，在高维图像重建任务中表现不佳。
- 论文提出一种基于策略搜索的训练框架，利用非参数编码器的自然梯度更新参数编码器，避免重参数化。
- 实验表明，该方法在ImageNet等数据集上优于现有方法，能从紧凑潜在空间重建高维数据。
📝 摘要（中文）
离散变分自编码器(VAE)中的离散潜在瓶颈提供了高比特效率，并且可以使用自回归离散分布进行建模，从而能够使用transformers进行参数高效的多模态搜索。然而，离散随机变量不允许精确的可微参数化；因此，离散VAE通常依赖于近似，例如Gumbel-Softmax重参数化或straight-through梯度估计，或者采用高方差的无梯度方法，例如REINFORCE，这些方法在高维任务（如图像重建）上的成功有限。受到策略搜索中常用技术的启发，我们提出了一种离散VAE的训练框架，该框架利用非参数编码器的自然梯度来更新参数编码器，而无需重参数化。我们的方法结合了自动步长自适应和基于transformer的编码器，可以扩展到具有挑战性的数据集（如ImageNet），并且在从紧凑潜在空间重建高维数据方面优于近似重参数化方法和基于量化的离散自编码器。
🔬 方法详解
问题定义：离散变分自编码器（Discrete VAEs）旨在学习数据的离散潜在表示，这对于压缩和生成模型具有重要意义。然而，由于离散变量的不可微性，传统的基于梯度下降的训练方法难以直接应用。现有方法，如Gumbel-Softmax重参数化和REINFORCE算法，要么引入近似误差，要么具有高方差，导致训练不稳定，在高维数据（如图像）重建任务中效果不佳。因此，如何有效地训练离散VAE，使其能够从紧凑的离散潜在空间中重建高维数据，是一个亟待解决的问题。
核心思路：该论文的核心思路是借鉴策略搜索的思想，将离散VAE的训练过程视为一个策略优化问题。具体来说，将编码器视为一个策略，其目标是选择最佳的离散潜在变量，以最大化重建数据的概率。通过利用非参数编码器的自然梯度来更新参数编码器，避免了直接对离散变量进行梯度估计的困难。这种方法允许使用更有效的梯度信息，从而提高训练的稳定性和收敛速度。
技术框架：该方法的技术框架主要包括以下几个模块：1) 编码器：使用一个参数化的编码器（例如，基于Transformer的编码器）将输入数据映射到离散潜在空间。2) 非参数编码器：维护一个非参数的编码器，用于估计每个离散潜在变量的概率分布。3) 解码器：使用一个解码器将离散潜在变量映射回原始数据空间。4) 策略搜索：使用策略搜索算法，利用非参数编码器的自然梯度来更新参数编码器。整个训练流程包括：首先，使用参数化编码器对输入数据进行编码，得到离散潜在变量；然后，使用解码器重建数据；接着，使用非参数编码器估计离散潜在变量的概率分布；最后，使用策略搜索算法，根据重建误差和概率分布，更新参数化编码器。
关键创新：该论文最重要的技术创新点在于使用策略搜索的思想来训练离散VAE，并利用非参数编码器的自然梯度来更新参数编码器。与传统的基于重参数化或REINFORCE的方法相比，该方法避免了近似误差和高方差问题，能够更有效地利用梯度信息，从而提高训练的稳定性和收敛速度。此外，结合自动步长自适应和基于Transformer的编码器，使得该方法能够扩展到具有挑战性的数据集（如ImageNet）。
关键设计：该方法的一些关键设计包括：1) 非参数编码器：使用一个基于K近邻（KNN）的非参数编码器来估计离散潜在变量的概率分布。2) 自然梯度：使用Fisher信息矩阵来计算自然梯度，从而更好地适应参数空间的几何结构。3) 自动步长自适应：使用一种自动步长自适应算法来调整参数更新的步长，以提高训练的稳定性。4) 损失函数：使用重建误差和KL散度的加权和作为损失函数，其中KL散度用于约束离散潜在变量的分布。
🖼️ 关键图片
📊 实验亮点
实验结果表明，该方法在ImageNet数据集上优于现有的离散VAE方法，包括基于Gumbel-Softmax重参数化的方法和基于量化的自编码器。具体来说，该方法能够以更紧凑的离散潜在空间重建高维图像，并且重建质量更高。例如，在相同的比特率下，该方法的重建图像的PSNR值比现有方法高出显著幅度。
🎯 应用场景
该研究成果可应用于图像压缩、图像生成、视频编码等领域。通过学习数据的离散潜在表示，可以实现高效的数据压缩和生成。此外，该方法还可以用于无监督表示学习，为下游任务提供更好的特征表示。未来，该方法有望扩展到其他类型的数据，如文本、音频等，并应用于更广泛的领域。
📄 摘要（原文）
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.