Source-linked AI summary

Controlling Conditional Language Models without Catastrophic Forgetting

Tomasz Korbak, Hady Elsahar, German Kruszewski, Marc Dymetman

arXiv:2112.00791v2cs.LGcs.CL

TL;DR

Pretrained generative models can miss downstream requirements, while control approaches face limited ground-truth data and catastrophic forgetting. This paper extends distributional policy gradients to conditional energy-based models and evaluates the approach across translation, summarization, and code generation. The authors report that CDPG improves desired characteristics while retaining fluency and diversity of the original models.

  • Problem

    Control objectives often lack ground-truth training data, and methods that enforce them can cause catastrophic forgetting; existing EBM-based DPG is limited to unconditional generation.

    Source: §1 · p. 2
  • Method

    Conditional DPG extends DPG to approximate conditional energy-based models using context-specific target distributions that balance control objectives with closeness to the original model.

    Source: §1 · p. 2
  • Results

    CDPG fine-tuning yields desired characteristics such as improved factual consistency or compilability while retaining the original model’s fluency and diversity.

    Source: §5 · p. 9
  • Takeaways & Limitations

    CDPG shifts model distributions in a minimally invasive way and does not require ground-truth training data.

    Source: §5 · p. 9

    Evaluation on dialogue and controls for Python-function semantics remains future work.

    Source: §5 · p. 9

Abstract

from arXiv · show

Machine learning is shifting towards general-purpose pretrained generative models, trained in a self-supervised manner on large amounts of data, which can then be applied to solve a large number of tasks. However, due to their generic training methodology, these models often fail to meet some of the downstream requirements (e.g., hallucinations in abstractive summarization or style violations in code generation). This raises the important question of how to adapt pre-trained generative models to meet all requirements without destroying their general capabilities ("catastrophic forgetting"). Recent work has proposed to solve this problem by representing task-specific requirements through energy-based models (EBMs) and approximating these EBMs using distributional policy gradients (DPG). Despite its effectiveness, this approach is however limited to unconditional distributions. In this paper, we extend DPG to conditional tasks by proposing Conditional DPG (CDPG). We evaluate CDPG on four different control objectives across three tasks (translation, summarization and code generation) and two pretrained models (T5 and GPT-Neo). Our results show that fine-tuning using CDPG robustly moves these pretrained models closer towards meeting control objectives and -- in contrast with baseline approaches -- does not result in catastrophic forgetting.

1. Introduction

Conditional DPG extends distributional policy gradients to conditional energy-based models, addressing both the scarcity of supervised control data and catastrophic forgetting. Across translation, summarization, and code generation, the paper reports improved control while preserving important qualities of the original models.

1. Introduction

  • Control objectives often lack ground-truth training data, while many approaches risk catastrophic forgetting beyond the objective’s scope.The available signal may be only an indicator of whether a sample satisfies the objective.Source: §1 · p. 2
  • CDPG extends DPG to approximate conditional EBMs, defining context-specific target distributions that balance control objectives against closeness to the original model.The paper evaluates translation, summarization, and code generation using T5 and GPT-Neo.Source: §1 · p. 2
  • CDPG makes translation samples satisfying the numeral-as-digits constraint 116 times more likely, while the toy objective leaves other translation aspects unchanged.The constraint asks that numeral nouns such as “two” be translated as digits.Source: §1 · p. 2
  • For summarization, CDPG increases correct and relevant named entities and Rouge score without degrading diversity or quality, unlike RL approaches.The control objective operationalizes factual faithfulness as entity-level factual consistency.Source: §1 · p. 2
  • 40% of GPT-Neo Python functions compile, with four PEP8 violations on average; CDPG improves compilability and compliance without reducing diversity or causing degeneration.The code-generation control objectives target compilable Python functions and PEP8 compliance.Source: §1 · p. 2
  • The paper evaluates controls for number format, factuality, Python compilability, and PEP8 compliance across translation, summarization, and code generation.The evaluation uses T5 and GPT-Neo.Source: §1 · p. 2

2. Method

The method represents control with per-context energy-based target distributions that stay close to a pretrained conditional model. CDPG trains one seq2seq model against these targets on average across contexts, estimating context-dependent partition functions by sampling.

