disco: a toolkit for Distributional Control of Generative Models (2023)

Abstract: Pre-trained language models and other generative models have revolutionized NLP and beyond. However, these models tend to reproduce undesirable biases present in their training data. Also, they may overlook patterns that are important but challenging to capture. To address these limitations, researchers have introduced distributional control techniques. These techniques, not limited to language, allow controlling the prevalence (i.e. expectations) of any features of interest in the model’s outputs. Despite their potential, the widespread adoption of these techniques has been hindered by the difficulty in adapting the complex, disconnected code. Here, we present disco, an open-source Python library that brings these techniques to the broader public

## introduction the advent of pre-trained generative models has had a paradigm-shifting impact in natural language processing @xcite @xcite , but also in other fields such as speech processing @xcite , code generation @xcite , computer vision @xcite @xcite , among others. the common thread in these models is that of training a probability distribution over a given space of interest (text, images, audio, etc.) using large corpora, which can then be used to generate samples in this space. in particular, in nlp, these models have found applications not only in traditional tasks such as summarization @xcite , but also opened new capabilities through few-shot learning @xcite . however, the models may suffer from deficiencies stemming both from replicating some patterns in the training data that are not desirable such as offensiveness @xcite or unequal treatment @xcite , but also from failing to replicate other more desirable patterns which are also present in the data but are hard to capture by the neural network model, such as truthful information @xcite . for these reasons, there is a growing interest in controlling the generations to align with human values @xcite . @xcite proposed a comprehensive framework to tackle these issues that they coined "generation under distributional control" or gdc. this framework builds on the idea introduced by parshakova et al. (2019b) that we can decouple the problems of describing the target distribution representing the aligned generative model (i.e., the what) from the problem of approximating it (i.e., the how). in particular, they design the target distribution by fixing the desired expectations of some features of interest while avoiding catastrophic forgetting, and approximate it using the dpg algorithm @xcite ). yet, other target distributions are possible. for example, @xcite showed that reinforcement learning from human feedback or rlhf @xcite @xcite could also be framed as approximating a well-defined target distribution, highlighting the generality and flexibility of the distributional approach. here, we present disco, a user-friendly library that provides developers, researchers, and practitioners easy access to state-of-the-art distributional control techniques. in what follows, we provide an overview of the gdc theoretical framework and its associated techniques before introducing the toolkit, with an overview of some design choices and a quick tour of its capabilities. we then suggest possible applications and apply disco to three experimental use cases. ## background let's assume a pre-trained generative model a(•) that defines a probability distribution over a sample space x such that we can efficiently compute the probability a(x) for any element @xmath0. under the gdc framework, controlling the generative model amounts to defining a new probability distribution p * (x) to sample from. this probability distribution is such that 1. it meets the control conditions: given a vector of n pre-defined real-valued functions (or features) ϕ(x) = [ϕ i (x)] @xmath1...n defined over @xmath2, p * is constrained such that each moment (i.e. expectation) µ i . = e x∼p * ϕ i (x) matches a desired value μi ; and 2. it avoids catastrophic forgetting: p * is the distribution that minimizes kl divergence from a among all distributions p ′ ∈ c satisfying the previous constraints p * . and where the moments e x∼p distr (•;λ) ϕ(x) are computed through self-normalized importance sampling (snis; owen, 2013) using a(•) or any other proposal distribution @xcite . ## design and implementation expressing preferences to express either pointwise or distributional preferences, distributions support the constrain() method, which given a list of features ϕ i (x) and their corresponding moments μi , returns a representation of the target distribution that respects the constraints while deviating minimally from the original model. 2 features can be defined using the scorer class, which accepts a function or a lambda abstraction taking a sample s and a context c as arguments and returning a score. an important class of features are boolean features, represented by the booleanscorer class. while general features can only be used to define distributional constraints, boolean features can also be used to define pointwise constraints. for example, we can score the presence of the string "amazing" in the sample s, as follows: conditional features can be expressed simply by taking the context c into account. next, we can define an ebm with a pointwise constraint requiring that all our samples must include (the string) "amazing" by setting the target moment of a booleanscorer feature to 1: distributional constraints are enforced by specifying any real-valued target moment or using nonbinary features. finally, an rlhf-like target distribution with regularization parameter beta and a reward scorer can be defined in the following way. in all cases, the resulting target is a posi-tivescorer representing the target distribution as an ebm. crucially, it is not an instance of distribution since it does not allow sampling. fine-tuning the model to tune a distribution to approximate the target ebm so that we can use it to generate samples, disco provides a set of tuner classes, notably the dpgtuner and fdpgtuner for the unconditional case, and cdpgtuner and fcdpgtuner for the conditional case. the difference between using vanilla dpg/cdpg and f -dpg/f -cdpg for tuning is that whereas the former are restricted to minimizing the kl divergence to the target, f -(c)dpg can be used to minimize any f -divergence (defaulting to jensen-shannon, which often works well in practice). note that we treat the unconditional case as a particular instance of the conditional one in which there is a single fixed context, the reason why (f)dpgtuner is also a (f)cdpgtuner. conditional tuning only requires further specifying a distribution of possible contexts on which the model will be conditioned. this is done with a contextdistribution, such as for instance the datasetcontextdistribution, which samples contexts from huggingface datasets @xcite . the tuner reports a number of metrics that are useful to monitor the training progress. a number of logger classes are provided to keep track of these metrics, including json, w&b, neptune or custom loggers. one of the most important reported metrics includes the estimate of the model's divergence to the target, [kl/js/tv]_target_model, as measured by kl, js or tv, respectively, one of which is typically the quantity being optimized. other metrics can include the features moments and the divergence from the base model if they are requested. improving the approximation with mc sampling after the tuning is done, model is now a better approximation to the target ebm, but it is not guaranteed to perfectly match this distribution. while further training can improve the situation, another alternative is using quasi-rejection sampling (qrs; @xcite , a monte-carlo sampling technique that allows to trade-off sampling efficiency for a higher fidelity to the target distribution -a higher value of beta yields a better approximation at a higher computational cost by retaining a smaller fraction of samples. compilability/style constraints on code generation language models trained on clean code data can still generate code that does not compile or, even if it does, can fail to meet style standards. @xcite showed that it was possible to effectively improve code generation models on both accounts by using pointwise constraints on the result coming from the python compiler and of an off-the-shelf linter. limiting hallucinations seq2seq models such as those used in summarization or nmt have a common failure mode by which they generate information not originally present in the source document (aka "hallucinations"). entity-level factual consistency (nan et al., 2021) is a family of measures that detect whether produced entities were included in the source, and whether they are part of the target in the dataset. @xcite showed that gdc could be successfully applied to improve on these metrics. below, we reproduce part of the experiments. debiasing language models gdc can address bias in language models by defining a feature detecting a population of interest, and setting the target moments of the feature to the desired value. @xcite experimented with reducing gender bias, while @xcite use this technique to balance the "regard" score among different religious groups. ## showcase experiments this section presents a selection of experiments to showcase a few use cases of disco, along with code snippets illustrating their implementation.