什么是Top-p采样与Top-k采样？大模型推理时如何同时设置？解析Transformers库源代码

什么是Top-p采样与Top-k采样？大模型推理时如何同时设置？

在自然语言生成（NLG）和大规模语言模型推理中，Top-k采样 和 Top-p采样 是两种常用的解码策略，用于控制生成文本的多样性和质量。它们在生成过程中对词汇选择有不同的约束方式，可以通过调整生成过程的随机性来得到更符合需求的输出。接下来，我们将分别介绍这两种采样方法，并讨论它们如何在大模型推理时同时设置。

一、Top-k采样

1.1 什么是Top-k采样？

Top-k采样 是一种基于概率的词汇选择策略。在每次生成一个词时，模型会根据输出的概率分布选出前k个最可能的词，并从这k个词中随机选择一个作为输出。

具体来说，Top-k采样的步骤如下：

1. 模型生成当前时刻的每个词的概率分布。
2. 根据概率分布选出概率最高的前k个词。
3. 在这k个词中，根据它们的概率分布进行随机选择。

Top-k采样的优点：

• 它通过限制候选词汇的数量来避免低概率词汇对结果的影响，从而控制输出的多样性。
• 能够减少生成文本中的不合适词汇，避免模型输出语法不正确或者语义不清的内容。

Top-k采样的缺点：

• 它依然可能存在生成内容缺乏连贯性的风险，尤其是k值较小时，模型选择的候选词较少，可能使生成结果过于单一。

1.2 示例

假设我们有一个词汇的概率分布：

• word1: 0.4
• word2: 0.3
• word3: 0.2
• word4: 0.1

如果我们设置 k=2，那么只有 word1 和 word2 会被考虑。接下来，从这两个词中随机选择一个作为下一个生成的词。

二、Top-p采样（又叫核采样）

2.1 什么是Top-p采样？

Top-p采样，又叫核采样（Nucleus Sampling），是一种基于累积概率的采样策略。与Top-k采样不同，Top-p采样不是固定选择前k个最可能的词，而是选择一个累计概率为p的词集。

具体来说，Top-p采样的步骤如下：

1. 对所有可能的词按照概率从高到低排序。
2. 累加排序后的概率，直到累计概率大于或等于p为止。
3. 从这些词中随机选择一个词作为输出。

Top-p采样的优点：

• 它自适应地选择词汇数量，避免了Top-k采样中人为限制候选词数量的问题。对于低概率的词汇，Top-p采样会自动排除它们。
• 通过动态调整候选词的数量，能够在控制生成多样性的同时保证输出的合理性和连贯性。

Top-p采样的缺点：

• 由于选择的候选词集是动态的，这使得Top-p采样在不同的生成步骤中可能会有不同数量的候选词，因此它的控制方式相对更复杂。

2.2 示例

假设我们有一个词汇的概率分布：

• word1: 0.4
• word2: 0.3
• word3: 0.2
• word4: 0.1

如果我们设置 p=0.7，那么从概率分布中选出的词将包括 word1 和 word2，因为它们的累计概率（0.4 + 0.3 = 0.7）达到了设定的p值。因此，从这两个词中随机选择一个作为下一个生成的词。

三、Top-k与Top-p同时设置时的采样方法

在实际应用中，Top-k采样 和 Top-p采样 可以一起使用，结合两者的优点来优化生成效果。当Top-k和Top-p同时设置时，模型的词汇选择会受到两者的限制。具体的采样步骤如下：

1. 应用Top-k约束：首先，限制模型候选词汇的数量为k个，即只考虑概率最高的前k个词。
2. 应用Top-p约束：然后，在这k个词中，进一步选择使得累计概率之和大于或等于p的词集。
3. 从这个经过双重筛选的候选词集合中进行随机选择。

这种组合方式既能控制生成文本的多样性（通过Top-p）和合理性（通过Top-k），同时也能避免单独使用Top-k时可能产生的过于单一的输出，或者单独使用Top-p时生成的低概率内容。

3.1 示例

假设我们有一个词汇的概率分布：