2. Method

  • For each context, control is framed as finding a distribution that satisfies a binary scorer while minimizing KL divergence from the pretrained model, represented by a conditional EBM.The target distribution is the normalized version of the per-context EBM.Source: §2 · p. 3
  • Sampling and MAP decoding from the target are intractable because it represents potentially infinitely many per-context models, and the whole-sequence scorer prevents autoregressive factorization.The scorer is defined over the complete output sequence rather than its parts.Source: §2 · p. 3
  • CDPG trains a single seq2seq model to approximate the per-context targets on average by minimizing expected cross-entropy over contexts.Earlier DPG approximated unconditional EBMs and is not directly usable for a conditional model covering many contexts.Source: §2 · p. 3
  • The gradient estimate samples outputs from the seq2seq model and uses the EBM score, weighted by the context-specific partition function, as a pseudoreward term.The objective remains expected cross-entropy over contexts, not expected reward.Source: §2 · pp. 3–4
  • Because the partition function varies by context, CDPG estimates each Zc using importance sampling over model-generated samples for that context.The described procedure draws M samples from πθ(x|c) and reweights them by their likelihood.Source: §2 · p. 4
  • Training samples N contexts per epoch, draws M outputs per context, stores context-output pairs with their partition estimates, then shuffles the buffer for gradient updates.The learning rate, N, and M are hyperparameters.Source: §2 · p. 4

3. Experiments

Across translation, summarization, and code generation, CDPG is compared with DPG and reinforcement-learning baselines on held-out contexts, measuring both target-distribution fit and distance from the pretrained model.

3. Experiments

  • Experiments used held-out contexts and measured expected divergence from both the target distributions and the original pretrained model.The evaluation covered four control objectives across translation, summarization, and code generation.Source: §3 · p. 4

3.1. Baselines

  • The DPG-like ablation uses a running-mean constant instead of context-specific Z_c, producing biased gradient estimates; Reinforce maximizes constraint reward, while Ziegler adds an adaptive KL penalty.The baselines include Reinforce and Ziegler variants of reward-based fine-tuning.Source: §3.1 · p. 4

3.2. Translation

  • Constraint satisfaction rose from 0.006 to 0.7 with CDPG, which nearly halved expected DKL(pc, πθ) and outperformed DPG.CDPG also stayed closer to the original distribution than Reinforce and incurred almost no BLEU-4 drop.Source: Figure 2 · p. 5, §3.2 · p. 5

3.3. Summarization

  • CDPG alone stably decreased expected DKL(pc, πθ), while improving named-entity quantity, source precision, target recall, Distinct-2, and ROUGE-L.The authors suggest the ROUGE-L increase may reflect reduced bias toward including too few factually correct entities.Source: Figure 3 · p. 6, §3.3 · p. 6

3.4. Code generation

  • In code generation, CDPG raised compilability from around 40% to around 65% and reduced average PEP8 violations.The PEP8 objective also increased compilability, as many style violations are also compilation errors.Source: Figure 4 · p. 7, §3.4 · p. 7
  • CDPG and its DPG-like ablation moderately diverged from the original model while preserving code length and AST-node statistics; Reinforce shortened functions and Ziegler produced degenerate, repetitive samples.The CDPG–DPG gap was smaller for code generation than summarization, consistent with lower partition-function variation making the constant approximation less biased.Source: §3.4 · p. 7

3.5. Qualitative analysis

  • Across code errors and summarization entities, CDPG reduced error frequencies and increased frequencies across entity types, including rare ones.CDPG and DPG also matched original token frequencies more closely, while Ziegler and Reinforce had shorter tails of rare tokens.Source: §3.5 · p. 8

4. Related work

The related work contrasts CDPG with approaches tailored to particular translation constraints and reviews methods for factual summarization and compilable code. Earlier translation methods either train jointly with a log-linear model or do not generalize to arbitrary constraints.

Injecting prior information in machine translation

  • Earlier translation methods either train jointly with a log-linear model or target specific terminology constraints rather than arbitrary constraints.Zhang et al. (2017) regularize neural translation toward a log-linear model; constrained decoding and source-context methods address terminology constraints [Chatterjee et al., 2017; Hasler et al., 2018; Dinu et al., 2019; ibn Alam et al., 2021].Source: Injecting prior information in machine… · pp. 8–9
  • Figure 6 plots token frequency against rank for samples from CDPG and baselines; longer tails indicate greater sample diversity.The plotted samples number 10,280.Source: Figure 6 · p. 9

Reducing hallucinations in summarization

  • Summarization work includes reranking with entailment classifiers and reinforcement-learning fine-tuning, while entity-level factual consistency measures whether summary entities appear in the source.Maynez et al. (2020) found hallucinations in more than 70% of single-sentence summaries; Nan et al. (2021) introduced entity-level factual consistency.Source: Injecting prior information in machine… · p. 9

Controllable code generation

  • Prior work used semantic constraints in neural probabilistic context-free grammars for unconditional generation, penalized syntactically invalid SQL in policy-gradient natural-language-to-SQL training, and applied DPG to improve compilability of unconditional code language models.Source: Injecting prior information in machine… · p. 9

