disco: a toolkit for controlling language models and other generative models

Language models were practically unknown to the general public until late 2022 yet we’ve all been using them on a daily basis for many years, like in the predictive keyboard application on our mobile phones.  Today, this has changed with generative AI language models going mainstream thanks to well-known conversational agents such as ChatGPT.

Language models belong to the wider class of “generative models”, which can be used to generate new content that resembles the data that the system was trained on. Language models are trained to predict human-written texts obtained from sources such as books, Internet forums, Wikipedia, etc. These models can capture a lot of useful information from their training data yet there can be obstacles, which are more or less serious, in directly using them for generation after training. First, the model is trained to produce likely continuations for any given input text or “prompt”. Therefore, when given a question, like “How to cook a carbonara sauce?,” it might continue with more questions or add extra context to the question without actually answering it. This is the reason why commercial models are often further trained to prefer continuations that are relevant for the user, using reinforcement learning from human feedback (RLHF). Second, because the model replicates the characteristics of Internet text, it will also capture all sorts of undesirable behaviour that lies within this text, including social bias and toxic language.

disco distributional control

At Naver Labs Europe, we’ve developed an elegant framework that we call “Generation under Distributional Control” or GDC that can tackle all of these issues [1,2,3,4,5]. To make this framework accessible to all, we’ve also created a library, disco, which leverages it to provide a simple and intuitive interface for expressing preferences and tuning the models to match the desired preferences.

A developer can use disco to define any feature they want to control (for example, the topic of a sentence, its toxicity, the gender of mentioned people, or its usefulness) and decide how prevalent that feature should be in the produced generations (for example, all samples should be of a given topic, none should contain toxic language, half of them should mention a person of a certain gender, and they should have maximal usefulness) to then train the model to meet those preferences while minimally deviating from the original model. disco natively supports a wide variety of language models through the HuggingFace interface, including GPT-2, OPT, BART, FLAN-T5, Pythia, LLaMA, among many others. Other generative models can be supported simply by writing appropriate wrappers for them.

disco’s main design principle, inherited from the GDC framework, is to separate the problem of representing what the user wants from how to get there. For the first part, the constrain method allows to specify the features that the user wants to control for and the expected value of these features. The result of this operation is an energy-based model or EBM. This object represents precisely the target distribution that the language model should follow. However, this model cannot be directly used for generation because it is not “autoregressive”, a crucial property that enables efficient generation in language models.

The second step is to tune a regular language model in autoregressive form to approximate the target EBM. A variety of algorithms can be used for this purpose, including distributional policy gradient or DPG [1], which minimizes the forward KL divergence to the target EBM (where “forward” means that the target distribution is in the first argument of the KL), or its generalization, \(\mathcal{f}\)-DPG [6], which minimizes any \(\mathcal{f}\)-divergence (including the forward KL, the reverse KL, jensen-shannon and total variation distance). Notably, RLHF is a particular case of this framework in which the objective is to minimize the reverse KL to an EBM of a certain form (more details in [4,6]). For this reason, disco is more general than other toolkits that center on RL methods only as it allows not only to express preferences on maximizing some scalar reward but also expressing distributional preferences of any sort.

Once we’ve obtained a fine-tuned model, we can sample generated content directly from it. Notably, we can also improve the quality of our samples without any additional tuning by using quasi-rejection sampling or QRS [5]. This algorithm keeps some of the samples and rejects others depending on a parameter that can be tuned to find the best trade-off between efficiency and quality.

To conclude, disco is a powerful and versatile toolkit for controlling language models and other generative models. Whether you’re a researcher looking to fine-tune a language model or a developer looking to build custom generative models, disco provides a simple and intuitive interface for working with them.

Get started today by running pip install disco-generation to install disco and heading to https://github.com/naver/disco for a quick overview.

References

[1] Parshakova, Tetiana, Jean-Marc Andreoli, and Marc Dymetman. “Distributional Reinforcement Learning for Energy-Based Sequential Models.” arXiv preprint*.*

[2] Khalifa, Muhammad, Hady Elsahar, and Marc Dymetman. “A distributional approach to controlled text generation.” In Proceedings of ICLR 2021.

[3] Korbak, Tomasz, Hady Elsahar, Germán Kruszewski, and Marc Dymetman. “Controlling Conditional Language Models with Distributional Policy Gradients.” In Proceedings of ICML 2022.

[4] Korbak, Tomasz, Hady Elsahar, Germán Kruszewski, and Marc Dymetman. “On Reinforcement Learning and Distribution Matching for Fine-Tuning Language Models with no Catastrophic Forgetting.” In Proceedings of NeurIPS 2022.

[5] Eikema, Bryan, Germán Kruszewski, Christopher R. Dance, Hady Elsahar, and Marc Dymetman. “An approximate sampler for energy-based models with divergence diagnostics.” In Transactions of Machine Learning Research (November 2022).

[6] Go, Dongyoung, Tomasz Korbak, Germán Kruszewski, Jos Rozen, Nahyeon Ryu, and Marc Dymetman. “Aligning language models with preferences through f-divergence minimization.”  To appear in ICML 2023.