• word1: 0.4
• word2: 0.3
• word3: 0.2
• word4: 0.1

现在我们设置 k=2 和 p=0.7：

1. Top-k筛选：根据Top-k策略，选出概率最高的两个词 word1 和 word2。
2. Top-p筛选：根据Top-p策略，累计概率是0.7，因此 word1 和 word2 被选中，因为它们的累计概率为0.7。

最终，我们从 word1 和 word2 中随机选择一个作为下一个生成的词。

3.2 为什么要同时使用Top-k和Top-p？

• 避免过于狭隘的选择范围：Top-k可能会限制候选词的数量，导致生成文本的多样性降低。通过引入Top-p，可以保证在保证概率约束的同时，仍然有一定的多样性。
• 提高生成结果的合理性：Top-p能够自动控制候选词的数量，而Top-k能够避免低概率词汇的干扰。两者结合使用，可以避免生成不符合上下文或语法不合的词。

四、总结

• Top-k采样通过固定选择前k个最有可能的词，控制候选词的数量，通常适用于需要严格控制生成内容的场景。
• Top-p采样则基于累计概率动态选择候选词集，能够自适应地调整候选词的数量，适用于生成多样性和合理性之间的平衡。
• Top-k和Top-p同时设置时，先通过Top-k筛选出k个候选词，再通过Top-p限制累积概率，保证了生成内容的多样性和合理性。

在实际应用中，选择Top-k、Top-p或两者结合的采样策略，取决于生成任务的具体需求以及所需的生成效果。在大模型推理中，通过调整这两种策略，可以显著提高生成结果的质量和多样性。

Transformers库的代码实现

先用TopKLogitsWarper进行top k采样，然后用TopPLogitsWarper进行top p采样：

https://github.com/huggingface/transformers/blob/main/src/transformers/generation/utils.py

if generation_config.top_k is not None and generation_config.top_k != 0:
    processors.append(
        TopKLogitsWarper(top_k=generation_config.top_k, min_tokens_to_keep=min_tokens_to_keep)
    )
if generation_config.top_p is not None and generation_config.top_p < 1.0:
    processors.append(
        TopPLogitsWarper(top_p=generation_config.top_p, min_tokens_to_keep=min_tokens_to_keep)
    )

其中两个类的具体实现如下：

https://github.com/huggingface/transformers/blob/main/src/transformers/generation/logits_process.py

首先是top k

class TopKLogitsWarper(LogitsProcessor):
    r"""
    [`LogitsProcessor`] that performs top-k, i.e. restricting to the k highest probability elements. Often used
    together with [`TemperatureLogitsWarper`] and [`TopPLogitsWarper`].

    Args:
        top_k (`int`):
            The number of highest probability vocabulary tokens to keep for top-k-filtering.
        filter_value (`float`, *optional*, defaults to -inf):
            All filtered values will be set to this float value.
        min_tokens_to_keep (`int`, *optional*, defaults to 1):
            Minimum number of tokens that cannot be filtered.

    Examples:

    ```python
    >>> from transformers import AutoTokenizer, AutoModelForCausalLM, set_seed

    >>> set_seed(1)
    >>> model = AutoModelForCausalLM.from_pretrained("distilbert/distilgpt2")
    >>> tokenizer = AutoTokenizer.from_pretrained("distilbert/distilgpt2")

    >>> inputs = tokenizer("A sequence: A, B, C, D", return_tensors="pt")

    >>> # With sampling, the output is unexpected -- sometimes too unexpected.
    >>> outputs = model.generate(**inputs, do_sample=True)
    >>> print(tokenizer.batch_decode(outputs, skip_special_tokens=True)[0])
    A sequence: A, B, C, D, E — S — O, P — R

    >>> # With `top_k` sampling, the output gets restricted the k most likely tokens.
    >>> # Pro tip: In practice, LLMs use `top_k` in the 5-50 range.
    >>> outputs = model.generate(**inputs, do_sample=True, top_k=2)
    >>> print(tokenizer.batch_decode(outputs, skip_special_tokens=True)[0])
    A sequence: A, B, C, D, E, F, G, H, I
    ```
    """

    def __init__(self, top_k: int, filter_value: float = -float("Inf"), min_tokens_to_keep: int = 1):
        if not isinstance(top_k, int) or top_k <= 0:
            raise ValueError(f"`top_k` has to be a strictly positive integer, but is {top_k}")

        self.top_k = max(top_k, min_tokens_to_keep)
        self.filter_value = filter_value

    @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING)
    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
        top_k = min(self.top_k, scores.size(-1))  # Safety check
        # Remove all tokens with a probability less than the last token of the top-k
        indices_to_remove = scores < torch.topk(scores, top_k)[0][..., -1, None]
        scores_processed = scores.masked_fill(indices_to_remove, self.filter_value)
        return scores_processed

