Three Penalties in LLM Sampling

Three Penalties in LLM Sampling
Photo by Ivan Torres / Unsplash

Hello! This is the Qualiteg Product Development Team.

Today, the whole team went out for lunch at an Italian restaurant. The three-cheese pizza was a big hit.

So today, instead of three cheeses, we would like to introduce three penalties used in LLMs.

ChatStream ships with a rich set of preset sampling algorithms, including penalties. By choosing sampling that suits the model and the purpose, you can generate more natural responses.

The Role of Penalties in Text Generation

What exactly is a "penalty," one of the important concepts in natural language processing (NLP) and the world of LLMs?

What is a penalty?

A penalty is a mechanism used to adjust the appearance of specific tokens (words or strings) when an LLM generates text.

For a text generation model to produce natural and diverse sentences, it needs to avoid repeating the same words or phrases over and over. This is where penalties come in.

It is fairly common for an LLM to generate similar sentences or words repeatedly, so we set appropriate penalties.

Types of Penalties

There are several types of penalties in text generation, but here we will explain three particularly important ones.

  • Repetition Penalty:
    • Scope: All tokens or phrases that appear repeatedly.
    • Penalty criterion: Based on all past tokens.
    • Usage example: Applies a penalty to every token generated so far.
  • Frequency Penalty:
    • Scope: The frequency of occurrence of specific tokens.
    • Penalty criterion: Based on the number of times each token has occurred.
    • Usage example: Tracks how many times a specific token has been generated during generation and applies a penalty based on that frequency.
  • Presence Penalty:
    • Scope: Tokens that have been generated at least once.
    • Penalty criterion: Based on whether a token has been generated at least once.
    • Usage example: Applies a penalty to already-generated tokens when they reappear.

Implementing the Penalties

Now, let's actually implement the penalties.

Implementing penalties as ChatStream sampling classes

A sampling class overrides AbstractLogitsProcessor.

from chatstream.token_samplers.logits_processor import AbstractLogitsProcessor

AbstractLogitsProcessoris a simple abstract class that looks like this:


from abc import ABC, abstractmethod


class AbstractLogitsProcessor(ABC):
    @abstractmethod
    def process(self, logits, params):
        pass

    @abstractmethod
    def get_name(self):
        pass

Repetition Penalty

Let's start by implementing the Repetition Penalty.

We allow two calculation methods, multiplicative and subtractive: if a token has already been generated, its log probability is either divided by the penalty value (multiplicative) or has the penalty value subtracted from it (subtractive). All past tokens are subject to the penalty.

from chatstream.token_samplers.logits_processor import AbstractLogitsProcessor


class RepetitionPenaltyProcessor(AbstractLogitsProcessor):
    """
    A processor that applies a penalty to repeated tokens.

    This class suppresses repetition by adjusting the logits of tokens that have been used in the past.
    The penalty can be applied either multiplicatively or subtractively, as specified by a parameter.

    The basic multiplicative penalty calculation is  logits[token_id] /= penalty , which lowers the logit values
    of token_ids that have appeared before, reducing their probability and thereby suppressing repeated output of the same token.

    """

    def __init__(self):
        pass

    def process(self, logits, params):

        past_tokens = params.get("past_tokens", None)
        penalty = params.get("penalty", None)
        penalty_method = params.get("penalty_method", "multiplicative")

        # Apply the penalty to the logits of past tokens
        if penalty is not None and past_tokens is not None:
            # Create a copy of logits (guarantees the logits passed as an argument are not modified)
            adjusted_logits = logits.clone()
            # Check the type of the penalty value
            if not isinstance(penalty, (int, float)):
                raise ValueError(f"penalty should be a scalar value, but got {penalty}({type(penalty)})")

            # Update logits according to the penalty method
            if penalty_method == "multiplicative":
                if penalty != 1.0:
                    for token_id in set(past_tokens):
                        adjusted_logits[token_id] /= penalty

            elif penalty_method == "subtractive":
                for token_id in set(past_tokens):
                    adjusted_logits[token_id] -= penalty
            else:
                raise ValueError(f"Unknown penalty_method: {penalty_method}")
        else:
            adjusted_logits=logits

        return {"name": "RepetitionPenaltyProcessor", "type": "logits", "logits": adjusted_logits}


    def get_name(self):
        return "rep_penalty"

