Source-linked AI summary
A Distributional Approach to Controlled Text Generation
Muhammad Khalifa, Hady Elsahar, Marc Dymetman
TL;DR
Existing pointwise controlled-generation methods cannot enforce collective distributional conditions. GDC specifies pointwise and distributional constraints while minimizing divergence from a pretrained LM, then trains an autoregressive policy to approximate the resulting distribution; the reported pointwise experiments show better constraint satisfaction without the degeneration seen in baselines.
Problem
Most existing controlled-generation methods focus on individual outputs and cannot enforce desired collective statistical properties across generations.
Source: §1 · p. 2Method
GDC constrains feature expectations under a target distribution while minimizing its KL divergence from the pretrained LM, then uses KL-adaptive Distributional Policy Gradient to train an autoregressive policy approximating that distribution.
Source: §1 · p. 2, §2.3 · p. 5Results
GDC is reported to better satisfy pointwise constraints while avoiding degeneration; baselines can approach 100% constraint satisfaction only at severe cost to quality and sample diversity.
Source: §4 · p. 9Takeaways & Limitations
The decomposition is presented as a methodological advantage, though efficiently sampling from the optimal EBM remains a computational challenge.
Source: §4 · p. 9End policies still generate some samples that do not satisfy the constraints; using MCMC to reduce this residual gap is left for future work.
Source: §4 · p. 9
Abstract
from arXiv · showhide
We propose a Distributional Approach for addressing Controlled Text Generation from pre-trained Language Models (LMs). This approach permits to specify, in a single formal framework, both "pointwise" and "distributional" constraints over the target LM -- to our knowledge, the first model with such generality -- while minimizing KL divergence from the initial LM distribution. The optimal target distribution is then uniquely determined as an explicit EBM (Energy-Based Model) representation. From that optimal representation we then train a target controlled Autoregressive LM through an adaptive distributional variant of Policy Gradient. We conduct a first set of experiments over pointwise constraints showing the advantages of our approach over a set of baselines, in terms of obtaining a controlled LM balancing constraint satisfaction with divergence from the initial LM. We then perform experiments over distributional constraints, a unique feature of our approach, demonstrating its potential as a remedy to the problem of Bias in Language Models. Through an ablation study, we show the effectiveness of our adaptive technique for obtaining faster convergence. (Code available at https://github.com/naver/gdc)
1 INTRODUCTION
Prior controlled-generation methods focus on individual outputs and cannot enforce collective distributional conditions, a gap relevant to addressing bias in pretrained language models. GDC formulates control as constraints on a target distribution with minimal KL divergence from the pretrained model, then approximates its optimal EBM with an autoregressive policy.
1 INTRODUCTION
- Prior controlled-generation methods focus on individual outputs and cannot enforce distributional conditions over the set of generations.Such collective control is described as key to addressing social biases in language models trained on uncurated web corpora.Source: §1 · p. 2
- GDC formalizes pointwise and distributional control as constraints on a target distribution while minimizing KL divergence from the pretrained LM.The framework is presented as a generalization of Maximum Entropy and yields a unique target distribution represented as an EBM.Source: §1 · p. 2
- The optimal EBM is not directly sampleable, so KL-adaptive DPG trains an autoregressive policy to approximate it using earlier approximations to speed later rounds.The final policy can be assessed for its divergence from both the optimal target distribution and the initial LM.Source: §1 · pp. 2–3
- The approach separates determining a pivot optimal EBM from the harder task of exploiting it at inference time, helping diagnose the distinct challenges.The paper contrasts this decoupling with NLP EBM work that incorporates EBM representations inside neural-network training loops.Source: §1 · p. 2
- Experiments cover pointwise and distributional conditions, assessing divergence from GPT-2, fluency, and diversity, with better performance than strong baselines.The distributional experiments are presented as showing potential for addressing bias in pretrained language models.Source: §1 · p. 3
2 FORMALIZATION
GDC defines a target distribution through feature-moment constraints and selects the constraint-satisfying distribution closest to the pretrained model in KL divergence. It represents that target as an EBM and trains an autoregressive policy to approximate it.
2 FORMALIZATION
- GDC chooses a distribution satisfying specified feature-moment targets while minimizing KL divergence from the initial model.The moments are expectations of predefined real-valued feature functions under the desired distribution.Source: §2 · p. 3
- Setting a binary sports feature's target to 1.0 makes every output discuss sports, while a female-character target of 0.5 requires half the outputs to mention one.These illustrate pointwise and distributional constraints, respectively.Source: §2 · p. 3
2.1 CONSTRAINTS, INFORMATION GEOMETRY, EXPONENTIAL FAMILIES
- When the constraint set is nonempty, the optimal target distribution is unique and has an exponential-family EBM representation on the constraint support.The result also permits arbitrary-precision approximation; for hybrid constraints, some parameters may diverge, and thresholding them suffices for a good approximation.Source: §2.1 · p. 4
- The solution satisfies a Pythagorean identity for KL divergence and generalizes maximum entropy, which arises when the initial distribution is uniform.The framework's generality supports pointwise, distributional, and hybrid constraints.Source: §2 · p. 3, §2.1 · p. 4
2.2 FROM MOMENT CONSTRAINTS TO EBM
- To estimate EBM parameters, the method samples from the initial model, applies importance weights, estimates feature moments with self-normalized importance sampling, and uses SGD to match target moments.The resulting representation is an unnormalized EBM of the form P(x) = a(x) exp ⟨λ, φ(x)⟩.Source: §2.2 · pp. 4–5
2.3 FROM EBM TO AUTOREGRESSIVE POLICY
- Because the EBM specifies the target but does not provide samples when its normalizer is unknown, the method trains an autoregressive policy to approximate that target.Distributional Policy Gradient minimizes cross-entropy, using a proposal distribution whose support contains the target's support.Source: §2.3 · p. 5
- KL-adaptive DPG starts with the initial model as its proposal and periodically replaces it with the current policy when that policy is closer to the target.The proposal's closeness to the target is linked to algorithm efficiency.Source: §2.3 · pp. 5–6
- An ablation study reports that the adaptive step obtains faster convergence.The paper's experiments stop parameter optimization when the estimated moment error falls below 0.01.Source: §2.3 · pp. 5–6
3 EXPERIMENTS, RESULTS, AND EVALUATION
The experiments evaluate pointwise, distributional, and hybrid constraints using measures of constraint satisfaction, divergence, fluency, and diversity. GDC performs well across the pointwise comparisons and moves biography-generation statistics toward specified distributional targets.
3.1 EVALUATION METRICS
- Evaluation tracked constraint expectations, forward KL from optimal p, reverse KL from GPT-2, Self-BLEU across samples, and token/type frequencies, alongside Distinct scores.Figure 2 aggregates four metrics across 17 pointwise experiments at every 10 gradient updates.Source: §3.1 · p. 6, Figure 2 · p. 7
3.2 POINTWISE CONSTRAINTS EXPERIMENTS
- The pointwise experiments covered nine single-word constraints, four word lists, and four classifier-based constraints, compared with REINFORCE, REINFORCE_P(x), and ZIEGLER.REINFORCE_P(x) optimizes the energy-model score, whereas ZIEGLER combines reward with a KL penalty from the pretrained model.Source: §3.2 · p. 6
- REINFORCE reached high constraint rewards but diverged sharply from GPT-2 and showed degeneration, while REINFORCE_P(x) concentrated probability on a few sequences and had the baselines’ highest Self-BLEU-5.ZIEGLER had slightly better reward than GDC in the aggregated experiments, but inferior diversity and much worse stability; GDC was superior on target-distribution KL and stability.Source: §3.2 · p. 7
- Across pointwise experiments, GDC smoothly increased constraint satisfaction while maintaining the lowest divergence from GPT-2, the closest fit to target p, and the best overall diversity scores.At 30K steps, the average constraint expectation was still increasing but had not reached 100%.Source: §3.2 · p. 7, Figure 2 · p. 7
3.3 DISTRIBUTIONAL AND HYBRID CONSTRAINTS EXPERIMENTS
- GDC raised female biographies from 7.4% to 35.6%, and shifted science, art, business, and sports expectations to 20.3%, 31.6%, 10.2%, and 11.9%, respectively.The specified targets were 50% female biographies and, for the four professions, 40%, 40%, 10%, and 10%, respectively.Source: §3.3 · p. 8
- In four hybrid experiments, GDC increased profession-specific pointwise satisfaction and female-biography distributional satisfaction simultaneously toward their specified expectations.The constraints required the selected profession for each biography and a 50% female-biography expectation.Source: §3.3 · p. 8
4 DISCUSSION
The discussion presents GDC as a framework for pointwise and distributional control, while emphasizing the remaining challenge of efficiently sampling from its optimal EBM. The appendix explains how transitivity supports incrementally adding constraints and reusing earlier solutions.
4 DISCUSSION
- GDC is presented as the first framework handling both distributional and pointwise constraints, with a two-stage approach that determines an optimal EBM before approximating it with an autoregressive policy.The authors identify efficient sampling from the optimal EBM as the remaining computational challenge.Source: §4 · pp. 8–9
- Table 2 presents experiments on distributional and hybrid constraints as evidence of GDC's generality across these constraint types.Source: Table 2 · p. 9
- GDC's pointwise-constraint method avoids degeneration while satisfying constraints, whereas baselines can approach 100% satisfaction at a severe cost to quality and diversity.Generating a single sentence that satisfies the constraint would trivially yield 100% satisfaction if quality and diversity did not matter.Source: §4 · p. 9
A.2 THE CASE OF POINTWISE CONSTRAINTS IN §2.2
- For purely pointwise constraints, the exponential factor is constant on the allowed support, making the target distribution proportional to a(x)b(x).Source: §A.2 · p. 16
A.3 INCREMENTALLY ADDING NEW CONSTRAINTS
- Adding constraints can be done incrementally: project first onto the manifold satisfying the initial constraints, then project that result onto the manifold satisfying all constraints.The lambdas from the first projection can be reused as the initial lambdas for the expanded constraint set.Source: §A.3 · p. 16, §A.3.1 · p. 16
A.3.1 TRANSITIVITY PROPERTY OF GENERALIZED MAXENT
- The appendix proves transitivity by applying the Pythagorean identity for KL projections, establishing that the staged projection equals the direct projection onto the stricter constraint manifold.Source: §A.3.1 · p. 17
A.3.2 TRANSITIVITY AND AUTOREGRESSIVE POLICY
- For the autoregressive policy, the authors expect that initializing from the policy approximating the earlier projection will help approximate the new target, but they cannot formally prove it.The expected benefit follows because the earlier projection is closer to the new target than the initial distribution; approximation by the policy prevents a formal guarantee.Source: §A.3.2 · p. 17
A.1 COMMENTS ON THEOREM 1
The appendix relates Theorem 1 to earlier information-projection results and details how KL-adaptive DPG assesses proposals when the target is represented by an unnormalized EBM. It estimates the partition function and distribution distances using importance sampling and a moving average.
A.1 COMMENTS ON THEOREM 1
- Theorem 1 reformulates results by Csiszár & Shields (2004), with its properties corresponding to their Remark 3.1 and Theorem 3.2.The cited work assumes finite X, though the passage notes that extensions to infinite countable or continuous spaces are possible.Source: §A.1 · p. 16
B.1 DETAILS ON KL-ADAPTIVITY
- Because the target p is only implicitly represented by the unnormalized EBM P, directly sampling from p to assess proposal similarity is not straightforward.Source: §B.1 · p. 17
- The appendix describes using importance sampling with a sampleable proposal q to estimate the partition function Z and compute KL divergence to a policy.The authors compare KL divergence and total variation distance because they identify no obvious best metric for assessing an importance-sampling proposal.Source: §B.1 · pp. 17–18
- A moving-average estimate of Z uses unbiased per-iteration estimates, pooling samples across training as the proposal q evolves closer to p.The authors say the improving proposal also increases the quality of the importance-sampling estimates of KL divergence and Z.Source: §B.1 · p. 18
B.2 ABLATION ON ADAPTIVITY
The ablation finds that adaptive DPG converges faster than the non-adaptive variant, with the efficiency gap widening for rarer constraints. The KL- and TVD-based adaptive variants have similar convergence rates.
B.2 ABLATION ON ADAPTIVITY
- The ablation compares DPG-KLD, which updates the proposal when the trained policy has lower KL divergence to p, DPG-TVD using total variation distance, and non-Adaptive, which keeps the initial proposal fixed.Source: §B.2 · p. 20
- They run three single-word pointwise experiments at three rarity levels in the original GPT-2 distribution—“Vampire” (1/104), “Paris” (1/103), and “US” (1/102)—using three seeds and 10k gradient updates each.Source: §B.2 · p. 20
- Adaptive DPG converges significantly faster than non-Adaptive, with the efficiency gap growing as constraints become rarer; DPG-KLD and DPG-TVD converge at similar rates.Source: Figure 6 · p. 20, §B.2 · p. 20
C CAN STANDARD SUPERVISION FULLY SATISFY THE CONSTRAINTS?
Standard supervised fine-tuning did not reach high constraint satisfaction without overfitting, while GDC’s optimal distribution differs from the REINFORCE and Ziegler solutions. The results suggest that strict constraints can be difficult to impose while staying close to the initial model.
C CAN STANDARD SUPERVISION FULLY SATISFY THE CONSTRAINTS?
- The supervised result suggests a relationship between reaching 100% constraint satisfaction and overfitting, making strict constraints difficult to impose without moving far from the initial model.The authors note that this would be especially difficult for a hash-based predicate that filters out, on average, one sentence in two.Source: C CAN STANDARD SUPERVISION FULLY… · p. 21
- 0.5625 was the highest constraint expectation achieved by the supervised model selected for best validation loss; higher values coincided with higher validation loss and overfitting.The experiment fine-tuned GPT-2 on samples satisfying the “amazing” constraint.Source: Figure 7 · p. 21, C CAN STANDARD SUPERVISION FULLY… · p. 21
D.1 ILLUSTRATION COMPARING GDC, REINFORCE, AND ZIEGLER
- GDC’s optimal distribution lies on the pointwise constraint manifold while minimizing KL divergence from the initial model.REINFORCE allows any distribution on that manifold, whereas Ziegler’s solution generally assigns positive probability outside it.Source: Figure 8 · p. 22, §D.1 · p. 22
D.2 COMPARISON AGAINST FURTHER BASELINES
In sentiment-control experiments, GDC outperformed PPLM and CTRL on sentiment control and perplexity while achieving comparable diversity. The extended related work contrasts individual-output reward optimization with distributional control of collections of outputs.
D.2 COMPARISON AGAINST FURTHER BASELINES
- GDC is trained without prefix knowledge and can use any prefix at test time, unlike PPLM, which updates GPT-2 hidden states for a given prefix.The comparison used five prefixes and 500 continuations per approach.Source: §D.2 · p. 22
- CTRL’s sentiment control depends on review-specific control codes, restricting generation to positive or negative reviews rather than other text types.At lower temperatures, CTRL mostly generates reviews.Source: §D.2 · p. 23
- GDC achieved better positive/negative sentiment control and lower perplexity than PPLM and CTRL, with comparable diversity.GDC also outperformed PPLM on the Distn diversity metrics in the positive sentiment task.Source: Table 3 · p. 23, §D.2 · p. 23
E RELATED WORK EXTENDED
- Prior global-reward approaches steer autoregressive models toward sequence-level rewards, while weighted decoding approximates global constraints with a biased estimator and heavy search.The cited work includes RL methods across translation, summarization, image-to-text, dialogue, and video captioning.Source: E RELATED WORK EXTENDED · p. 25
- Reward optimization can produce degeneration, and combining fluency rewards with target rewards is described as non-trivial.Prior work also uses KL penalties to limit deviation from the original policy.Source: E RELATED WORK EXTENDED · p. 25
- The paper distinguishes pointwise control of individual outputs from distributional control, where properties of the collection of outputs are objectives in their own right.It notes that focusing on a single good output can fit conditional seq2seq tasks but is ill-suited to language models where sample diversity is required.Source: E RELATED WORK EXTENDED · pp. 25–26
- Earlier text EBM applications include adding global factors to autoregressive LMs to lower training-data perplexity and inference with pretrained non-autoregressive translation models.The related work also cites a survey of EBMs for text.Source: E RELATED WORK EXTENDED · p. 26
F HYPERPARAMETERS AND TRAINING DETAILS
Experiments used pretrained GPT-2 small, with periodic sampling for policy evaluation and task-specific classifiers or word lists for constraint measurements.
- All experiments started from pretrained GPT-2 small with 117M parameters and fine-tuned for 3K gradient-update steps.Training used two Nvidia V100 GPUs; the longest model took about 72 hours.Source: F HYPERPARAMETERS AND TRAINING DETAILS · p. 27
- Policy evaluation sampled 2,048 sequences of 40 tokens every 10 minibatch updates using nucleus sampling with top-p = 0.9.A separate set of 2,048 sequences was sampled for accurate DKL-based metric estimates.Source: F HYPERPARAMETERS AND TRAINING DETAILS · p. 27
- Distributional and hybrid experiments used GPT-2 small fine-tuned on 700K Wikipedia biographies, with a pronoun-based female-biography discriminator and profession word lists.Sentiment and clickbait measurements used pretrained classifier heads over GPT-2 medium; pointwise experiments used four word lists from the PPLM repository.Source: F HYPERPARAMETERS AND TRAINING DETAILS · p. 27
G DISTRIBUTIONAL AND HYBRID CONTROL EXPERIMENTS FOR DEBIASING LANGUAGE MODELS
GDC experiments target demographic and profession expectations with distributional or hybrid constraints. The reported examples show changes in generated biographies toward the specified targets.
G DISTRIBUTIONAL AND HYBRID CONTROL EXPERIMENTS FOR DEBIASING LANGUAGE MODELS
- Using top-p sampling (p = 0.9) on 20,480 examples, GPT-2bio generated around 7% female biographies, with profession proportions of 1% Science, 10% Art, 10% Business&Politics, and 20% Sports.Women’s biographies make up only 18.58% of English Wikipedia biographies.Source: G DISTRIBUTIONAL AND HYBRID CONTROL… · p. 28
- With a 50% female-biography target, GDC increased female biographies from 7% to 36.7% in the single distributional-constraint experiment.The passage contrasts this with a pointwise target of 100%, which would maximize female biographies at the expense of other demographics.Source: Figure 9 · p. 28
- Under four distributional constraints, science biographies rose from 1.5% to 20.3%, art from 10% to 31.6%, business stayed near its 10% target, and sports fell from 19.6% to 11.9%.The targets were 40% for science and art and 10% for sports and business.Source: Figure 10 · p. 29
- With female-biography and pointwise art constraints, female biographies increased from 7.4% to 36.6% and Art professions from 11.4% to 88.6%.The specified targets were 50% for female biographies and 100% for Art.Source: Figure 11 · p. 29
- With female-biography and pointwise sports constraints, female biographies increased from 7.4% to 31.9% and Sports professions from 17.5% to 92.9%.With pointwise business constraints, the corresponding values were 7.4% to 37.7% for female biographies and 10.1% to 82.4% for Business professions.Source: Figure 12 · p. 29, Figure 13 · p. 30
- With female-biography and pointwise science constraints, female biographies increased from 7.4% to 28.8% and Science professions from 1.2% to 74.7%.The specified targets were 50% for female biographies and 100% for Science.Source: Figure 14 · p. 30
H EXTRA DETAILS ON POINTWISE EXPERIMENTS
The appendix compares GDC and three baselines by how quickly they approach the desired distribution across single-word, word-list, and classifier-based controls. In each setting, the captions report better convergence for GDC.
H.1 APPROXIMATING THE DESIRED DISTRIBUTION
- For single-word constraints spanning three rarity levels, GDC shows much better convergence than the three baselines when approximating the desired distribution.The curves cover nine constraints with rarity levels 1/100, 1/1000, and 1/10000.Source: Figure 15 · p. 32
- For kitchen, fantasy, politics, and computers word-lists, GDC shows much better convergence than the three baselines in approximating the desired distribution.Source: Figure 16 · p. 33
- Across four sentiment and click-bait classifiers, GDC shows much better convergence than the three baselines in approximating the desired distribution.The classifiers cover very positive, positive, very negative sentiment, and click-bait.Source: Figure 17 · p. 33
H.2 MORE DETAILS ON POINT-WISE CONSTRAINTS EXPERIMENTS
The figures track evaluation metrics over training steps for single-word, word-list, and classifier-based pointwise controls. They cover nine single-word constraints and multiple word-list and classifier settings.
H.3 POINTWISE CONSTRAINTS
- The figures plot evaluation metrics against training steps for the kitchen, fantasy, politics, and computers word-list constraints.Politics is also the constraint shown in Figure 31.Source: Figure 27 · p. 43, Figure 28 · p. 44, Figure 29 · p. 45
- The figures plot evaluation metrics against training steps for positive-sentiment, very-negative-sentiment, and click-bait classifier controls.Source: Figure 32 · p. 48, Figure 33 · p. 49, Figure 34 · p. 50
H.4 TOKEN FREQUENCY ANALYSIS
Token-frequency plots and generated examples examine diversity across 17 pointwise experiments. The analysis reports very low diversity for REINFORCE baselines and highest diversity for GDC among the baselines.
H.4 TOKEN FREQUENCY ANALYSIS
- Across 17 pointwise experiments, the analysis filters 4,000 constraint-satisfying sequences per policy for token-frequency comparisons.The accompanying plots show token frequency against token rank, with a longer tail indicating more diverse generations.Source: §H.4 · p. 51, Figure 35 · p. 51, Figure 36 · p. 52
- Vanilla REINFORCE has very low generation diversity, with degeneration examples composed of a few repeated tokens.REINFORCEP(x) often concentrates sequence probability on a single, often fluent constraint-satisfying sequence; ZIEGLER also has low sample diversity in 5 of 17 experiments, while GDC has the highest diversity among baselines.Source: §H.4 · p. 51
H.5 GENERATION EXAMPLES
- The single-word example tables show sampled generations, constraint satisfaction, repetition counts, and token-frequency highlighting for seven words.The words are Wikileaks, Vampire, amusing, Paris, restaurant, amazing, and Canada.Source: Table 8 · p. 53, Table 9 · p. 54, Table 10 · p. 55
- The single-word example table for China and the table for US show sampled generations, constraint satisfaction, repetition counts, and token-frequency highlighting.Source: Table 15 · p. 60, Table 16 · p. 61
- The word-list example tables show sampled generations, repetition counts, and token-frequency highlighting for kitchen, fantasy, politics, and computers.Source: Table 17 · p. 62, Table 18 · p. 63, Table 19 · p. 64
- The classifier-based example tables show sampled generations, repetition counts, and token-frequency highlighting for very positive, positive, and very negative sentiment and click-bait controls.Source: Table 21 · p. 66, Table 22 · p. 67, Table 23 · p. 68
Related papers
This paper cites
- Language Models are Few-Shot Learners2020
- PyTorch: An Imperative Style, High-Performance Deep Learning Library2019
- HuggingFace's Transformers: State-of-the-art Natural Language Processing2019
- How NOT To Evaluate Your Dialogue System: An Empirical Study of Unsupervised Evaluation Metrics for Dialogue Response Generation2016
Cited by
- Combating Instruction Conflict via Energy-Driven Latent Conflict Detection2026
- The Reward Was in Your Data All Along: Correcting Flow Matching with Discriminator-Guided RL2026
- Understanding the Effects of RLHF on LLM Generalisation and Diversity2023
- Ethical and social risks of harm from Language Models2021