PaperlayerView exact revision on arXiv

Source-linked AI summary

Controlled Text Generation via Language Model Arithmetic

Jasper Dekoninck, Marc Fischer, Luca Beurer-Kellner, Martin Vechev

arXiv:2311.14479v3cs.CL

TL;DR

Precise control over vocabulary, style, and other text attributes is important, but existing prompting and controlled-generation methods have limitations in controllability and efficiency. The paper introduces formula-based model arithmetic for composing and biasing language models, with generalized speculative sampling for efficient generation. Its evaluation reports improved toxicity reduction over prior approaches.

  • Problem

    Prompting makes attribute strength ambiguous, while controlled-generation methods can make strength unintuitive and lack natural, controlled combinations of attributes.

    Source: §1 · p. 1

  • Method

    Model arithmetic composes and biases multiple language models through precise formulas, and extends speculative sampling to defer expensive model calls.

    Source: §1 · pp. 2–3

  • Results

    Model arithmetic subsumes prior approaches for toxicity reduction and significantly outperforms them.

    Source: §1 · p. 3, §7 · p. 9

  • Takeaways & Limitations

    The framework can incorporate prompting or fine-tuning and express prior controlled-generation techniques as formulas.

    Source: §1 · pp. 2–3

    Model arithmetic can also increase toxic content and potentially avoid built-in safety filters.

    Source: BROADER IMPACT · p. 10

Abstract

from arXiv · show

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.

1 INTRODUCTION

LLM customization matters across vocabulary, style, and emotional expression, but existing prompting and controlled-generation approaches have limits in precise, efficient, and composable control. The paper introduces formula-based model arithmetic and extends speculative sampling to make multi-model generation more efficient.

1 INTRODUCTION

  • Existing prompting is ambiguous about attribute strength, while controlled-generation strength parameters can become unintuitive and do not naturally support controlled combinations of attributes.These approaches may also require specific training data and evaluate multiple models at inference time.Source: §1 · pp. 1–2
  • Model arithmetic composes multiple models and attributes into a precisely controlled formula-based model, and can express prior controlled text generation approaches.The illustrative example combines child-like and magical language, subtracts an adult-conditioned model, and biases formality with a classifier.Source: §1 · pp. 2–3, Figure 1 · p. 2
  • Generalized speculative sampling postpones expensive model calls in arithmetic formulas, enabling multi-model generation with marginal overhead over a single model.The resulting speedup also extends to prior controlled-generation techniques expressible in model arithmetic.Source: §1 · pp. 2–3
  • Model arithmetic outperforms prior approaches in toxicity reduction, and its extended speculative sampling reduces model calls by up to 64%.The evaluation also finds it more expressive than prior controlled text generation work.Source: §1 · p. 3

2 BACKGROUND

This section introduces probability distributions and token-by-token language modeling, then describes controlled text generation and speculative sampling. These concepts provide background for the paper's methods.

2 BACKGROUND

  • A discrete probability distribution assigns a probability to each element in a finite set, typically tokens or subwords in language modeling.Token probabilities can be conditioned on the preceding sequence.Source: §2 · p. 3
  • Autoregressive language models generate text token by token by modeling each token's probability conditioned on preceding tokens.Source: §2 · p. 3
  • Controlled text generation biases a language model's distribution toward an attribute, often using a strength parameter and a classifier, fine-tuned model, or alternate prompt.Many approaches draw on a variant of Bayes' rule.Source: §2 · p. 3
  • Speculative sampling proposes several tokens with a small model, then validates them with a larger model; rejected proposals are resampled while accepted proposals preserve the larger model's distribution.A single validation call can validate the proposed sequence and may also provide the next token.Source: §2 · p. 3

3 MODEL ARITHMETIC

Model arithmetic constructs a sampleable output distribution by combining input distributions through weighted KL-divergence optimization. Its operators support attribute steering, classifier guidance, and nonlinear combinations of model distributions.