代码解释：TopKLogitsWarper 类

TopKLogitsWarper 是一个继承自 LogitsProcessor 的类，用于对生成过程中预测的词汇概率进行 Top-k筛选。Top-k采样的基本思想是：每次生成时，仅保留概率最高的 k 个词，其他词的概率被设为一个非常小的值（通常是负无穷 -inf）。这样可以保证在生成过程中只有概率最高的词有可能被采样，避免生成低概率的无意义词汇。

下面是 TopKLogitsWarper 类的详细解释：

1. 构造函数：__init__(self, top_k, filter_value, min_tokens_to_keep)

def __init__(self, top_k: int, filter_value: float = -float("Inf"), min_tokens_to_keep: int = 1):
    if not isinstance(top_k, int) or top_k <= 0:
        raise ValueError(f"`top_k` has to be a strictly positive integer, but is {top_k}")

    self.top_k = max(top_k, min_tokens_to_keep)
    self.filter_value = filter_value

• top_k：是一个正整数，表示在每次生成时要保留的最大词汇数量。这个参数决定了每次生成时考虑哪些最有可能的词。
• filter_value：是一个浮动值，用于填充被过滤掉的词的概率，通常是负无穷 -inf，意味着这些词不再被选中。
• min_tokens_to_keep：确保至少保留一定数量的词，防止在某些情况下词汇数量太少。

该构造函数首先检查 top_k 是否为正整数。如果是，则将其设置为类的属性；否则，抛出一个 ValueError。同时，top_k 会和 min_tokens_to_keep 做比较，确保 top_k 至少为 min_tokens_to_keep。

2. __call__ 方法：def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor

def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
    top_k = min(self.top_k, scores.size(-1))  # Safety check
    # Remove all tokens with a probability less than the last token of the top-k
    indices_to_remove = scores < torch.topk(scores, top_k)[0][..., -1, None]
    scores_processed = scores.masked_fill(indices_to_remove, self.filter_value)
    return scores_processed

• input_ids：是输入的token ID序列，通常是前文生成的tokens。
• scores：是模型输出的词汇预测分数（logits），其形状通常为 (batch_size, vocab_size)，表示对于每个token预测所有词汇的logits。
• top_k：根据 top_k 的值，确定要保留多少个词的分数。这里使用 min(self.top_k, scores.size(-1)) 来确保不会超出词汇表的大小。
• torch.topk(scores, top_k)：得到词汇中前 top_k 个最大值和它们的索引。scores < torch.topk(scores, top_k)[0][..., -1, None] 语句表示筛选出所有小于前 top_k 个最大值的词（即概率较小的词）。
• scores.masked_fill(indices_to_remove, self.filter_value)：将所有低于前 top_k 的词的分数替换为 filter_value，通常是 -inf，这些词在后续生成中将不会被选中。

最终，这个方法返回的是处理后的分数，其中低于Top-k的词的分数被设为 -inf，使得它们不会被选中。

示例代码解释：outputs = model.generate(**inputs, do_sample=True, top_k=2)

>>> outputs = model.generate(**inputs, do_sample=True, top_k=2)
>>> print(tokenizer.batch_decode(outputs, skip_special_tokens=True)[0])
A sequence: A, B, C, D, E, F, G, H, I

在这段代码中：

