Repetition Penalty - LLM Generation & Decoding

Repetition Penalty

What you'll build

The CTRL repetition penalty (Keskar et al., 2019), as implemented by Hugging Face's RepetitionPenaltyLogitsProcessor: make every token that already appeared in the output less likely, by shrinking its logit.

Background

Small or greedily decoded models fall into loops ("I think that I think that I think..."). The penalty pushes the logits of already-generated tokens down before sampling. It is applied to the raw logits, before temperature, top-k and top-p.

Formula

For every distinct token id t in generated_tokens:

logits[t] = logits[t] / penalty    if logits[t] > 0
logits[t] = logits[t] * penalty    otherwise (logits[t] <= 0)

With penalty > 1, both branches lower the logit: a positive logit moves toward 0, a negative one moves further from 0. Dividing a negative logit instead would raise it, which is the classic bug.

Inputs

• logits: float array (vocab_size,)
• generated_tokens: list of int token ids (may contain repeats, may be empty)
• penalty: float >= 1.0, default 1.2

Expected Return Format

penalized_logits: new float array (vocab_size,); the input array must not be modified.

Edge cases and conventions

• Each distinct token is penalized once, however often it appears: [0, 0, 0] penalizes token 0 once, as in CTRL and Hugging Face. (For a count-based penalty, see the follow-up.)
• penalty = 1.0 leaves the logits unchanged.
• A logit of exactly 0 stays 0.

Constraints

NumPy only. Copy the input first (np.array(logits, dtype=float) or .copy()).

Follow-up

• OpenAI's presence_penalty and frequency_penalty subtract alpha_presence * [count > 0] + alpha_frequency * count instead. Implement it; how is it different from dividing?
• Why can a strong repetition penalty hurt code generation?

Input: apply_repetition_penalty(np.array([2.0, -1.0, 3.0, 0.5]), [0, 1], penalty=1.5)
Output: [1.3333, -1.5, 3.0, 0.5]

Input: apply_repetition_penalty(np.array([2.0, -1.0, 3.0]), [0, 0, 1, 1, 1], penalty=2.0)
Output: [1.0, -2.0, 3.0]

import numpy as np

def apply_repetition_penalty(logits, generated_tokens, penalty=1.2):
    """
    Apply the CTRL repetition penalty to previously generated tokens.

    Args:
        logits: Raw logits (vocab_size,)
        generated_tokens: List of previously generated token ids (may repeat)
        penalty: Repetition penalty factor (>= 1.0)

    Returns:
        penalized_logits: New array (vocab_size,); each distinct generated
                          token: positive logit / penalty, else * penalty
    """
    # Your code here
    pass