3 MODEL ARITHMETIC

  • Weighted KL optimization combines arbitrary input distributions into an output distribution, with a softmax solution under mild technical assumptions.The theorem also applies when conditioning includes future tokens, not only autoregressive context.Source: §3 · p. 4
  • Model arithmetic represents distributions through weighted log-probability vectors, enabling arithmetic combinations of prompts, attributes, language models, and classifiers.The notation omits normalization and treats vector multiplication element-wise.Source: §3 · p. 4
  • Linear combinations steer output toward attributes with positive coefficients and away with negative coefficients, and can express prior CTG approaches.For formal-style generation, a composite using a negative coefficient for a prompted formal model produces more formal output than standard prompting; Table 2 gives the example.Source: §3 · p. 5, Table 2 · p. 5
  • Binary classifiers can guide generation toward attributes such as human-like text, but classifier guidance is computationally approximated by evaluating only the top-k tokens.Prior classifier-guidance approaches are generally restricted to top-k sampling because evaluating every token is infeasible.Source: §3 · p. 5
  • The union operator combines token probabilities nonlinearly using their maximum, allowing uncommon or disparate attributes to coexist when linear combinations cannot raise low-likelihood tokens enough.In the example, prompting collapses to an alien-like output, linear combination remains mostly human-like, and union interpolates both attributes; Table 4 shows the example.Source: §3 · pp. 5–6, Table 4 · p. 6
  • Union-based subtraction limits biasing to tokens where the attribute distribution exceeds the original, keeping other tokens unchanged up to normalization; the paper also derives an intersection operator.The authors state this localized biasing enables better toxicity reduction than linear combination.Source: §3 · p. 6

4 SPECULATIVE SAMPLING

The section extends speculative sampling to arithmetic formulas by treating formula terms as staged proposal models, postponing expensive evaluations until validation. Speculative factors control token batches, while classifier terms remain limited to one-token steps.

4 SPECULATIVE SAMPLING

  • Partial formulas induce proposal distributions that can be iteratively validated until samples match the distribution of the full formula F.The extension relies on a lemma that successive speculative-sampling steps between distributions preserve the final target distribution.Source: §4 · p. 6
  • Each term’s speculative factor s_i sets how many tokens are sampled before that term is evaluated and used for validation.Terms can be reordered during inference to evaluate only those needed for current validation; classifier terms require s_i = 1.Source: §4 · p. 6
  • The supersede operator uses a small model’s output until the larger model has been computed, expressing standard speculative sampling in model arithmetic.The two-model form is supersede(m, M).Source: §4 · p. 7

5 EVALUATION

The evaluation examines toxicity reduction, attribute control, and speculative-sampling efficiency. It reports lower toxicity without sacrificing fluency at substantial bias strengths, smooth attribute control with low perplexity cost, and fewer model calls with faster inference.

5.1 TOXICITY REDUCTION

  • At bias strength 0.6, PREADD has higher perplexity than the union operator at 0.96, while union plus a classifier attains the lowest toxicity across models with similar or lower perplexity.Perplexity degrades only at the highest tested biasing strength.Source: §5.1 · p. 7

5.2 FINE-GRAINED CONTROL

  • Attribute presence generally rises smoothly with relative λ_i, and the curves are mostly linear; sports rises significantly only after its relative strength exceeds 1.0.The sports formulation counterweights M_helpful, and its attribute presence can exceed the standard-prompting value.Source: §5.2 · p. 8, Figure 2 · p. 8
  • Maximum perplexity across formulas and attributes is 6.2, versus 4.8 for the prompted models; educational fluency is slightly affected only at the highest evaluated strengths.The passage attributes this exception to heavy use of a fluency-agnostic classifier.Source: §5.2 · p. 8

5.3 SPECULATIVE SAMPLING

  • With up to four constituent models, speculative sampling reduces calls to at most 1.44 per token and boosts inference speed by up to 2.24 times.The evaluation optimizes speculative factors using a calibration run of 10 samples.Source: §5.3 · p. 9, Table 7 · p. 9
  • For F = M + λMa, Fig. 3 shows model calls per token increase with λ as KL divergence from original model M rises, decreasing acceptance probability in speculative sampling.Source: §5.3 · p. 9

6 RELATED WORK

The related-work discussion contrasts model arithmetic with controlled-generation approaches requiring decoding-time gradient steps or training data, and describes complementary speculative-sampling extensions.