• top_k=2 表示在每一步生成时，仅考虑概率最高的 2 个词汇。换句话说，生成时每次会有两个候选词可供选择，其他词的概率被设为 -inf，不能被选中。
• do_sample=True 表示启用采样机制。没有启用采样时，生成模型会选择概率最大的词（即贪心算法）。

为什么输出多了 E, F, G, H, I？

top_k=2 并不是要求输出的每个词只能从前两个候选词中选，而是每次生成时，从所有词汇中，模型只会计算前2个最大概率的词的可能性。由于在 do_sample=True 的情况下，生成过程是基于概率的，并且每次采样都可能选择概率较高的词，即便它们是 top_k 中的第二个候选词。

生成过程的解释：

1. 首先输入 “A sequence: A, B, C, D”。
2. 在 top_k=2 的情况下，模型会根据当前生成的文本，计算所有词汇的概率，选择前 2 个概率最高的词（比如 E 和 F，可能是当前最合适的词）。
3. 在每一步生成中，选择其中一个词，比如 E。
4. 随后，模型继续计算概率，再选择前2个词，可能包括 F，并选择其中一个词。
5. 由于 do_sample=True，生成过程中存在随机性，模型有一定的机会选择更高概率的词，导致最终生成文本可能比我们预期的更长。

因此，虽然我们设置了 top_k=2，但因为采样过程的随机性和生成长度的逐步递增，最终输出的文本可能是 “A sequence: A, B, C, D, E, F, G, H, I”，并不是只生成 E 这一个词，而是继续扩展到了 I。

总结

• Top-k采样 通过限制每次生成时的候选词数目，帮助控制生成的多样性和质量。
• Top-k=2 只是每次生成时，从概率最高的2个词中选择一个。因为采样机制的随机性和生成过程的逐步扩展，最终输出可能生成多个词，甚至超过设置的 top_k。
• 结合 do_sample=True，生成过程是随机的，所以生成结果可能包含更多的词汇，而不仅仅是前两个候选词中的某个。

希望这些解释能够帮助你更好地理解 TopKLogitsWarper 以及为什么 top_k=2 会输出超过两个词的原因。

详细解析scores < torch.topk(scores, top_k)[0][…, -1, None]

这句代码 indices_to_remove = scores < torch.topk(scores, top_k)[0][..., -1, None] 是 TopKLogitsWarper 中非常关键的一行，用于实现 Top-k 筛选。它的目的是从模型输出的分数中移除那些不在 Top-k 范围内的词（即低概率的词），确保只有概率最高的前 k 个词能够参与生成。

详细解析

我们分解这行代码来理解其细节：

indices_to_remove = scores < torch.topk(scores, top_k)[0][..., -1, None]

1. scores：

   • scores 是一个 浮动类型的张量，通常是形状为 (batch_size, vocab_size)，表示模型在每个生成步骤中预测各个词汇的 “logits”（即每个词的概率分数）。
   • 这些分数还没有经过 softmax 处理，因此它们的数值代表的是相对概率大小，但不是实际的概率值。

2. torch.topk(scores, top_k)：

   • torch.topk(scores, top_k) 计算的是 scores 中的前 top_k 个最大值，并返回这 top_k 个最大值以及它们的索引。
   • 这会返回一个包含两个元素的元组，第一个元素是最大值的张量，第二个元素是这些最大值对应的索引。
   • torch.topk(scores, top_k) 的输出形式是 (values, indices)，其中 values 是形状为 (batch_size, top_k) 的张量，包含每个批次的前 top_k 个最大分数；indices 是形状为 (batch_size, top_k) 的张量，包含这些最大分数的词汇索引。

3. torch.topk(scores, top_k)[0]：

   • torch.topk(scores, top_k)[0] 选取的是 最大值 部分，即前 top_k 个最大分数，形状为 (batch_size, top_k)。每一行表示该批次中前 top_k 个最大概率的词的分数。
   • 例如，如果 top_k=3，且某一批次的 scores 为 [0.1, 0.5, 0.3, 0.2]，top_k=3 之后，前 3 个最大值会是 [0.5, 0.3, 0.2]，它们对应的索引是 [1, 2, 3]。

