Air-Decoding: Attribute Distribution Reconstruction for Decoding-Time Controllable Text Generation (2023)

Abstract: Controllable text generation (CTG) aims to generate text with desired attributes, and decoding-time-based methods have shown promising performance on this task. However, in this paper, we identify the phenomenon of Attribute Collapse for the first time. It causes the fluency of generated text to rapidly decrease when the control strength exceeds a critical value, rendering the text completely unusable. This limitation hinders the effectiveness of decoding methods in achieving high levels of controllability. To address this problem, we propose a novel lightweight decoding framework named Air-Decoding. Its main idea is reconstructing the attribute distributions to balance the weights between attribute words and non-attribute words to generate more fluent text. Specifically, we train prefixes by prefix-tuning to obtain attribute distributions. Then we design a novel attribute distribution reconstruction method to balance the obtained distributions and use the reconstructed distributions to guide language models for generation, effectively avoiding the issue of Attribute Collapse. Experiments on multiple CTG tasks prove that our method achieves a new state-of-the-art control performance.

## introduction controllable text generation (ctg) aims to produce texts with specific attributes(e.g. sentiment, topic, non-toxicity). in this field, mainstream methods using transformer-based casual language models (clms) like gpt-2 @xcite and gpt-3 @xcite have achieved fairly good attribute control. even so, it remains challenging to control the generated text to simultaneously satisfy certain attributes and maintain a reasonable coherence @xcite . current ctg methods can be divided into three categories @xcite . the first category achieves attribute controllability by retraining the whole parameters of clms @xcite . these methods obtain impressive control effects, but as the scale of clms increases, the computational cost becomes excessively large. the second category fine-tunes prefixes or prompts to control clms in generating texts with specified attributes @xcite @xcite . these methods have low computational cost and fast generation speed, but the control prefixes or control prompts absorb the features of the training corpus and are prone to overfitting, resulting in poor generality of the generated text. the third category is the decoding-time approach, which guides the model to generate the desired attribute text by adjusting the output probability distribution of the model during the decoding stage. for example, most methods use an attribute distribution obtained by a well-trained classifier to control the clms' output distribution and add an exponential term to the attribute distribution as a control strength @xcite @xcite . due to no direct fine-tuning of the models used for the generation, these methods exhibit re-markable generalization performance and are easy to achieve high control effectiveness when increasing control strength, which is hardly achievable by the first two methods. however, decoding-time approaches suffer from attribute collapse, which refers to the phenomenon that when the control strength increases to a certain critical value, the fluency of the generated text will rapidly decrease. as illustrated in figure 1 (left), where the fluency of generated texts would severely deteriorate as the control strength increases. we found the reason lies in the imbalanced attribute distribution, where the weights of attribute words are significantly higher than those of non-attribute words. as a result, when given a high control strength, the final output distribution will deviate from the original language model's output distribution and excessively favor the attribute distribution. this leads to the generated tokens favoring specified attributes but ignoring basic semantics and syntax, resulting in generated text with high attribute accuracy but compromised fluency. we propose air-decoding, a novel lightweight decoding framework to address the issue of attribute collapse in traditional decoding-time approaches. our method could obtain a more balanced attribute distribution in a low-resource manner, which helps us achieve better generation results while reducing resource consumption. concretely, we first use a lightweight approach to obtain corresponding conditional language models for each attribute, which serves as our attribute classifier to generate the original word-level attribute distribution. then, we propose an attribute distribution reconstruction method to reconstruct the original attribute distribution, which avoids exaggerating the weights of attribute words and excessively diminishing the weight of non-attribute words, resulting in a more balanced attribute distribution. as shown in figure 1 (right), after applying our air-decoding framework, the perplexity could remain within a stable range as the control strength increases, proving that generated texts based on the reconstructed attribute distribution ensure high attribute accuracy and good semantic and syntactic structures. our main contributions are as follows: ## related work in recent years, numerous controllable text generation methods based on pre-trained language models have emerged @xcite @xcite . these methods can be roughly categorized into three types: retrain/refact, prefix/prompt-tuning, and decodingtime approach @xcite . retrain, refactor. retraining and refactoring aim to retrain a conditional language model from scratch or change its architecture. @xcite use 55 attribute control codes to finetune a 1.63 billion-parameter transformer to control generation. @xcite added a control block to gpt-2's architecture and retrained the model using self-supervised learning methods with special control codes. @xcite introduce pointer, a modified transformer model for lexically constrained text generation. ## methodology we propose air-decoding, a decoding-time-based ctg framework. its main idea is reconstructing the attribute distributions to make the weights of attribute words and non-attribute words more balanced to generate more fluent text. specifically, we propose prefix-conditional lm (pc-lm) which employs prefix-tuning to enable gpt-2 to acquire attribute distributions. then, we design an attribute distribution reconstruction method to balance the obtained distributions, which will be used to guide a frozen gpt-2. the overall framework is illustrated in figure 2 . ## conclusion we first identify the phenomenon of attribute collapse and propose a novel lightweight framework, air-decoding, to address this issue in decodingtime ctg. specifically, we train prefixes as prefix-conditional lms (pc-lms), followed by a novel attribute distributional reconstruction method to ensure a more balanced attribute distribution obtained by the pc-lms. then we use the reconstructed attribute distribution to guide the generation of clms and achieve remarkable ctg performance. experimental results on three typical ctg tasks demonstrate that our method not only achieves high attribute control but also excellent text fluency, effectively solving the problem of attribute collapse in traditional decoding-time-based methods.