5. Conclusion

The paper presents CDPG as a way to fine-tune conditional language models for arbitrary constraints without ground-truth training data. It concludes that controlled models retain fluency and diversity, and identifies broader task and objective coverage as future work.

5. Conclusion

  • CDPG fine-tunes conditional language models to satisfy arbitrary constraints without ground-truth training data, while shifting distributions minimally.The resulting models combine desired properties such as improved factual consistency or compilability with the original model’s fluency and diversity.Source: §5 · p. 9
  • Future work could evaluate CDPG on dialogue, test semantic rather than syntactic code constraints, and extend it to conditional exponential-form EBMs for distributional constraints.The latter could represent desired expected values for features of generated samples.Source: §5 · p. 9

A. Details of metric and score calculation

The appendix describes how the paper estimates KL divergences and operationalizes its translation, summarization, and code-generation metrics. Its measures include importance sampling, entity checks, syntax compilation, lint violations, and AST size.

A.1. KL divergences

  • KL calculations use batches of contexts and model samples; the conditional partition function is importance-sampled, while the reverse KL can be computed directly from model samples.A small epsilon stabilizes partition estimates, and KL sampling uses pure ancestral sampling rather than top-p or beam search.Source: §A.1 · p. 14

A.2. Translation

  • Translation constraints use a lookup mapping numeral nouns from one to nine to digits, and translation BLEU-4 is computed with SacreBLEU.A constraint is satisfied when each numeral noun in the source has a corresponding digit in the translation.Source: §A.2 · p. 14

A.3. Summarization

  • Summarization factual consistency requires at least four named entities, all present in the source; entity extraction uses SpaCy categories and excludes dates, times, and numerals.The named-entity categories include PERSON, FAC, GPE, ORG, NORP, LOC, and EVENT.Source: §A.3 · p. 15

A.4. Code generation

  • Code compilability is syntactic validity under Python’s compile check, PEP8 violations are counted with pycodestyle, and AST node counts include only compilable samples.Compilation checks the concatenated signature and generated function, with text after the function declaration removed.Source: §A.4 · p. 15

A.5. Normalized standard deviations for Zc across tasks

The appendix normalizes the standard deviation of conditional partition functions by their average. It links lower normalized variability with a smaller performance gap between CDPG and DPG on Figures 3–4.

A.5. Normalized standard deviations for Zc across tasks

  • Lower normalized standard deviation of Z_c explains a closer performance gap between CDPG and DPG on Figures 3–4.Normalized standard deviation is std(Z_c)/avg(Z_c).Source: §A.5 · p. 15

B. Hyperparmeters and implementation details

This section reports implementation details, hyperparameter references, and normalized standard deviations for selected control objectives.

B. Hyperparmeters and implementation details

  • Each training run took approximately 5 days on 2 Nvidia V100 GPUs.The models were implemented using PyTorch and HuggingFace Transformers.Source: B. Hyperparmeters and implementation… · p. 15
  • Figure 7 reports normalized standard deviations of Zc for summarization factual consistency, code compilability, and PEP8 objectives.The figure concerns Zc values associated with unconditional EBMs in the codomain of conditional EBMs.Source: Figure 7 · p. 16
  • Table 1 lists hyperparameters for code-generation experiments.Table 2 lists hyperparameters for translation and summarization experiments.Source: Table 1 · p. 16, Table 2 · p. 16

C.1. Translation with terminology constraint

The translation appendix presents generated translations for randomly selected evaluation sentences and specifies the decoding method.

C.1. Translation with terminology constraint

  • Tables 3–7 show translations for 5 randomly chosen source sentences from the evaluation set.The models were fine-tuned using Conditional DPG and three baselines.Source: §C.1 · p. 17
  • Translations in Tables 3–7 were generated with beam search.Source: §C.1 · p. 17

C.2. Summarization with entity-level factual consistency constraint

The appendix provides generated examples for summarization and code generation under the stated control constraints, with separate sampling procedures for each task.

C.2. Summarization with entity-level factual consistency constraint

  • Tables 8–12 present beam-search summaries of five randomly chosen documents from evaluation set Ctest, generated by models fine-tuned using Conditional DPG and three baselines.Source: §C.2 · p. 17

C.3. Code generation with compilability constraint

  • Tables 13–15 show functions generated for 3 randomly chosen evaluation signatures using pure ancestral sampling.The models were fine-tuned for compilability using Conditional DPG and three baselines.Source: §C.3 · p. 17

C.4. Code generation with PEP8 constraint

The appendix presents code-generation examples from models fine-tuned to satisfy a PEP8 constraint.

C.4. Code generation with PEP8 constraint

Loading 2112.00791v2…