4. [..., -1, None]：

   • [..., -1, None] 通过切片操作选择了 每个批次 的第 top_k 个最大值，即每一行的 最后一个最大值（最小的最大值）。
   • 这是因为在 torch.topk 中，top_k 的最大值排在第一位，第二大排在第二位，依此类推。所以，[..., -1, None] 取的是每一行的 第 top_k 个最大值（即最小的那个大值）。
   • -1 代表每行的最后一个元素，None 用来保持张量的形状一致，确保结果仍然是一个二维张量，方便后续广播操作。

5. scores < torch.topk(scores, top_k)[0][..., -1, None]：

   • 这部分代码实际上是在检查每个词的分数是否小于该批次的第 top_k 个最大分数。如果是，则这个词会被 移除（即给它的分数设为 filter_value，通常为 -inf）。

   • 具体来说，

     scores < torch.topk(scores, top_k)[0][..., -1, None]
     

     生成一个布尔类型的张量

     indices_to_remove
     

     ，其中值为

     True
     

     的位置表示这些词的分数小于第

     top_k

     个最大分数。

     • 如果某个词的分数小于该批次的第 top_k 个最大分数，意味着它不在 Top-k 范围内，所以我们希望把它的分数设为负无穷，从而在后续的采样中排除它。

举个例子

假设有以下的 scores 张量，表示某次生成的每个词的概率分数（top_k=3）：

scores = torch.tensor([[0.1, 0.5, 0.3, 0.2]])  # shape: (1, 4)

• 这个张量表示有 1 个样本（batch size = 1），词汇表有 4 个词，分数分别为 [0.1, 0.5, 0.3, 0.2]。

执行 torch.topk(scores, top_k=3) 后，得到：

top_k_values = torch.topk(scores, 3)  # 获取前 3 个最大值
print(top_k_values)
# output: torch.return_types.topk(
#     values=tensor([[0.5, 0.3, 0.2]]), 
#     indices=tensor([[1, 2, 3]]))

• top_k_values[0] 即 [0.5, 0.3, 0.2]，这就是前 3 个最大值。

然后，使用切片 [..., -1, None] 选取第 top_k 个最大值（即最小的那个最大值）：

print(top_k_values[0][..., -1, None])  # 最小的最大值
# output: tensor([[0.2]])

最后，执行 scores < top_k_values[0][..., -1, None] 进行比较，得到：

indices_to_remove = scores < top_k_values[0][..., -1, None]
print(indices_to_remove)
# output: tensor([[True, False, False, False]])

• 这里的比较结果表明，第一个词（0.1）的分数小于 0.2（即第 3 大的分数），因此被标记为 True，表示需要移除；其他的词（0.5, 0.3, 0.2）则不需要移除。

总结

• 最大值 在这里是指 每个批次的前 top_k 个最大概率分数。
• 通过 torch.topk(scores, top_k)[0][..., -1, None] 获取 第 top_k 个最大值（即前 top_k 个词中概率最小的那个值）。
• 然后通过比较 scores < top_k_value 来确定哪些词的概率小于这个值，从而将这些词排除在后续生成过程中。

这种做法的核心是：在生成时，只保留概率最高的前 top_k 个词，过滤掉其他低概率的词。

代码解释：TopPLogitsWarper 类

下面是TopP的源代码

