[Paper Review] Deterministic Decoding for Discrete Data in Variational Autoencoders

Daniil Polykovskiy, Dmitry Vetrov|arXiv (Cornell University)|Mar 4, 2020

Gaussian Processes and Bayesian Inference30 references4 citations

TL;DR

This paper proposes Deterministic Decoding Variational Autoencoders (DD-VAE), which replace stochastic decoding with argmax-based deterministic generation to improve disentangled representation learning in discrete sequential data. By using bounded support proposal distributions and a continuous relaxation of the argmax, DD-VAE achieves higher reconstruction accuracy and better downstream performance in molecular property optimization compared to standard VAEs.

ABSTRACT

Variational autoencoders are prominent generative models for modeling discrete data. However, with flexible decoders, they tend to ignore the latent codes. In this paper, we study a VAE model with a deterministic decoder (DD-VAE) for sequential data that selects the highest-scoring tokens instead of sampling. Deterministic decoding solely relies on latent codes as the only way to produce diverse objects, which improves the structure of the learned manifold. To implement DD-VAE, we propose a new class of bounded support proposal distributions and derive Kullback-Leibler divergence for Gaussian and uniform priors. We also study a continuous relaxation of deterministic decoding objective function and analyze the relation of reconstruction accuracy and relaxation parameters. We demonstrate the performance of DD-VAE on multiple datasets, including molecular generation and optimization problems.

Motivation & Objective

• To address the issue of latent code ignoring in standard VAEs with flexible stochastic decoders.
• To improve disentangled and structured representation learning in discrete sequential data by enforcing deterministic decoding.
• To develop a continuous relaxation of the argmax decoding objective for differentiable training.
• To demonstrate improved performance in downstream tasks such as molecular generation and Bayesian optimization.
• To provide theoretical analysis of KL divergence and reconstruction under bounded support proposals.

Proposed method

• Proposes a deterministic decoder that selects the highest-scoring token at each step using argmax over the decoder’s output distribution.
• Introduces bounded support proposal distributions (e.g., Gaussian and uniform) to enable lossless auto-encoding, overcoming the limitations of full-support proposals.
• Derives a continuous relaxation of the argmax objective using the Gumbel-Softmax trick to allow backpropagation through the deterministic decoding path.
• Derives the evidence lower bound (ELBO) for DD-VAE and proves convergence of the relaxed objective to the optimal solution of the original problem.
• Uses a temperature hyperparameter τ to balance reconstruction and KL terms, with log-linear annealing shown to improve training stability.
• Employs a sparse Gaussian process on learned latent codes for Bayesian optimization of molecular properties in downstream tasks.

Experimental results

Research questions

• RQ1 Can deterministic decoding in VAEs prevent latent code ignoring and improve representation learning in discrete sequential data?
• RQ2 What type of proposal distribution supports lossless auto-encoding in deterministic VAEs?
• RQ3 How can the argmax operation in deterministic decoding be relaxed for end-to-end training with gradient backpropagation?
• RQ4 Does DD-VAE yield better downstream performance in molecular generation and property optimization than standard VAEs?
• RQ5 How do different prior distributions (Gaussian vs. uniform) and relaxation hyperparameters affect model performance?

Key findings

• DD-VAE achieves 89.89% sequence reconstruction accuracy on the ZINC dataset, outperforming standard VAEs and other baselines in reconstruction and validity.
• On Bayesian optimization for molecular property scoring, DD-VAE with tricube proposal found molecules with a top-1 score of 5.86, significantly outperforming standard VAEs.
• The use of bounded support proposals (e.g., tricube) enables lossless auto-encoding, which is impossible with full-support distributions like Gaussian.
• DD-VAE with a uniform prior leads to more evenly distributed latent codes, improving performance in downstream classification and visualization tasks.
• The continuous relaxation of the argmax objective converges to the optimal solution of the original non-relaxed problem, ensuring training stability and correctness.
• Log-linear annealing of the temperature hyperparameter τ yields better performance than linear annealing, with initial τ chosen to balance gradient scales of reconstruction and KL terms.