An Invariant Learning Characterization of Controlled Text Generation (2023)

Abstract: Controlled generation refers to the problem of creating text that contains stylistic or semantic attributes of interest. Many approaches reduce this problem to training a predictor of the desired attribute. For example, researchers hoping to deploy a large language model to produce non-toxic content may use a toxicity classifier to filter generated text. In practice, the generated text to classify, which is determined by user prompts, may come from a wide range of distributions. In this paper, we show that the performance of controlled generation may be poor if the distributions of text in response to user prompts differ from the distribution the predictor was trained on. To address this problem, we cast controlled generation under distribution shift as an invariant learning problem: the most effective predictor should be invariant across multiple text environments. We then discuss a natural solution that arises from this characterization and propose heuristics for selecting natural environments. We study this characterization and the proposed method empirically using both synthetic and real data. Experiments demonstrate both the challenge of distribution shift in controlled generation and the potential of invariance methods in this setting.

## introduction the development of large language models (llms) has changed the landscape of research in nlp. simply by conditioning on a prompt, an llm can produce fluent and readable text. by using different and well-thought-out prompts, it can be adapted to many applications @xcite @xcite @xcite . but this increase in adaptability has also led to a greater need for controlled generation, to be able to generate text from an llm that adheres to certain attributes. for example, suppose we want to use an llm as a chatbot and deploy it to a large set of users. they might prompt the model in many different ways, such as by asking for advice, information, or just playing with its capabilities. we would like the users to freely explore the chatbot, but we also want to ensure that the text it generates is not toxic -that is, not rude, disrespectful, or unreasonable. how can we allow users to freely prompt it, but ensure that the llm does not produce toxic text? there have been many approaches to solving this problem, each trying to ensure that the text produced by a prompted llm adheres to the attribute, e.g., that it is not toxic @xcite @xcite @xcite . here we build on the simple method of filtering. filtering reduces the problem of controlled generation to one of building a good classifier of the targeted attribute. first we collect a dataset of texts that is labeled as to whether each is toxic, and we use this data to fit a toxicity classifier. when a user prompts the llm to produce a sample of text, we use the fitted classifier to filter its results. we collect multiple texts from the prompted llm, but only retain one that is classified as non-toxic. filtering is a simple and direct approach to controlled generation, but it is only as effective as the fitted classifier. in this paper, we argue that a classifier that might perform well in a classical ml setting will likely perform worse in the context of a prompted llm. the reason is that classical ml tacitly assumes that the future unlabeled text comes from a similar distribution as the training data. but, when used in the context of controlled generation, the unlabeled text to classify may come from any distribution as it is determined by a user's prompt. compounding the problem, we hope the classifier will work well for many different prompts and thus many different distributions of unlabeled texts. in this paper, we characterize controlled text generation as an out-of-distribution generalization problem. this characterization highlights that distribution shift is an inherent aspect of controlled text generation and it suggests that methods addressing out-of-distribution generalization can be used in the context of controlled generation. concretely, we employ recent algorithms for multi-environment learning @xcite @xcite @xcite . these are methods that analyze multiple related datasets, called "environments," to weed out spurious correlations and find patterns that are consistent across distributions of text. we develop two approaches to create these environments from common text classification datasets, and we demonstrate that invariant methods can be effective for controlled text generation. 1 ## characterizing controlled generation in this section, we review controllable text generation and illustrate the problem of distribution shifts in this setting. ## controlled generation with invariant learning section 2 describes how the task of controlled generation reduces to finding a predictor p θ(y|x) to approximate the ground truth relationship between text and attribute, p * (y|x). the predictor p θ(y|x) is typically fitted by minimizing the training distribution risk, however, the predictor p θ(y|x) that is most effective for a deployment distribution p h (y|x) is the minimizer of the deployment distribution risk, thus, for a predictor p θ(y|x) to generalize to many deployment distributions, it should not be trained to minimize the training distribution risk (eq. 5). instead, a good predictor p θ(y|x) should have a low value for r h ( θ) for many prompts h. even if there is only a single deployment distribution of interest, yielding a predictor that performs well for many prompts h will increase the quality of controlled generations for the single prompt. invariant learning. we cast the task of finding a generalizable predictor as an invariant learning problem. invariant learning refers to a class of methods developed to address distribution shifts @xcite @xcite @xcite . these methods posit that features are drawn from multiple distributions, or "environments," but the relationship between label and features is invariant across environments. the motivation is that if a predictor is optimal across environments seen during training, then it will generalize better to future unseen environments. to adapt invariant learning for controlled generation, we note that each deployment distribution p h (x) defines a new environment, indexed by h. since the true relationship between text and attribute p * (y|x) is invariant across distributions of x, the attribute predictor p θ(y|x) should also be invariant in order to generalize to unseen deployment distributions p h (x). the optimal invariant predictor will yield the desired controlled generations p h, θ(x|y) = p * h (x|y). formally, we adapt the data generating process from peters et al. @xcite and arjovsky et al. @xcite for controlled generation: practical optimization. in practice, solving eq. 8 is challenging because each constraint calls an inner optimization @xcite . instead, we find invariant predictors by relying on algorithms developed to approximate eq. 8. these methods add a regularizer to the empirical risk loss (eq. 5) to encourage invariance. see app. a for a description of the three methods we employ in the empirical study. these methods all rely on a hyperparameter, β, that balances the tradeoff between empirical risk and the invariance regularizer. the best way to select this hyperparameter remains an open question @xcite . in section 6, we consider two ways of selecting β. the first is to use a held-out training environment @xcite , while the second relies on samples from the deployment distribution. ## constructing multiple environments invariant learning relies on multiple data environments. in many settings, labeled environments are not available. this section describes how to build environments from passively collected data. recall that a training environment is a collection of data drawn from an environment distribution, we consider two approaches for creating environments. the first uses existing auxiliary labels to split data into environments. the second is a method we propose for creating environments that does not necessarily rely on auxiliary labels. ## related work controlled generation. generating text while controlling for specific attributes is a central problem in nlp @xcite . various approaches include modeling the conditional distribution directly @xcite @xcite ; fine-tuning an existing language model to make use of the observed text and labels @xcite @xcite ; and prompt engineering @xcite . the challenge of modeling the conditional distribution directly is that this limits the use of pre-trained models. there is little theoretical understanding of prompting or fine-tuning, which makes it difficult to predict the robustness of models on unseen data. similar to this paper, another line of work makes use of filtering-based controlled generation (eq. 4) and focuses on training a discriminator p θ(y | x). the discriminator is then used to modify the model activation @xcite or the decoding weights at the token level @xcite @xcite or simply through rejection sampling @xcite . this paper differs from existing work in that we identify a distribution shift problem inherent to prompting that has been overlooked in prior papers. ## experiments we empirically investigate distribution shifts in controlled text generation and assess the effectiveness of invariance methods. this paper studies a filtering-based approach to controlled generation, where each method corresponds to a different classifier. thus, the effectiveness of these methods is determined by the predictive performance of the classifier under distribution shifts. the study includes two settings: an idealized setting involving synthetic data where the distribution shift is known, and another with real world data where a distribution shift is induced but its exact form is unknown. training data and predictors. for both settings, we use training data from civilcomments dataset of comments submitted to an online news platform. the comments are annotated for toxicity and other semantic features such as mention of identity attributes (e.g., race or religion). we compare empirical risk minimization (erm, eq. 5) to invariance-based approaches. in the idealized settings, we use one invariance method, v-rex (eq. 12). in the real world setting, we additionally include mmd @xcite and coral @xcite . we fine-tune bert @xcite on a subset of civilcomments to optimize each objective. dataset, training, and hyperparameter details are in app. b. ## limitations & potential risks there are two main limitations to this work. first, we focus on the "filtering" approach to controlled generation. while this formulation clarifies what a distribution is, it can be computationally expensive to do rejection sampling in practice. a promising area of future research is the application of these invariance principles to the design of large language models. second, achieving true invariance, i.e., generalizing to any arbitrary distribution of text, is a challenging open problem. the purpose of this paper is not to solve this problem. rather, we illustrate that controlled generation is an important application area for invariance methods. an exciting area of future work is to use prompted language models to construct well-defined distribution shift benchmarks for domain generalization methods. controlled text generation has the potential to have large impacts on society, both positive and negative. one potential source of risk is misuse. although we focus on the detection and removal of toxicity, the method we developed can also be applied to the generation of dangerous and toxic content. in addition, this paper does not address other biases (such as gender or social bias) that may already be present in language models. the use of a toxicity filter may compound the problem of decreased diversity in generated text if there is a correlation between social biases and toxicity.