Frequency Penalty

Frequency Penaltylooks like this.

Each time a token appears, based on the number of times that token has occurred, its log probability is either cumulatively divided by the penalty value (multiplicative) or has the penalty value cumulatively subtracted from it (subtractive).

class FrequencyPenaltyProcessor(AbstractLogitsProcessor):
    """
    A processor that applies a penalty based on the frequency of generated tokens.

    This class tracks how many times each token has appeared during generation and penalizes frequently occurring tokens.
    The penalty can be applied either multiplicatively or subtractively, as specified by a parameter.
    """

    def __init__(self):
        self.token_counts = {}

    def process(self, logits, params):
        penalty = params.get("penalty", None)
        penalty_method = params.get("penalty_method", "multiplicative")

        # Create a copy of logits
        adjusted_logits = logits.clone()

        if penalty is not None:
            if not isinstance(penalty, (int, float)):
                raise ValueError(f"penalty should be a scalar value, but got {penalty}({type(penalty)})")

            for token_id, count in self.token_counts.items():
                if penalty_method == "multiplicative":
                    adjusted_logits[token_id] /= (penalty ** count)
                elif penalty_method == "subtractive":
                    adjusted_logits[token_id] -= (penalty * count)
                else:
                    raise ValueError(f"Unknown penalty_method: {penalty_method}")

        return {"name": "FrequencyPenaltyProcessor", "type": "logits", "logits": adjusted_logits}

    def update_token_counts(self, token_ids):
        for token_id in token_ids:
            if token_id in self.token_counts:
                self.token_counts[token_id] += 1
            else:
                self.token_counts[token_id] = 1

    def get_name(self):
        return "freq_penalty"

Presence Penalty

The Presence Penalty looks like this.

Based on whether a token has been generated at least once, the log probability of any token that has already been generated is either divided by the penalty value (multiplicative) or has the penalty value subtracted from it (subtractive).

class PresencePenaltyProcessor(AbstractLogitsProcessor):
    """
    A processor that applies a penalty based on the presence of generated tokens.

    This class tracks whether a given token has already appeared and penalizes tokens that are present.
    The penalty can be applied either multiplicatively or subtractively, as specified by a parameter.
    """

    def __init__(self):
        self.seen_tokens = set()

    def process(self, logits, params):
        penalty = params.get("penalty", None)
        penalty_method = params.get("penalty_method", "multiplicative")

        # Create a copy of logits
        adjusted_logits = logits.clone()

        if penalty is not None:
            if not isinstance(penalty, (int, float)):
                raise ValueError(f"penalty should be a scalar value, but got {penalty}({type(penalty)})")

            for token_id in self.seen_tokens:
                if penalty_method == "multiplicative":
                    adjusted_logits[token_id] /= penalty
                elif penalty_method == "subtractive":
                    adjusted_logits[token_id] -= penalty
                else:
                    raise ValueError(f"Unknown penalty_method: {penalty_method}")

        return {"name": "PresencePenaltyProcessor", "type": "logits", "logits": adjusted_logits}

    def update_seen_tokens(self, token_ids):
        self.seen_tokens.update(token_ids)

    def get_name(self):
        return "presence_penalty"

Penalties and Other Sampling Parameters

Now, let's also look at how penalties relate to other sampling parameters.

First, a recap of the penalties we implemented above:

Penalties

  1. Repetition Penalty: Prevents repetition of specific tokens or phrases.
  2. Frequency Penalty: Applies a penalty based on the frequency of generated tokens, suppressing tokens that appear often.
  3. Presence Penalty: Prevents tokens that have already been generated from appearing again.

Top-k, Top-p, Temperature

Next up are these three. Three again, as it happens.

These three are especially common sampling techniques.

top-k, top-p, and temperature are parameters that control the quality and diversity of the generated text. In brief:

Top-k

Purpose: Considers only the k most probable tokens.

  • Method: Selects the k highest-probability tokens and randomly picks the next token from among them.
  • Calculation:
    1. Sort the token probabilities in descending order.
    2. Select the top k tokens.
    3. Choose the next token from among them.

