Adversarial Imitation Learning with Controllable Rewards for Text Generation

Abstract

Supervised fine-tuning of large language models (LMs) does not always provide good text-generation performance in terms of quality and diversity. This is because such models maximize the likelihood of correct subsequent words based on previous contexts encountered in the training phase, instead of evaluating the entire structure of the generated texts. In this context, fine-tuning methods for LMs using adversarial imitation learning (AIL) have been proposed to improve the trade-off relationship between quality and diversity. This method leverages the evaluations of the discriminators without requiring manually designed metrics. Previously proposed AIL methods cannot control the shapes of the reward functions and constrain updates of LMs using fixed ranges, independent of the quality, e.g., proximal policy optimization. This study proposes a combination of an AIL method and an approximation of mixture distributions (AMDAIL), synergizing with LMs for text generation. AMDAIL exhibits two features: (1) controlling the distribution of the bounded reward values by varying the shape of the bounded reward function, and (2) a variable constraint to promote updates using the confidence of the discriminator as the quality of the texts. The proposed method exhibits stable behavior in the training phases and improves the trade-off relationship between the quality and diversity in the inference phases.