class TopPLogitsWarper(LogitsProcessor):
    """
    [`LogitsProcessor`] that performs top-p, i.e. restricting to top tokens summing to prob_cut_off <= prob_cut_off.
    Often used together with [`TemperatureLogitsWarper`] and [`TopKLogitsWarper`].

    Args:
        top_p (`float`):
            If set to < 1, only the smallest set of most probable tokens with probabilities that add up to `top_p` or
            higher are kept for generation.
        filter_value (`float`, *optional*, defaults to -inf):
            All filtered values will be set to this float value.
        min_tokens_to_keep (`int`, *optional*, defaults to 1):
            Minimum number of tokens that cannot be filtered.

    Examples:

    ```python
    >>> from transformers import AutoTokenizer, AutoModelForCausalLM, set_seed

    >>> set_seed(1)
    >>> model = AutoModelForCausalLM.from_pretrained("distilbert/distilgpt2")
    >>> tokenizer = AutoTokenizer.from_pretrained("distilbert/distilgpt2")

    >>> inputs = tokenizer("A sequence: 1, 2", return_tensors="pt")

    >>> # With sampling, the output is unexpected -- sometimes too unexpected.
    >>> outputs = model.generate(**inputs, do_sample=True)
    >>> print(tokenizer.batch_decode(outputs, skip_special_tokens=True)[0])
    A sequence: 1, 2, 3 | < 4 (left-hand pointer) ;
    <BLANKLINE>
    <BLANKLINE>

    >>> # With `top_p` sampling, the output gets restricted to high-probability tokens.
    >>> # Pro tip: In practice, LLMs use `top_p` in the 0.9-0.95 range.
    >>> outputs = model.generate(**inputs, do_sample=True, top_p=0.1)
    >>> print(tokenizer.batch_decode(outputs, skip_special_tokens=True)[0])
    A sequence: 1, 2, 3, 4, 5, 6, 7, 8, 9
    ```
    """

    def __init__(self, top_p: float, filter_value: float = -float("Inf"), min_tokens_to_keep: int = 1):
        top_p = float(top_p)
        if top_p < 0 or top_p > 1.0:
            raise ValueError(f"`top_p` has to be a float > 0 and < 1, but is {top_p}")
        if not isinstance(min_tokens_to_keep, int) or (min_tokens_to_keep < 1):
            raise ValueError(f"`min_tokens_to_keep` has to be a positive integer, but is {min_tokens_to_keep}")

        self.top_p = top_p
        self.filter_value = filter_value
        self.min_tokens_to_keep = min_tokens_to_keep

    @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING)
    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
        sorted_logits, sorted_indices = torch.sort(scores, descending=True)
        cumulative_probs = sorted_logits.softmax(dim=-1).cumsum(dim=-1)

        # Remove tokens with cumulative top_p above the threshold (token with 0 are kept)
        sorted_indices_to_remove = cumulative_probs <= (1 - self.top_p)
        # Keep at least min_tokens_to_keep
        sorted_indices_to_remove[..., -self.min_tokens_to_keep :] = 0

        # scatter sorted tensors to original indexing
        indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
        scores_processed = scores.masked_fill(indices_to_remove, self.filter_value)
        return scores_processed

这段代码是 TopPLogitsWarper 类的实现，用于实现 Top-p 筛选，这是一种常见的用于生成模型的采样方法。其目的是限制生成过程中所选择的候选词集的累积概率，使得总和不超过设定的 top_p 阈值，确保生成的结果更具多样性和质量。

TopPLogitsWarper 会将概率分布按从高到低排序，然后从最高概率的词开始累加，直到这些词的概率和达到或超过指定的 top_p 阈值。然后，只有这些高概率的词才有机会被选中，而其他的低概率词则会被屏蔽（即其概率被设为 filter_value）。

代码解释：TopPLogitsWarper 类

class TopPLogitsWarper(LogitsProcessor):
    def __init__(self, top_p: float, filter_value: float = -float("Inf"), min_tokens_to_keep: int = 1):
        top_p = float(top_p)
        if top_p < 0 or top_p > 1.0:
            raise ValueError(f"`top_p` has to be a float > 0 and < 1, but is {top_p}")
        if not isinstance(min_tokens_to_keep, int) or (min_tokens_to_keep < 1):
            raise ValueError(f"`min_tokens_to_keep` has to be a positive integer, but is {min_tokens_to_keep}")

        self.top_p = top_p
        self.filter_value = filter_value
        self.min_tokens_to_keep = min_tokens_to_keep