6 RELATED WORK

  • Some CTG approaches impose hard constraints, use classifier gradients, or reduce mutual information, but require costly decoding-time gradient steps or training data.The discussion cites Meng et al. (2022), Kumar et al. (2021; 2022), Dathathri et al. (2020), and Yang et al. (2023).Source: §6 · p. 9
  • Other speculative-sampling work stages multiple small models or uses several at once with tree-based proposals, approaches compatible with the paper’s supersede extension.The cited approaches are Spector and Re (2023) and Miao et al. (2023).Source: §6 · p. 9
  • Tree-based proposal-sequence generation is presented as an orthogonal improvement that can be incorporated into the supersede operator.The related-work passage describes both staged and concurrent multiple-small-model approaches.Source: §6 · p. 9

7 CONCLUSION

The conclusion presents model arithmetic as a formula-based framework for composing LLMs and controlled-generation attributes. It says the framework subsumes prior toxicity-reduction approaches and that speculative sampling reduces multi-model computational overhead.

7 CONCLUSION

  • Model arithmetic composes multiple LLMs and controlled-generation attributes through a principled, formula-based framework.Source: §7 · p. 9
  • A novel union operator lets model arithmetic subsume prior CTG approaches for toxicity reduction and significantly outperform them.Source: §7 · p. 9
  • A speculative-sampling procedure for arithmetic formulas heavily reduces the computational overhead associated with multi-model CTG.Source: §7 · p. 9

BROADER IMPACT

Model arithmetic offers precise control and expressiveness, but that flexibility can also be used to generate undesirable attributes. The appendices show how prior work fits the framework and define an intersection operator.

BROADER IMPACT

  • Model arithmetic can generate undesirable attributes, including increased toxic content that could potentially avoid built-in safety filters.The authors judge its benefits to outweigh these potential risks because it supports more precise control and expressiveness for inclusive, controlled content.Source: BROADER IMPACT · p. 10

A PRIOR WORK

  • Table 8 shows how multiple prior works can be expressed within the model-arithmetic framework.The table notes that some classifier-based methods require TopK sampling for efficiency, while other sampling choices are unrestricted.Source: A PRIOR WORK · p. 16, Table 8 · p. 16

B INTERSECTION

  • The intersection operator assigns high probability to a token only when both input distributions assign it high probability.Its solution is σ(min(log Q1, log Q2)).Source: B INTERSECTION · p. 16

C ATTRIBUTION

  • The figure icons are credited to Freepik via flaticon.com.Mchild, Madult, Mmagic, and Cformal are attributed to Freepik.Source: C ATTRIBUTION · p. 16

D PROOFS

The appendices provide proofs for the model-arithmetic optimization and classifier formula, establish speculative sampling across distributions, and describe the algorithm and factor-selection procedure.

D.1 SOLUTION MINIMIZATION PROBLEM

  • The optimization proof assumes normalization and a nonnegative weighted sum; the distributions must also share a token set and be nonzero on it.The first assumption ensures proper normalization, while the shared support condition makes KL divergence well-defined.Source: §D.1 · p. 17

D.2 CLASSIFIER FORMULA

  • The classifier-formula proof reduces the optimization to an equivalent objective, with additive constants that do not affect its solution.The derivation uses the KL-divergence definition and identifies the expected cross-entropy term.Source: §D.2 · p. 18

D.3 SPECULATIVE SAMPLING ON N DISTRIBUTIONS

  • The proof for speculative sampling on n distributions proceeds by induction, extending the two-distribution case to n+1 distributions.It applies the procedure successively to adjacent distributions to obtain a sample from the final distribution.Source: §D.3 · p. 18

E SPECULATIVE ALGORITHM FOR MODEL ARITHMETIC

  • Algorithm 2 tracks model evaluation positions, stored term evaluations, and normalized distributions while applying speculative sampling to a partial formula.It accepts a proposed token or rejects it and restarts after updating the relevant stored states.Source: E SPECULATIVE ALGORITHM FOR MODEL… · p. 19

E.1 DETERMINING SPECULATIVE FACTORS

  • Factor selection fixes s1 = 1, estimates proposal acceptance from 10 samples, and optimizes s2 using expected compute per accepted token.For multiple terms, the procedure assumes each term’s expected acceptance probability is constant regardless of how many models have already run; it approximates model-call compute as equal, despite key-value storage.Source: §E.1 · p. 20

