Critic-Guided Decoding for Controlled Text Generation

TL;DR

This paper introduces CriticControl, a novel decoding method that combines reinforcement learning and weighted decoding to improve controlled text generation in terms of coherence, control, and generalization.

Contribution

CriticControl is a new method that trains a language model steering critic using non-differentiable rewards, enhancing control and efficiency in text generation.

Findings

Outperforms previous methods in topic, sentiment, and detoxification tasks.

Generates more coherent and well-controlled texts.

Shows superior zero-shot generalization ability.

Abstract

Steering language generation towards objectives or away from undesired content has been a long-standing goal in utilizing language models (LM). Recent work has demonstrated reinforcement learning and weighted decoding as effective approaches to achieve a higher level of language control and quality with pros and cons. In this work, we propose a novel critic decoding method for controlled language generation (CriticControl) that combines the strengths of reinforcement learning and weighted decoding. Specifically, we adopt the actor-critic framework to train an LM-steering critic from non-differentiable reward models. And similar to weighted decoding, our method freezes the language model and manipulates the output token distribution using called critic, improving training efficiency and stability. Evaluation of our method on three controlled generation tasks, namely topic control, sentiment control, and detoxification, shows that our approach generates more coherent and well-controlled texts than previous methods. In addition, CriticControl demonstrates superior generalization ability in zero-shot settings. Human evaluation studies also corroborate our findings.