1. __init__
   

   方法:

   • top_p：这是限制词汇累积概率和的阈值，介于 0 到 1 之间。例如，top_p=0.9 会保留概率累积和达到 90% 的最有可能的词，其他词都会被滤除。
   • filter_value：这个参数用于设定过滤掉的词的分数值，通常设为负无穷（-inf），以便完全排除这些词。
   • min_tokens_to_keep：确保至少保留一定数量的高概率词，避免因为阈值过低而只剩下非常少的候选词。

    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
        sorted_logits, sorted_indices = torch.sort(scores, descending=False)
        cumulative_probs = sorted_logits.softmax(dim=-1).cumsum(dim=-1)

1. __call__
   

   方法:

   • scores：模型生成的词汇概率分数。
   • sorted_logits, sorted_indices：torch.sort(scores, descending=False) 会对分数进行升序排序，sorted_logits 是排序后的分数，sorted_indices 是排序后的索引（即哪个分数对应哪个词汇）。
   • cumulative_probs：对排序后的 logits 应用 softmax 函数后，得到每个词的概率。然后使用 cumsum(dim=-1) 计算累积概率。累积概率可以帮助我们确定保留哪些词汇。例如，累积概率为 [0.1, 0.4, 0.7, 1.0]，表示前两个词的累积概率已经达到 0.4，前 3 个达到 0.7。

        sorted_indices_to_remove = cumulative_probs <= (1 - self.top_p)
        sorted_indices_to_remove[..., -self.min_tokens_to_keep :] = 0

1. 确定需要移除的词:
   • sorted_indices_to_remove 是一个布尔数组，标记了哪些词的累积概率和超出了 top_p 的阈值。具体来说，cumulative_probs <= (1 - self.top_p) 会将累积概率和超过 top_p 阈值的词标记为 True（即需要移除的词）。
   • 但是，我们需要确保至少保留 min_tokens_to_keep 个词，因此通过 sorted_indices_to_remove[..., -self.min_tokens_to_keep :] = 0 将最小数量的词的移除标志设置为 False，即这些词不会被移除。

        indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
        scores_processed = scores.masked_fill(indices_to_remove, self.filter_value)
        return scores_processed

1. 恢复原始索引的顺序:
   • scatter(1, sorted_indices, sorted_indices_to_remove)：scatter 操作将布尔标记（哪些词需要移除）映射回原始的词汇顺序。由于我们之前对 scores 进行了排序，原来的索引已经丢失，因此我们需要使用 sorted_indices 来重新恢复到原始的顺序。具体过程可以参考笔者的另一篇博客：PyTorch 中scatter函数：以Transformers库Top P采样为例
   • scores.masked_fill(indices_to_remove, self.filter_value)：最后，通过 masked_fill 函数将需要移除的词的分数设置为 filter_value（通常是 -inf）。这样，在生成时，模型会忽略这些词，只考虑剩下的高概率词。

代码执行过程

假设我们有以下的 scores 张量，表示每个词的概率分数：

scores = torch.tensor([[0.1, 0.5, 0.3, 0.2]])  # shape: (1, 4)

并设置 top_p = 0.7，min_tokens_to_keep = 1。

1. 排序 scores 和计算累积概率：
   • sorted_logits = [0.1, 0.2, 0.3, 0.5]
   • cumulative_probs = [0.1, 0.3, 0.6, 1.0] （假设 softmax 已经计算过）
2. 根据 top_p 阈值 0.7 计算移除的词：
   • cumulative_probs <= (1 - top_p) 等价于 cumulative_probs <= 0.3，标记需要移除的词。[True, True, False, False]。
3. 保留至少 min_tokens_to_keep = 1 个词：
   • sorted_indices_to_remove[..., -self.min_tokens_to_keep :] = 0，即最后一个词不会被移除，因此更新为 [True, True, False, False]。
4. 恢复原索引顺序：
   • 使用 scatter 恢复到原始顺序，得到最终的 indices_to_remove。
5. 设置过滤值：
   • 使用 masked_fill 将被移除的词的分数设置为 -inf，这样模型在生成时不会选择这些词。

总结

