Controlled Text Generation via Language Model Arithmetic

TL;DR

This paper tackles the challenge of controlling text generation without retraining by introducing model arithmetic, a principled framework to compose and bias multiple language models and attributes via formulas over probability distributions. It advances CTG by enabling precise, interpretable control through operators like linear combinations, classifiers, and a union operator, and by unifying prior CTG techniques under a single mathematical framework. A key methodological contribution is the weighted KL-optimality formulation, which yields the composite distribution P by combining input distributions Q_i with context-dependent weights, and is complemented by an extension of speculative sampling that reduces runtime overhead for multi-model expressions. The authors demonstrate substantial toxicity reduction and fine-grained control across multiple models and tasks, while achieving significant speedups (up to 64% reduction in model calls in some settings) through generalized speculative sampling. The work offers practical impact for deploying customizable, attribute-aware LLMs efficiently and without dataset-specific retraining, and provides open-source resources to foster reproducibility and broader adoption.

Abstract

As Large Language Models (LLMs) are deployed more widely, customization with respect to vocabulary, style, and character becomes more important. In this work, we introduce model arithmetic, a novel inference framework for composing and biasing LLMs without the need for model (re)training or highly specific datasets. In addition, the framework allows for more precise control of generated text than direct prompting and prior controlled text generation (CTG) techniques. Using model arithmetic, we can express prior CTG techniques as simple formulas and naturally extend them to new and more effective formulations. Further, we show that speculative sampling, a technique for efficient LLM sampling, extends to our setting. This enables highly efficient text generation with multiple composed models with only marginal overhead over a single model. Our empirical evaluation demonstrates that model arithmetic allows fine-grained control of generated text while outperforming state-of-the-art on the task of toxicity reduction. We release an open source easy-to-use implementation of our framework at https://github.com/eth-sri/language-model-arithmetic.

Figures (5)

• Figure 1: Overview of model arithmetic using an illustrative example. We outline the procedure for generating a fairy tale (left) using the models $\bm{{M_\text{child}}}$, $\bm{{M_\text{adult}}}$, $\bm{{M_\text{magic}}}$ that produce text conditioned on the attributes child, adult, and magic, respectively and $\bm{{C_\text{formal}}}$ a classifier for the formality of text. The right table shows example outputs for different (partial) formulas. Image attribution in \ref{['appendix:attribution']}.
• Figure 1: Overview of Model Arithmetic where $\mathcal{I}_1(x) := [Q_1(x) > Q_2(x)]$ and $\mathcal{I}_2(x) := 1 - \mathcal{I}_1(x)$, $\bm{{C}}$ is a classifier and $U$ the uniform distribution.
• Figure 2: Attribute presence for several attributes and formulas. The dashed line indicates the value of the attribute when prompting the model to use the attribute.
• Figure 3: Model calls per token with speculative sampling for $\bm{{M}} + \lambda \bm{{M_a}}$, $\lambda \in [0.1, 1.0]$.
• Figure 4: Perplexity for several attributes and formulas. Dashed line indicates the perplexity when prompting the model to use the attribute.

Theorems & Definitions (8)

• Theorem 1: Weighted KL-Optimality
• Lemma 1
• proof
• Lemma 2
• proof
• proof
• Lemma 3
• proof