F SENTIMENT CONTROL

In sentiment-control experiments, the union operator outperforms baselines, and combining it with a classifier yields stronger GPT-4 preference and sentiment results. The combined formula also has lower perplexity than PREADD in all cases.

F SENTIMENT CONTROL

  • The evaluation uses 1,000 positive and 1,000 negative IMDB reviews, truncated to 32 tokens, and asks models to continue with the opposite sentiment.GPT-4 selects the best response for sentiment and relevance.Source: F SENTIMENT CONTROL · p. 21
  • The union operator significantly outperforms all baselines: GPT-4 prefers it in every evaluated setting, with an average preference advantage of 5% over the closest-following baseline, PREADD.Source: F SENTIMENT CONTROL · p. 21
  • 12% average GPT-4 preference over PREADD: the union-plus-classifier formula wins in all cases and has lower perplexity than PREADD throughout.The combined formula also produces much higher measured sentiment than the other methods.Source: F SENTIMENT CONTROL · p. 21, Table 12 · p. 22

G EXPERIMENTAL DETAILS

The appendix records the experimental settings for toxicity, sentiment, GPT-4 preference judgments, and attribute evaluation. It also reports a capacity-dependent toxicity result for GPT-2.

G EXPERIMENTAL DETAILS

G.1 TRAINING AND PROMPTING DETAILS TOXICITY REDUCTION

  • Toxicity evaluation uses 2,000 filtered /pol/ messages with toxicity scores above 0.5, and stops completions at 32 tokens or specified stop markers.The tested models are Llama-2-13b, Pythia-12b, and MPT-7b; FUDGE uses a classifier fine-tuned on Jigsaw data.Source: §G.1 · p. 23

G.2 TRAINING AND PROMPTING DETAILS SENTIMENT CONTROL

  • Sentiment evaluation uses 1,000 positive and 1,000 negative IMDB reviews, truncated to 32 input tokens, with completions stopped after 64 tokens or stop markers.The conditioned models receive a positive- or negative-sentiment sentence in the prompt.Source: §G.2 · pp. 23–24

G.3 GPT-4 PROMPTING DETAILS

  • GPT-4 compares two methods’ continuations, with their order randomized to avoid selection bias; reported preferences are percentages of wins, losses, or draws.It is prompted to choose the better completion or neither when it cannot decide.Source: §G.3 · p. 24

G.4 PROMPTING DETAILS ATTRIBUTES

  • Attribute evaluations use classifiers for formality, happiness, and topics, plus normalized average word length for simplicity.The attribute experiments use the Llama-2-Chat prompt template and distinct system prompts.Source: §G.4 · pp. 24–25, Table 13

H.1 TOXICITY REDUCTION FOR GPT-2

  • For the smallest GPT-2 model, FUDGE beats the union operator and PREADD; the union operator beats previous methods on the other evaluated models.Combining classifiers with the union operator substantially outperforms previous methods in toxicity reduction.Source: §H.1 · p. 26

H.2 PERPLEXITY RESULTS FOR FINE-GRAINED CONTROL

Perplexity remains relatively low across attributes and changes smoothly with model parameters, although very high educational strength can reduce fluency. The supplied examples also show outputs from different formulas and extreme strengths.

H.2 PERPLEXITY RESULTS FOR FINE-GRAINED CONTROL

  • Perplexity changes smoothly with model parameters and remains relatively low across attributes; manual inspection finds coherent, fluent text except at very high educational strengths.At high classifier strengths (λ = 12), educational outputs may replace words with nonexistent ones.Source: §H.2 · p. 26
  • Table 15 showcases outputs for formulas F1 and F2 at the extreme strengths tested, with formulas combining attribute models and classifiers.Source: Table 15 · p. 27

I OUTPUT EXAMPLES

  • Examples compare different formulas on the same input prompt, including outputs from M, Mformal, M + Mformal, and 2Mformal − M.The examples use Llama-2-Chat with 13B parameters and several prompt templates, generating one response per method.Source: Table 16 · p. 28, §H.2 · p. 28, I OUTPUT EXAMPLES · p. 28

Related papers

Loading 2311.14479v3…