A Contrastive Framework for Neural Text Generation

In plain words

A language model that writes text tends to get stuck. It repeats the same phrase, or it loops back on what it just wrote. The authors trace much of this to the way the model stores words internally: the stored forms of different words end up bunched together, so the model struggles to tell them apart.

They offer two fixes that work as a pair. The first is a training change: an extra term that pushes the stored forms of different words apart while the model learns. The second is a new way of choosing the next word while writing. The writer looks only at a short list of likely next words, and it prefers the one that is likely but does not look too much like the text already written.

The idea is simple, and the authors report that it gives more varied and more on-topic text than the usual choices, which are picking the single most likely word or sampling at random.

TL;DR

• Degenerate text (repetition, dull loops) comes from token representations that are too similar to each other, which the paper calls anisotropy §1.
• SimCTG = MLE loss plus a margin-based contrastive loss on token states, with no extra weights §3.1, eq. 4.
• Contrastive search chooses among the top-k next tokens by trading model probability against a similarity penalty, with a single weight alpha §3.2, eq. 5.
• Author-reported on WikiText-103 with GPT-2 small: best or tied-best diversity, MAUVE and coherence among the decoders tested §4.2, tab. 1.
• Author-reported human evaluation: SimCTG-large with contrastive search is highest, or tied highest, on all three criteria §4.3, tab. 2.

Problem

Maximum-likelihood training followed by greedy or beam search yields dull, repetitive open-ended text §2.2. Random sampling such as nucleus sampling reduces repetition but can drift away from the prefix and contradict it §2.2. Unlikelihood training changes the output distribution to suppress repeats, and the paper reports that this lowers perplexity and next-token accuracy §4.2. The paper's claim is that the root cause sits in the representation space: the last-layer token states of GPT-2 have cosine similarity above 0.95 within a sentence §1, fig. 1, so similar states lead to repeated tokens. This claim is the paper's own interpretation of its figures; it is not proven.

Method

Contrastive training (SimCTG)

The model is trained on the usual next-token cross-entropy (MLE) plus a contrastive term over the last-layer hidden states of all tokens in the same sequence §3.1, eq. 1, eq. 2. For every ordered pair of distinct positions (i, j), the loss adds a hinge: it is zero when the similarity of i to j is at least a margin rho below the self-similarity of i, and positive otherwise eq. 2. The per-sequence sum is divided by T(T-1), the number of ordered pairs eq. 2. The total loss is MLE plus this term, with no weight eq. 4. When rho = 0 the contrastive term is zero and SimCTG reduces to MLE §3.1. The paper states that the best margin in its experiments is rho = 0.5 §6.2.

The paper defines the similarity as cosine similarity eq. 3. The contrastive term adds no new parameters; the paper reports about a 1.5% training-compute overhead over MLE app. I, tab. 8.

# Contrastive term over one sequence  [eq. 2] [§3.1]
# Input:  hidden (T, d) last-layer token states; margin rho
# Output: scalar L_CL

def contrastive_loss(hidden, margin):
    # Cosine similarity between every pair of token states  [eq. 3]
    unit = hidden / norm(hidden, axis=-1)     # (T, d)
    sims = unit @ unit.T                      # (T, T)
    total = 0
    for i in range(T):
        for j in range(T):
            if j == i:
                continue
            # Hinge: token j should sit at least rho below token i's
            # self-similarity, which is 1.0  [eq. 2]
            total += max(0, margin - sims[i, i] + sims[i, j])
    # Average over the ordered pairs T * (T - 1)  [eq. 2]
    return total / (T * (T - 1))

# SimCTG objective  [eq. 4] [§3.1]
# Input:  logits (T, V), targets (T,), hidden (T, d), margin rho
# Output: scalar loss

def simctg_loss(logits, targets, hidden, margin):
    # Next-token cross-entropy, the MLE term  [eq. 1]
    mle = cross_entropy(logits, targets)
    # Contrastive term  [eq. 2]
    cl = contrastive_loss(hidden, margin)
    # Plain sum, no weighting  [eq. 4]
    return mle + cl

Contrastive search (decoding)

At each step the decoder takes the top-k tokens of the next-token distribution §3.2. For each candidate v it computes a score from two parts: the model probability of v, and a penalty equal to the largest cosine similarity between v's representation and the representations of all earlier tokens in the context §3.2. The candidate's representation is the last-layer state produced when v is appended to the context §3.2. The score is (1 - alpha) times the probability minus alpha times the penalty, and the candidate with the highest score is emitted eq. 5. With alpha = 0 the rule becomes greedy search §3.2. The paper says k is typically 3 to 10 §3.2.

The paper describes no algorithm box; the loop is given by eq. 5 alone.