Top-p (or Nucleus Sampling)

Purpose: Selects tokens until their cumulative probability reaches p (e.g., 0.9).

  • Method: Selects high-probability tokens until the cumulative probability reaches p, then randomly picks the next token from among them.
  • Calculation:
    1. Sort the token probabilities in descending order.
    2. Select tokens until the cumulative sum of probabilities exceeds p.
    3. Choose the next token from among them.

Temperature

Purpose: Controls the randomness of the generated text.

  • Method: temperature is used to adjust the probability distribution over tokens.
  • Calculation:
    1. Divide each token's log probability by temperature.
    2. This smooths the probability distribution (high temperature: the distribution flattens; low temperature: the distribution sharpens).

A Calculation Scenario with Penalties, top_k, top_p, and temperature

For example, let's consider a scenario where we generate the text Qualiteg May Change the World with ChatStream.

How penalties interact with Top-k, Top-p, and Temperature

Qualiteg May Change the World with ChatStream — in generating this sentence, let's look at how penalties are applied and how they work together with top-k, top-p, and temperature.

1. Applying the penalties

  • Applying the penalties:
    • The Repetition Penalty, Frequency Penalty, and Presence Penalty are applied, adjusting the log probabilities of specific tokens.
    • For example, if the word "Qualiteg" has already appeared several times, its log probability is lowered.

2. Applying Temperature

  • Applying Temperature:
    • The adjusted log probabilities are further scaled by temperature.
    • This changes the shape of the distribution, increasing or decreasing the randomness of the generated tokens.
    • With a high temperature (e.g., 1.2), the distribution flattens and randomness increases.
    • With a low temperature (e.g., 0.7), the distribution sharpens and the highest-probability token becomes more likely to be chosen.

3. Applying Top-k

  • Applying Top-k:
    • top-k is applied, leaving only the top k tokens as candidates.
    • As a result, the next token is chosen from among the most probable tokens.
    • Low-probability tokens are excluded, improving the quality of the generated text.

4. Applying Top-p

  • Applying Top-p:
    • top-p is applied, selecting tokens until the cumulative probability exceeds p.
    • As a result, the next token is randomly chosen from among the high-probability tokens.
    • This too excludes low-probability tokens, improving the quality of the generated text.

A Concrete Example of the Sampling Calculation

For example, let's walk through concretely how the calculation proceeds with the following settings.

  • Settings:

    • Repetition Penalty: 1.2
    • Frequency Penalty: 0.8
    • Presence Penalty: 1.5
    • Temperature: 0.7
    • Top-k: 50
    • Top-p: 0.9
  • Step 1: Compute the initial log probabilities
    First, the model computes the initial log probabilities (logits) for each token. Suppose we obtain the following initial log probabilities:

