Masked Generative Modeling with Enhanced Sampling Scheme

TL;DR

This paper introduces an Enhanced Sampling Scheme (ESS) for masked generative models that improves sample diversity and fidelity through a three-stage process, demonstrating significant performance gains across multiple datasets.

Contribution

The paper proposes a novel ESS method that explicitly balances diversity and fidelity in masked generative modeling, addressing limitations of previous sampling schemes.

Findings

ESS outperforms existing methods in diverse datasets

Significant improvements in sample quality and realism

Effective in both unconditional and class-conditional sampling

Abstract

This paper presents a novel sampling scheme for masked non-autoregressive generative modeling. We identify the limitations of TimeVQVAE, MaskGIT, and Token-Critic in their sampling processes, and propose Enhanced Sampling Scheme (ESS) to overcome these limitations. ESS explicitly ensures both sample diversity and fidelity, and consists of three stages: Naive Iterative Decoding, Critical Reverse Sampling, and Critical Resampling. ESS starts by sampling a token set using the naive iterative decoding as proposed in MaskGIT, ensuring sample diversity. Then, the token set undergoes the critical reverse sampling, masking tokens leading to unrealistic samples. After that, critical resampling reconstructs masked tokens until the final sampling step is reached to ensure high fidelity. Critical resampling uses confidence scores obtained from a self-Token-Critic to better measure the realism of sampled tokens, while critical reverse sampling uses the structure of the quantized latent vector space to discover unrealistic sample paths. We demonstrate significant performance gains of ESS in both unconditional sampling and class-conditional sampling using all the 128 datasets in the UCR Time Series archive.