# Contrastive search  [eq. 5] [§3.2]
# Input:  prefix tokens, k (candidates), alpha, decoding length L
# Output: continuation tokens

def contrastive_search(prefix, k, alpha, L):
    # Context holds everything written so far, prefix included
    context = prefix
    for step in range(L):
        # Top-k next tokens by model probability  [§3.2]
        candidates = topk(next_token_probs(context), k)
        for v in candidates:
            # Last-layer state of v when appended to the context  [§3.2]
            h_v = last_hidden_state(context + [v])
            # Penalty: largest cosine to any earlier token  [eq. 5]
            penalty = max(cosine(h_v, h_x) for h_x in states(context))
            # Model confidence minus weighted penalty  [eq. 5]
            score[v] = (1 - alpha) * prob(v) - alpha * penalty
        # Keep the best-scoring candidate  [eq. 5]
        next_token = argmax(score)
        context = context + [next_token]
    return context[len(prefix):]

Setup used for the document experiments

GPT-2 small (117M, 12 layers, 12 heads) is fine-tuned on WikiText-103 with SimCTG, MLE, or the unlikelihood objective §4. The paper says the method is architecture-agnostic §4; the authors list other architectures as future work app. A. The paper also reports runs on Transformers trained from scratch and on GPT-2-large app. D.

Hyperparameters

Document generation (WikiText-103) unless marked dialogue:

Dialogue generation (LCCC, DailyDialog):

Results

Author-reported. Setup: WikiText-103 test set; GPT-2 small (117M) unless noted; baselines are MLE and unlikelihood training with greedy, beam, nucleus and contrastive decoding. The gen-ppl column is measured with the SimCTG model for all methods §4.1.2, which the paper says gives the same conclusions as the other models app. F.

Language-model and generation metrics tab. 1:

Source: tab. 1. Columns are relabelled here; "Top-1 in prefix" is the paper's prediction-repetition rate, and the 2-, 3- and 4-gram columns are the generation repetition rates. The paper's bold marks SimCTG for ppl and acc (model-level numbers), SimCTG + contrastive for rep-2, rep-3, rep-4, diversity, MAUVE, coherence and gen-ppl, and the unlikelihood model for rep and wrep tab. 1, §4.2.

Gen-ppl measured by other models app. F, tab. 6 and tab. 7: with the MLE model, SimCTG + contrastive is 20.53 and MLE + contrastive is 9.43 against 24.86 for human text; with the unlikelihood model, the values are 23.47 and 15.94 against 29.62 for human text. The paper concludes that SimCTG + contrastive is closest to human gen-ppl under both measurers app. F.

Human evaluation on WikiText-103 tab. 2: 5-point Likert, 200 prefixes, five graders; stars mark significance versus nucleus sampling (sign test, p < 0.05) §4.3.

(*) significant vs. nucleus sampling with the same model. The paper states that SimCTG-large + contrastive is comparable to human text on fluency (sign test p > 0.4) §4.3.

Dialogue generation tab. 3: 200 contexts, five annotators, 5-point Likert. Stars mark significance versus greedy, beam and nucleus sampling §5.

Other results the paper reports:

• Other model sizes, WikiText-103 app. D, tab. 5. GPT-2-large SimCTG + contrastive: ppl 16.53, acc 43.47, diversity 0.95, MAUVE 0.96, coherence 0.626. Vanilla Transformer (117M, no pretraining) SimCTG + contrastive: ppl 26.55, acc 36.03, diversity 0.91, MAUVE 0.85, coherence 0.566. Conicity and self-similarity for MLE GPT-2-large are 0.46 and 0.20 tab. 5.
• Contrastive search on MLE models §5, app. D, app. H: on LCCC, MLE with contrastive search beats MLE with nucleus sampling on coherence, fluency and informativeness, and the paper says the Chinese and GPT-2-large MLE models are already isotropic §5.
• Training compute app. I, tab. 8: MLE 8.08e16 FLOPs, unlikelihood 8.91e16, SimCTG 8.20e16; 117M parameters for all three.
• Margin sweep §6.2, fig. 3: rho from 0.1 to 1.0; contrastive training lowered perplexity relative to MLE at every value tested, and the paper picks rho = 0.5.
• Latency §6.4, fig. 5: at small beam width b and k (both at most 6), contrastive search and beam search have nearly identical latency; above 6, contrastive search is faster.
• Diverse contrastive search app. L, tab. 11: first 2 tokens sampled with nucleus p = 0.95, the rest by contrastive search.

Paper vs code

Official code: github.com/yxuansu/SimCTG, MIT licence, commit 3a91449c8f3133f2c183e481e705b1dbb268589c. The paper's abstract links it §abstract. The repository contains the document-generation training scripts, the dialogue scripts and a package simctg.

