Source-linked AI summary
Quark: Controllable Text Generation with Reinforced Unlearning
Ximing Lu, Sean Welleck, Jack Hessel, Liwei Jiang, Lianhui Qin, Peter West, Prithviraj Ammanabrolu, Yejin Choi
TL;DR
Language models can reproduce toxic, repetitive, or otherwise unwanted behaviors, motivating methods that unlearn them without losing language quality. Quark alternates reward-guided exploration, quantile conditioning, and KL-regularized language-model training; across toxicity, sentiment, and repetition tasks, it outperforms strong baselines and PPO using standard language-modeling primitives.
Problem
Language models often generate toxic, repetitive, or otherwise undesirable text, while existing data-curation and detector-based approaches make post-hoc unlearning difficult.
Method
Quark alternates sampling, reward-based quantization with reward tokens, and conditional likelihood training regularized by KL divergence to the original model.
Results
Across toxicity, sentiment, and repetition tasks, Quark more effectively unlearns undesirable behaviors while preserving fluency and diversity than prior methods.
Takeaways & Limitations
Quark provides a simple reward-optimization approach for controllable unlearning that relies on standard language-modeling primitives.
Takeaways & Limitations
Quark inherits biases and limitations from its reward functions, including the Perspective API for toxicity.
Abstract
from arXiv · showhide
Large-scale language models often learn behaviors that are misaligned with user expectations. Generated text may contain offensive or toxic language, contain significant repetition, or be of a different sentiment than desired by the user. We consider the task of unlearning these misalignments by fine-tuning the language model on signals of what not to do. We introduce Quantized Reward Konditioning (Quark), an algorithm for optimizing a reward function that quantifies an (un)wanted property, while not straying too far from the original model. Quark alternates between (i) collecting samples with the current language model, (ii) sorting them into quantiles based on reward, with each quantile identified by a reward token prepended to the language model's input, and (iii) using a standard language modeling loss on samples from each quantile conditioned on its reward token, while remaining nearby the original language model via a KL-divergence penalty. By conditioning on a high-reward token at generation time, the model generates text that exhibits less of the unwanted property. For unlearning toxicity, negative sentiment, and repetition, our experiments show that Quark outperforms both strong baselines and state-of-the-art reinforcement learning methods like PPO (Schulman et al. 2017), while relying only on standard language modeling primitives.
1 Introduction
Large language models can generate toxic, repetitive, or otherwise undesirable text because pretraining likelihood objectives reproduce problematic web data. Quark addresses the difficulty of post-hoc behavioral unlearning while preserving coherent and fluent generation.
- Pretrained language models may generate toxic text reflecting social biases or repetitive, dull language.
- Post-hoc unlearning is needed because undesirable behaviors are diverse and difficult to avoid, control, or specify in advance.
- Curating desirable-behavior data is difficult, and additional maximum-likelihood training can cause overfitting or loss of general capabilities.
- Behavior detectors score complete samples rather than token-level outputs, making their preferences difficult to optimize directly with standard differentiation.
- Quark uses reward-based unlearning while constraining the model near its original distribution through a KL-divergence penalty.
2 Quark: Quantized Reward Konditioning
Quark alternates exploration, reward quantization, and conditional language-model learning to optimize desired properties while remaining close to the pretrained model. High-reward tokens guide both training and generation, using standard language-modeling primitives with a KL constraint.
- Quark training procedure: Quark iteratively explores with the current model, quantizes samples by reward, and learns from each reward-conditioned quantile.
- Initialization: The datapool stores input, output, and reward examples, initially sampled from the pretrained model or optionally supplied from supervised data.
- Quantization: Quantization sorts examples by reward into equally sized quantiles, each identified by a prepended reward token.
- Learning: Learning maximizes conditional likelihood on quantized samples while applying a KL penalty that limits deviation from the original model.
- Exploration and evaluation: During exploration and evaluation, conditioning on the highest-reward token samples completions expected to have higher reward.
- Relationship to prior work: Unlike PPO, Quark quantizes rewards instead of subtracting a baseline value function and avoids policy clipping and adaptive KL-schedule hyperparameters.
- Relationship to prior work: Quark uses reward-token embeddings as iteratively updated control codes, extending ideas from Decision Transformer and control-code methods.
3 Experiments
Quark is evaluated for unlearning toxicity, unwanted sentiment, and repetitive text while preserving language quality. Across these experiments, it generally outperforms prior methods and maintains fluency and diversity comparable to vanilla GPT-2.
- 3.1 Unlearning Toxicity from Language Models: Quark substantially reduces toxic completions compared with all baselines in both in-domain and out-of-domain settings.Human evaluations also rate Quark generations as less toxic, more topical, and more fluent than all other baselines.
- 3.1 Unlearning Toxicity from Language Models: Quark reduces toxicity while maintaining fluency and diversity at levels similar to vanilla GPT-2.The evaluation uses Perspective API toxicity metrics, GPT2-XL perplexity for fluency, and normalized unique n-gram counts for diversity.
- 3.2 Steering Away from Unwanted Sentiment of Generated Texts: The sentiment experiments train on OpenWebText prompts and evaluate neutral, positive, and negative test prompts using automatic metrics and pairwise human judgments.The automatic evaluation reports the percentage of positive continuations among 25 generations per prompt.
- 3.2 Steering Away from Unwanted Sentiment of Generated Texts: Quark steers generations away from both positive and negative unwanted sentiment more effectively than other baselines while remaining as fluent and diverse as vanilla GPT-2.Human evaluations consistently favor Quark for desired sentiment, topicality, and fluency.
- 3.3 Unlearning Degenerate Repetition: Figure 2 tracks WIKITEXT-103 validation performance across training steps and compares Quark with and without unlikelihood loss.The additional unlikelihood loss accelerates reward optimization, while Quark without it steadily improves reward.
- 3.3 Unlearning Degenerate Repetition: Combining Quark with unlikelihood loss decreases repetition beyond either method alone and improves human-rated fluency, coherence, and overall quality by 35%, 27%, and 29%.Quark without unlikelihood generally outperforms MLE and SimCTG, while unlikelihood alone remains stronger than Quark alone on repetition.
4 Model Ablations
Quark’s ablations show that its KL penalty, quantile partitioning, exploration strategy, control codes, and exploration frequency govern the trade-off between reward optimization and language quality.
- KL term: The exact token-level KL provides a controllable trade-off between language quality and reward maximization, unlike the point-wise approximation, which hurts both.The KL coefficient β = .05 is the default; increasing it improves language quality but slightly increases toxicity.
- Number of quantiles: More quantiles improve reward maximization and reduce toxicity, while slightly worsening language quality as the model moves farther from the original.Finer partitions raise the average reward in the best quantile, making high-reward conditioning more effective.
- Learning quantiles: Training only on the highest-reward quantile improves reward maximization and lowers toxicity but significantly reduces fluency and diversity.The authors speculate that this trade-off results from overfitting to the highest-reward sequences.
- Exploration control: Conditioning exploration on uniformly sampled reward tokens produces worse reward maximization and much higher toxicity than conditioning on the best reward token.Best-token conditioning focuses exploration on promising regions, whereas uniform exploration reduces discovery of better trajectories for the data pool.
- Control codes: Using the initial policy for exploration or removing control codes during training and exploration decreases performance, even when high-reward samples enter the data pool.These comparisons support the role of reward-code conditioning in both exploration and learning.
- Training dynamics: Across training, toxicity decreases monotonically for every quantile and is lower in better quantiles at any given iteration.Figure 6 tracks toxicity probability over training iterations from the worst to the best quantiles.
- Exploration frequency: With a fixed number of gradient updates, more exploration lowers toxicity and increases diversity, while perplexity first decreases and then increases.The authors associate the initial perplexity decrease with less overfitting and the later increase with the quality–reward trade-off.
5 Related Work
Quark builds on reinforcement learning, transformer-based offline decision methods, and prior approaches for modifying or controlling undesirable language-model behavior.
- Reinforcement Learning in NLP: Reinforcement learning has been applied broadly across classical NLP tasks and generation tasks including question answering, summarization, and machine translation.The related work situates reward optimization within a wide range of NLP applications.
- Reinforcement learning with transformers: Trajectory Transformer and Decision Transformer use transformers to generate high-reward action sequences from fixed datasets, unlike Quark, which learns through interaction with the environment.These methods are described as offline reinforcement learning approaches with pre-specified trajectories.
- Unlearning undesirable behaviors: Prior work addresses undesirable language-model behavior through post-hoc reinforcement learning, pre-training interventions, and inference-time controls such as prompts.Quark is positioned among methods that modify or control model behavior after or during training.
6 Conclusion
The paper presents Quark as an effective method for unlearning toxicity, repetition, and unwanted sentiment while preserving fluency and diversity, but emphasizes that it does not guarantee safety.
- Conclusion: Quark more effectively unlearns toxicity, repetition, and unwanted sentiment than prior work without sacrificing fluency and diversity.The conclusion also notes ablation studies examining the effects of model components.
- Risks and limitations: Quark may inherit biases from its reward-scoring process, and reward functions could be used to operationalize pernicious behaviors.The authors identify dual-use risks for controlled generation techniques.
- Risks and limitations: The paper does not present Quark as a guarantee of safety, zero toxicity, or outputs free from negative social biases.It is framed instead as a tool for encouraging language generators to behave in specific ways.
- Future directions: Future directions include simultaneous control of multiple rewards, more diverse human-preference rewards, and reducing the number of optimized parameters.These directions extend the method’s reward scope and seek more parameter-efficient training.
7 Additional Ethical Considerations
Quark raises dual-use concerns because it can steer models toward malicious behaviors, while reward functions may encode subtle social biases. The authors describe inference-time removal of reward-token embeddings as a partial mitigation and emphasize that Quark inherits reward-function limitations.
- Quark could steer language models toward malicious behaviors, like other controllable text-generation techniques.
- Removing embedding parameters for selected reward-token buckets can disable conditioning on those quantiles in public releases.The authors caution that this does not fully mitigate undesirable behavior.
- Reward functions may encode subtle social biases, especially when implemented as black-box APIs or difficult-to-interpret neural networks.
- For toxicity unlearning, Quark checkpoints inherit the biases and limitations of the Perspective API.Human evaluations were conducted on REALTOXICITYPROMPTS, but the authors state that Quark is not a panacea.
- Quark is presented as a tool for generating higher-reward outputs, with improvements dependent on more accurate, specific, and inclusive classifiers.
Checklist
The checklist reports coverage of reproducibility, ethics, assets, and human-subject procedures. It also records that stable confidence intervals were not obtained because computational resources limited repeated runs and random seeds.
- No stable confidence intervals were reported because the experiments lacked enough cross-validation splits and random seeds.The authors attribute this limitation to computational resource constraints and point to extensive ablations instead.
- The checklist states that the paper cites existing assets, does not release new datasets, and refers readers to original releases for license information.
- The authors state that code enabling Quark’s use will be released after internal review.The checklist says the code was planned for release before NeurIPS 2022.
- The experiments use public data, report consent and privacy considerations, and state that existing corpora do not contain personally identifying information to their knowledge.
- Crowdworking procedures included participant instructions, risk discussion, anonymized WorkerIDs, and no collection of personal information.
- Participants were paid over $15 USD per hour on average, with compensation details provided in the appendix.
A Human Evaluation Details
Human evaluations compare Quark with baselines across toxicity, sentiment, and WIKITEXT-103 completion quality. They use qualified raters, multiple model-generation comparisons, and validation or agreement measurements to assess reliability.
- Toxicity and topicality: 1,200 toxicity and topicality comparisons used 100 REALTOXICITYPROMPTS and WRITINGPROMPTS prompts, with three raters per comparison.Each comparison involved six model pairs and two sampled generations from each model.
- Toxicity and topicality: Toxicity evaluations asked which generation was less rude, disrespectful, or unreasonable, alongside topicality and fluency judgments.
- Sentiment: 2,400 sentiment comparisons used 100 positive and 100 negative prompts, with three raters per comparison.Each comparison involved six model pairs and two sampled generations from each model.
- Sentiment: Sentiment evaluations asked which generation had more positive or negative sentiment, while also rating topicality and fluency.
- WIKITEXT-103: WIKITEXT-103 evaluations presented Wikipedia snippets with model completions and used five-point Likert ratings for generation quality.The reported dimensions include fluency and overall quality; the interface is shown in Figure 9.
- WIKITEXT-103: 6.6K ratings across 2.2K HITs evaluated 19 algorithms, with ground-truth validation and 10% double annotation.Krippendorff’s α was .42 for overall, .35 for fluency, and .45 for coherence, described as moderate agreement.
B Experimental Details
The experimental appendix details baselines, training configurations, and evaluation setup for toxicity, sentiment, and repetition. It includes hyperparameter searches, GPU budgets, and the objectives used to compare repetition-control methods.
- Baselines: PPLM, GEDI, DAPT, DEXPERTS, and PPO provide comparison baselines for controllable generation, adaptation, decoding, and reward optimization.PPO is implemented by the authors; the other baseline results come from prior work.
- Toxicity and sentiment: Four NVIDIA Quadro RTX 8000 GPUs and about 100 GPU hours supported toxicity and sentiment training.
- Degenerate repetition: WIKITEXT-103 comparisons include MLE, Unlikelihood, and SimCTG, using test-set prefixes and greedy decoding.The objectives respectively represent standard maximum likelihood, unlikelihood training, and contrastive training.
- Degenerate repetition: Repetition experiments fine-tuned the MLE model with Quark and searched quantiles from [2, 10] and KL coefficient β from [0, 0.3].Training used four NVIDIA Quadro RTX 8000 GPUs and about 600 GPU hours.
C Details for Quark Implementation
Quark adds reward tokens to the language-model vocabulary and prepends them to prompts, while its loss combines conditioned language-model loss with a KL penalty. The implementation uses standard language-model libraries and, in practice, batched left-to-right training with Adam.
- Quark augments the vocabulary with K reward tokens and prepends the selected token to each prompt.The conditioned input is represented as (r_k; x).
- Quark’s loss can be implemented in standard language-modeling libraries such as Transformers and PyTorch.The supplied implementation uses AutoModelForSeq2SeqLM instances for the reference and conditioned models.
- The implementation combines the conditioned model loss with a KL-divergence term weighted by beta.The code returns quark.loss + beta * kl.
- In practice, the authors use a batched loss for left-to-right models such as GPT-2 and update parameters with Adam.Seq2SeqLMs are shown only for brevity in the example implementation.
D Qualitative results
The qualitative-results section presents example generations comparing Quark with baselines for toxicity, negative sentiment, and degenerate repetition.
- Table 10 presents example generations from unlearning toxicity with Quark and baselines.
- Table 11 presents example generations from unlearning negative sentiment with Quark and baselines.
- Table 12 presents example generations from unlearning degenerate repetition with Quark and baselines.