["Qualiteg": 2.0, "May": 1.5, "Change": 1.0, "the": 0.5, "World": 0.3, "with": 0.2, "ChatStream": 0.1]
  • Step 2: Apply the penalties
    Next, we apply each penalty. In the example below, assume "Qualiteg" has already appeared once and the other tokens are appearing for the first time.

    • Repetition Penalty: Since "Qualiteg" has already appeared, a penalty of 1.2 is applied to logits["Qualiteg"].
    logits["Qualiteg"] /= 1.2
    2.0 / 1.2 ≈ 1.67
    
    • Frequency Penalty: A penalty based on frequency of occurrence is applied. Since "Qualiteg" has appeared once, the frequency penalty is applied.
    logits["Qualiteg"] *= 0.8
    1.67 * 0.8 ≈ 1.34
    
    • Presence Penalty: Since "Qualiteg" is already present, a penalty of 1.5 is applied to logits["Qualiteg"].
    logits["Qualiteg"] /= 1.5
    1.34 / 1.5 ≈ 0.89
    

    Log probabilities after applying the penalties:

    ["Qualiteg": 0.89, "May": 1.5, "Change": 1.0, "the": 0.5, "World": 0.3, "with": 0.2, "ChatStream": 0.1]
    
  • Step 3: Apply Temperature
    Next, we apply temperature to adjust the probability distribution. With temperature set to 0.7, each log probability is divided by 0.7.

    logits["Qualiteg"] /= 0.7
    0.89 / 0.7 ≈ 1.27
    
    logits["May"] /= 0.7
    1.5 / 0.7 ≈ 2.14
    
    logits["Change"] /= 0.7
    1.0 / 0.7 ≈ 1.43
    
    logits["the"] /= 0.7
    0.5 / 0.7 ≈ 0.71
    
    logits["World"] /= 0.7
    0.3 / 0.7 ≈ 0.43
    

    Adjusted log probabilities:

    ["Qualiteg": 1.27, "May": 2.14, "Change": 1.43, "the": 0.71, "World": 0.43, "with": 0.29, "ChatStream": 0.14]
    
  • Step 4: Apply Top-k
    Next, we apply top-k. Here top-k=50, but we show only the top 5 tokens.

    Top tokens:

    ["May": 2.14, "Change": 1.43, "Qualiteg": 1.27, "the": 0.71, "World": 0.43]
    
  • Step 5: Apply Top-p
    Finally, we apply top-p. Here top-p=0.9. Tokens are selected until the cumulative probability exceeds 0.9.

    Computing the cumulative probability:

    • Computing the probabilities:

      • May: exp(2.14) ≈ 8.50
      • Change: exp(1.43) ≈ 4.18
      • Qualiteg: exp(1.27) ≈ 3.56
      • the: exp(0.71) ≈ 2.03
      • World: exp(0.43) ≈ 1.54

      Total: 8.50 + 4.18 + 3.56 + 2.03 + 1.54 ≈ 19.81

    • Computing the cumulative probabilities:

      • May: 8.50 / 19.81 ≈ 0.43
      • Change: 4.18 / 19.81 ≈ 0.21
      • Qualiteg: 3.56 / 19.81 ≈ 0.18
      • Cumulative probability so far: 0.43 + 0.21 + 0.18 ≈ 0.82
      • the: 2.03 / 19.81 ≈ 0.10 (cumulative probability: 0.82 + 0.10 ≈ 0.92)

    Since the cumulative probability has exceeded 0.9, the tokens up to the remain as candidates.

    ["May": 2.14, "Change": 1.43, "Qualiteg": 1.27, "the": 0.71]
    
  • Step 6: Select a token
    Finally, the next token is chosen at random from the remaining tokens. In this example, one is chosen from ["May", "Change", "Qualiteg", "the"].

In this scenario, we have seen how penalties, temperature, top-k, and top-p combine to influence text generation.

By following the actual calculation process, we hope it is now clear how penalties adjust the probability of specific tokens, temperature changes the shape of the distribution, and top-k and top-p narrow down the candidates, ultimately controlling the quality and diversity of the generated text.

Summary

Today, we looked at three types of penalties and the sampling techniques that surround them.

The PenaltyProcessor we built today can also be incorporated into ChatStream (although ChatStream already includes preset PenaltyProcessor implementations). Fundamentally, how logits are sampled is up to the service provider, so you can combine any processors you like in any order you like.
How to combine sampling classes (functions) and what actual values to use in order to obtain the most desirable output for a given model is something you determine by measuring the model's actual inputs and outputs. This area is a real accumulation of know-how, so if you are interested, please feel free to consult Qualiteg.

Comparison of Penalties

[Appendix] Comparison of Penalties

Penalty type Purpose How it is applied Example penalty
Repetition Penalty Prevents specific tokens or phrases from being repeated. Applies a penalty to the log probabilities (logits) of all previously generated tokens. For example, if a token has already been generated, its log probability is divided by the penalty value (multiplicative) or has the penalty value subtracted from it (subtractive).
Frequency Penalty Applies a penalty based on the frequency of generated tokens, suppressing tokens that appear often. Applies a penalty based on the number of times each token has been generated. Each time a token appears, its probability of appearing is reduced. Each time a token appears, its log probability is cumulatively divided by the penalty value (multiplicative) or has the penalty value cumulatively subtracted from it (subtractive).
Presence Penalty Prevents tokens that have already been generated from appearing again. Applies a penalty based on whether a token has been generated at least once. Tokens that have been generated once are penalized when they reappear. The log probability of a token that has been generated once is divided by the penalty value (multiplicative) or has the penalty value subtracted from it (subtractive).

Read more