Contrastive Search for Human-Level Text Generation in Transformers

Hugging Face has introduced Contrastive Search, a state-of-the-art decoding method for neural text generation now available in the transformers library for both PyTorch and TensorFlow. Contrastive Search addresses the critical trade-off between deterministic and stochastic decoding, enabling off-the-shelf language models to generate human-level, semantically coherent text without the repetitions common in greedy search or the incoherence found in nucleus sampling.

The Problem: Model Degeneration vs. Semantic Incoherence

Existing decoding methods generally fall into two categories, both of which possess inherent flaws that degrade text quality:

Deterministic Methods (Greedy and Beam Search)

Deterministic methods select the continuation with the highest likelihood. This frequently leads to model degeneration, where the generated text becomes unnatural and contains undesirable repetitions. For example, using GPT-2 Large with greedy search often results in the model repeating the same sentences multiple times.

Stochastic Methods (Top-k and Nucleus Sampling)

Stochastic methods introduce randomness to avoid repetition. While nucleus sampling (top-p) can eliminate repetitions, it often fails to maintain semantic coherence. This results in generated phrases that are logically disconnected from the prefix text. While lowering the temperature can mitigate this, it pushes the model back toward greedy search, creating a difficult trade-off between repetition and incoherence.

How Contrastive Search Works

Contrastive Search optimizes token selection by jointly considering model confidence and a degeneration penalty to ensure the output is both likely and distinct from previous context.

The Decoding Objective

When selecting the next token $x_{t}$ given a prefix $x_{<t}$, the method evaluates a set of top-k predictions ($V^{k}$). The selection is based on two primary components:

  1. Model Confidence: The probability of the candidate token $v$ as predicted by the language model.
  2. Degeneration Penalty: A measure of how discriminative the candidate $v$ is relative to the previous context $x_{<t}$. This is calculated as the maximum cosine similarity between the token representation of $v$ (computed by the model given the concatenation of the prefix and $v$) and the representations of all tokens already in the context.

These components are balanced by a hyperparameter $\alpha$. If $\alpha = 0$, the method reverts to vanilla greedy search. A higher $\alpha$ increases the penalty for tokens that are too similar to the existing context, thereby preventing repetitions.

Performance and Visual Evidence

Contrastive Search produces text that is grammatically fluent, semantically coherent, and factually grounded. In tests using GPT-2 Large and Meta's OPT-1.3b, the method generated long-form documents (up to 512 tokens) that maintained a consistent narrative and avoided the repetitive loops seen in greedy search.

Token Similarity Visualization

Visual analysis of token similarity matrices reveals the technical success of the degeneration penalty:

  • Greedy Search: Displays high similarity scores in off-diagonal entries, indicating the presence of repeated tokens and patterns.
  • Contrastive Search: High similarity scores appear primarily in diagonal entries, verifying that the degeneration problem has been successfully mitigated.

Implementation in Transformers

Contrastive Search is integrated into the transformers library (version 4.24.0 and above). Users can implement it using the generate method by specifying the following hyperparameters:

  • penalty_alpha: The $\alpha$ hyperparameter regulating the importance of the degeneration penalty.
  • top_k: The number of top predictions to consider during the search.

Example implementation for a GPT-2 model:

output = model.generate(input_ids, penalty_alpha=0.6, top_k=4, max_length=512)

Research Foundations

The implementation is based on two primary research papers:

  1. A Contrastive Framework for Neural Text Generation (NeurIPS 2022), which proposed the original framework.
  2. Contrastive Search Is What You Need For Neural Text Generation (2022), which demonstrated the method's effectiveness across 16 different languages using off-the-shelf models.

Sources