Nested-Wasserstein Self-Imitation Learning for Sequence Generation

TL;DR

This paper introduces a nested-Wasserstein self-imitation learning framework that enhances sequence generation by better semantic matching and exploration, outperforming traditional reinforcement learning methods.

Contribution

It proposes a novel nested-Wasserstein distance and self-imitation framework for sequence generation, improving semantic matching and exploration efficiency.

Findings

Consistently improves sequence generation performance

Enhances semantic matching in RL-based models

Effective on both unconditional and conditional tasks

Abstract

Reinforcement learning (RL) has been widely studied for improving sequence-generation models. However, the conventional rewards used for RL training typically cannot capture sufficient semantic information and therefore render model bias. Further, the sparse and delayed rewards make RL exploration inefficient. To alleviate these issues, we propose the concept of nested-Wasserstein distance for distributional semantic matching. To further exploit it, a novel nested-Wasserstein self-imitation learning framework is developed, encouraging the model to exploit historical high-rewarded sequences for enhanced exploration and better semantic matching. Our solution can be understood as approximately executing proximal policy optimization with Wasserstein trust-regions. Experiments on a variety of unconditional and conditional sequence-generation tasks demonstrate the proposed approach consistently leads to improved performance.