TopPLogitsWarper 类通过计算累积概率，确保在生成过程中仅考虑概率和达到 top_p 阈值的词。这种方法使得生成的文本更具多样性，而不会过度集中在少数几个高概率的词汇上，同时避免了只选择非常小的概率集的情况。

TOP P两种实现方式

sorted_logits, sorted_indices = torch.sort(logits, descending=True)
cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)

# Remove tokens with cumulative probability above the threshold (token with 0 are kept)
sorted_indices_to_remove = cumulative_probs > top_p

这样的实现方式和

sorted_logits, sorted_indices = torch.sort(scores, descending=False)
cumulative_probs = sorted_logits.softmax(dim=-1).cumsum(dim=-1)

# Remove tokens with cumulative top_p above the threshold (token with 0 are kept)
sorted_indices_to_remove = cumulative_probs <= (1 - self.top_p)

这种方式等价吗？

这两种实现方式在功能上并不完全等价，它们的主要区别在于排序顺序和累积概率的阈值条件。

第一个实现方式：

sorted_logits, sorted_indices = torch.sort(logits, descending=True)
cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
sorted_indices_to_remove = cumulative_probs > top_p

分析：

1. 排序方式：torch.sort(logits, descending=True) 对 logits 进行降序排序。这样，sorted_logits 会按从大到小的顺序排列。
2. 累积概率计算：计算累积概率时，softmax 会计算每个词的概率，cumsum 会将概率按顺序累加。
3. 筛选词汇：使用 cumulative_probs > top_p 来筛选词汇。这意味着累积概率超过 top_p 的词会被移除，保留的词是累积概率和小于等于 top_p 的词。

第二个实现方式：

sorted_logits, sorted_indices = torch.sort(scores, descending=False)
cumulative_probs = sorted_logits.softmax(dim=-1).cumsum(dim=-1)
sorted_indices_to_remove = cumulative_probs <= (1 - self.top_p)

分析：

1. 排序方式：torch.sort(scores, descending=False) 对 scores 进行升序排序。这样，sorted_logits 会按从小到大的顺序排列。
2. 累积概率计算：与第一个实现一样，先应用 softmax，然后计算累积概率。
3. 筛选词汇：使用 cumulative_probs <= (1 - top_p) 来筛选词汇。这意味着累积概率和小于等于 1 - top_p 的词会被移除。

区别总结：

1. 排序顺序的不同：
   • 第一个实现使用 descending=True 对 logits 进行降序排序，而第二个实现使用 descending=False 对 scores 进行升序排序。
   • 这会导致它们处理的候选词的顺序不同，从而影响累积概率的计算方式和最终筛选出来的词。
2. 筛选条件的不同：
   • 第一个实现使用 cumulative_probs > top_p，表示累积概率超过 top_p 的词会被移除。
   • 第二个实现使用 cumulative_probs <= (1 - top_p)，表示累积概率和小于等于 1 - top_p 的词会被移除。

两者的等价性分析：

这两种方式的筛选机制是有区别的。第一种方式会移除那些累积概率超过 top_p 的词，而第二种方式会移除那些累积概率和小于等于 1 - top_p 的词。虽然它们在一定条件下可能产生相同的结果（比如 top_p = 0.9），但它们的内部实现和累积概率阈值的含义是不同的。

等价条件：

• 如果我们对 top_p = 0.9 设置，在第一个实现中，筛选的是累积概率大于 0.9 的词，而在第二个实现中，筛选的是累积概率和小于等于 0.1 的词。两者在逻辑上是对等的，因为 1 - top_p = 0.1。
• 但这并不意味着它们始终是等价的。它们的排序和筛选方式是不同的，因此在某些情况下，它们的输出结果会有所不同。

总结：

虽然这两种方法在 top_p 值相同的情况下可能产生相似的效果，但由于排序方式和筛选条件的不同，它们不是完全等价的。如果需要确保两者行为完全一致，需要保证排序方式和筛选条件的一致性。

Reference

a top p implementation

Top p source paper

a medium article about top p

大模型推理的时候 top k 和 top p 同时设置的时候怎么采样？

后记

2024年12月27日20点58分于上海，在GPT4o mini大模型辅助下完成。