• Loss normalisation. The paper defines the contrastive term per sequence, divided by T(T-1) eq. 2. The code sums the hinge over the whole batch and divides by the total number of valid pairs in the batch loss_func.py. Longer sequences carry more pairs and so weigh more in the batch-level ratio.
• Hinge form. The code computes margin minus (gold self-similarity minus pair similarity), then ReLU loss_func.py. This matches eq. 2.
• Loss weight. The code uses MLE + CL with weight 1 trainer.py. This matches eq. 4.
• Optimizer. The paper says Adam §4. The code uses transformers AdamW with no weight-decay argument trainer.py.
• Schedule and clipping. Neither is stated in the paper's text. The code uses a linear schedule with 10% warmup trainer.py and clips the gradient norm at 1.0 train.py, trainer.py.
• Contrastive search. The code's per-step scoring matches eq. 5. The probability term is the full softmax probability of each top-k candidate, and the penalty is the max cosine to all context states utlisgpt.py. Candidate states are computed with a key-value cache, batched over k utlisgpt.py.
• Which hidden state. The code uses the last entry of the model's hidden_states utlisgpt.py, utlisgpt.py. The paper says "last layer". In the Hugging Face GPT-2 implementation, that entry is after the final layer norm; the paper does not say this.
• Stopping. The package's fast_contrastive_search runs for exactly decoding_len steps and only truncates at an end token when early_stop is set simctggpt.py. The paper does not mention stopping.
• Dialogue data. The paper's dialogue section gives 40k LCCC steps and 5k DailyDialog steps §5. The repository includes only the LCCC script train.sh, which matches the 40k figure. The DailyDialog configuration is not in the repository. The LCCC script also sets --min_len 10 train.sh; the paper does not mention it.
• Metrics code. I searched the Python and shell files for gen-ppl, MAUVE and SimCSE coherence code and found none, so the evaluation numbers cannot be checked against the release. The paper does not name the SimCSE checkpoint it used.
• Hyperparameter selection. The paper says all decoders were tuned on validation MAUVE §4. The repository has no tuning script.
• Efficiency figure. The paper's text gives 1.48% extra training compute over MLE app. I. The FLOP counts in tab. 8 (8.20e16 vs 8.08e16) give about 1.49%. The difference is rounding.

Pitfalls

• The gen-ppl formula is misprinted. The paper defines gen-ppl as 2 to the power f, where f is the average of log2 probabilities with no minus sign §4.1.2. As written, that number lies between 0 and 1. A correct perplexity needs the negative of f. Use the standard definition and not the printed one.
• Cosine similarity of token states is the whole method. Both the training term and the decoding penalty use cosine similarity of last-layer states. Use the same layer and the same normalisation at training and decoding time.
• Margin sensitivity. Too small a margin (0.1) or too large (1.0) gives a worse perplexity than 0.5 §6.2. The paper uses 0.5 for both document and dialogue runs train.sh, train.sh.
• Contrastive search without SimCTG works only in some settings. For English GPT-2 small, MLE + contrastive search is well behind SimCTG + contrastive search tab. 1. The paper finds the MLE-trained Chinese model and GPT-2-large already isotropic, which is why contrastive search works there without SimCTG §5, app. D, app. H. The authors say the Chinese case still needs rigorous testing for other languages app. H.
• Cost per step. Each step scores k candidates, so decoding costs roughly k forward positions per step. The paper's latency comparison is for batch size 1 §6.4. The k-candidate pass is batched in the code utlisgpt.py.
• Penalty memory grows with length. The penalty compares each candidate to every earlier token, so the context states must be kept for the whole output.
• Determinism. Plain contrastive search is deterministic. The paper's stochastic variant (nucleus sampling for the first 2 tokens, then contrastive search) is preliminary, in app. L.
• Batch-level loss ratio. The code's batch-level normalisation (above) differs from the per-sequence definition. Use one convention and state it.
• Human evaluation scale. Five graders from a third-party platform rated 200 prefixes §4.3. The paper reports Fleiss' kappa of 0.51 to 0.70 tab. 2; coherence agreement is the lowest.
• Chinese and dialogue numbers are not on WikiText. The dialogue setting uses contrastive k = 5 §5, while the document setting uses k = 8 §4. Do not carry k across tasks without checking.

Lineage

• Builds on: anisotropy analyses of contextual representations (Ethayarajh 2019 and related work cited in §1) and the unlikelihood objective of Welleck et al. §2.2.
• Compared against: greedy search, beam search (beam 10), nucleus sampling (Holtzman et al.; p = 0.95), unlikelihood training, and MLE training §4.
• Later versions: the paper's v3 (2022-09-26) is the source used here.
• Superseded by: not stated in this wiki.