对大模型输出的 logits 进行处理，从而控制文本的生成

对大模型输出的 logits 进行处理，从而控制文本的生成

flyfish

在文本生成任务中，模型输出的 logits 代表了每个词被选为下一个生成词的未归一化概率得分。通过对 logits 进行处理，可以精确地控制文本的生成

基本原理

在每一步生成过程中，模型会输出一个 logits 向量，其长度等于词汇表的大小，每个元素对应词汇表中一个词的得分。通常，会对 logits 应用 softmax 函数将其转换为概率分布，然后根据这个概率分布来选择下一个生成的词。而 logits_processor 就是在应用 softmax 函数之前，对 logits 进行修改，从而改变最终的概率分布和词的选择。

具体控制方式

1. 避免重复

• 重复惩罚（RepetitionPenaltyLogitsProcessor）
  • 机制：对于已经在生成文本中出现过的词，降低其 logits 的值。具体来说，会将这些词的 logits 除以一个大于 1 的惩罚系数，使得它们在后续生成中被选中的概率降低。
  • 示例：假设生成的文本中已经出现了“苹果”这个词，当模型再次预测下一个词时，“苹果”对应的 logits 会被惩罚，从而减少再次生成“苹果”的可能性，避免文本中出现过多重复内容。
• 不重复 n - gram（NoRepeatNGramLogitsProcessor）
  • 机制：检查生成的文本中是否已经存在某个 n - gram（连续的 n 个词），如果存在，则将可能导致该 n - gram 重复出现的词的 logits 设为负无穷。这样，在后续的概率计算中，这些词的概率会变为 0，不会被选中。
  • 示例：如果 n = 2，当前生成的文本是“我 喜欢”，那么在选择下一个词时，会避免选择那些会导致“我 喜欢”这个 2 - gram 重复出现的词，如“我”或“喜欢”，从而提高文本的多样性。

2. 控制生成长度

• 最小长度限制（MinLengthLogitsProcessor）
  • 机制：在生成的文本长度未达到指定的最小长度之前，将结束标记（EOS）的 logits 设为负无穷。这样，在 softmax 处理后，结束标记的概率会变为 0，模型不会选择结束生成，确保文本达到一定的长度。
  • 示例：如果设置最小长度为 10 个词，在生成的词数小于 10 时，结束标记的 logits 始终为负无穷，模型会继续生成，直到达到最小长度要求。
• 最小新标记数（MinNewTokensLengthLogitsProcessor）
  • 机制：类似于最小长度限制，不过是针对新生成的标记数量。在新生成的标记数未达到指定数量之前，降低结束标记的 logits，保证生成足够数量的新内容。

3. 采样策略调整

• 温度调整（TemperatureLogitsWarper）
  • 机制：将 logits 除以一个温度参数 temperature。温度越高，logits 之间的差异会被缩小，经过 softmax 处理后，概率分布会更加均匀，采样会更随机；温度越低，logits 之间的差异会被放大，概率分布会更集中，更倾向于选择概率最大的词。
  • 示例：当 temperature = 1 时，保持原始的 logits 分布；当 temperature > 1 时，模型可能会生成一些更具创意但可能不太准确的文本；当 temperature < 1 时，模型会更保守，生成的文本更符合常见的表达。
• Top - k 采样（TopKLogitsWarper）
  • 机制：只保留 logits 中概率最高的 k 个词，将其余词的 logits 设为负无穷。这样，在后续的采样中，只会从这 k 个词中选择下一个生成的词，限制了采样范围，提高了生成的稳定性。
  • 示例：如果 k = 5，模型会在每次生成时，只考虑概率最高的 5 个词，排除其他词的干扰。
• Top - p 采样（TopPLogitsWarper）
  • 机制：选择累积概率达到 p 的最小词集合，只保留这些词的 logits，其余词的 logits 设为负无穷。这种方法结合了概率和词的数量，既能控制采样范围，又能适应不同的概率分布。
  • 示例：如果 p = 0.9，模型会选择累积概率达到 0.9 的最小词集合，从这个集合中进行采样。

4. 约束生成内容

• 禁用词过滤（NoBadWordsLogitsProcessor）
  • 机制：将禁用词的 logits 设为负无穷，使得这些词在后续的概率计算中概率为 0，不会被选中，从而避免生成包含禁用词的文本。
  • 示例：如果禁用词列表中包含“脏话”，那么在生成过程中，“脏话”对应的 logits 会被设为负无穷，不会出现在生成的文本中。
• 前缀约束（PrefixConstrainedLogitsProcessor）
  • 机制：根据给定的前缀允许标记函数，限制生成的词必须符合特定的前缀约束。不符合约束的词的 logits 会被设为负无穷，从而保证生成的文本符合特定的前缀要求。
  • 示例：如果要求生成的文本必须以“今天”开头，那么在生成第一个词时，只有与“今天”相关的词的 logits 会被保留，其他词的 logits 会被设为负无穷。

配置参数

logits 说明

logits 是模型在进行分类或预测任务时，最后一层神经元的原始输出值，它是未经过归一化处理的数值。在文本生成场景中，logits 代表了模型预测词汇表中每个词作为下一个生成词的得分，这些得分反映了模型对每个词成为下一个词的相对可能性判断，但并非是概率值。

数学公式

1. 线性变换得到 logits

在许多深度学习模型中，logits 通常是通过对前一层的输出进行线性变换得到的。假设模型前一层的输出为向量 $\mathbf{h}$，权重矩阵为 $\mathbf{W}$，偏置向量为 $\mathbf{b}$，则 logits 向量 $\mathbf{z}$ 的计算公式如下：

$$\mathbf{z}=\mathbf{Wh}+\mathbf{b}$$

其中，$\mathbf{h}$ 是前一层输出的特征向量，维度通常为 $d_h$；$\mathbf{W}$ 是权重矩阵，维度为 $V\times d_h$，$V$ 是词汇表的大小；$\mathbf{b}$ 是偏置向量，维度为 $V$；$\mathbf{z}$ 是 logits 向量，维度为 $V$，每个元素 $z_i$ 对应词汇表中第 $i$ 个词的得分。

2. logits 转换为概率分布

为了将 logits 转换为概率分布，通常会使用 softmax 函数。softmax 函数可以将 logits 向量中的每个元素转换为一个在 $[0,1]$ 范围内的值，且所有元素之和为 1，符合概率分布的定义。softmax 函数的数学公式如下：

$$P(y_i)=\frac{e^{z_i}}{{\sum }_{j=1}^V{e}^{z_j}}$$

其中，$P(y_i)$ 是词汇表中第 $i$ 个词被选为下一个生成词的概率，$z_i$ 是 logits 向量中第 $i$ 个元素的值，$V$ 是词汇表的大小。

示例

假设词汇表大小 $V=3$，模型输出的 logits 向量为 $\mathbf{z}=[2,1,3]$，下面计算经过 softmax 函数处理后的概率分布：

首先，计算分母的值：

$$\sum _{j=1}^3{e}^{z_j}=e^2+e^1+e^3\approx 7.389+2.718+20.086=30.193$$

然后，分别计算每个词的概率：

$$P(y_1)=\frac{e^2}{30.193}\approx \frac{7.389}{30.193}\approx 0.245$$

$$P(y_2)=\frac{e^1}{30.193}\approx \frac{2.718}{30.193}\approx 0.090$$

$$P(y_3)=\frac{e^3}{30.193}\approx \frac{20.086}{30.193}\approx 0.665$$

可以看到，经过 softmax 函数处理后，得到了一个概率分布 $[0.245,0.090,0.665]$，表示词汇表中三个词被选为下一个生成词的概率。

在模型中的作用

在文本生成任务中，模型会根据 logits 转换后的概率分布来选择下一个生成的词。常见的选择方法有贪心搜索（选择概率最大的词）、采样搜索（根据概率分布随机采样）等。同时，logits_processor 会对 logits 进行调整，从而影响最终的概率分布和词的选择，以控制文本生成的行为和质量。

代码说明

logits_processor 是 _get_logits_processor 方法的一个参数，它是一个可选的 LogitsProcessorList 对象。这个方法会根据 GenerationConfig 中的各种配置参数，创建一系列不同的 LogitsProcessor 实例，并将它们添加到 processors 列表中。最后，如果传入了 logits_processor，还会将其与新创建的处理器列表进行合并。

def _get_logits_processor(
        self,
        generation_config: GenerationConfig,
        input_ids_seq_length: int,
        encoder_input_ids: torch.LongTensor,
        prefix_allowed_tokens_fn: Callable[[int, torch.Tensor], List[int]],
        logits_processor: Optional[LogitsProcessorList],
        device: str = None,
        model_kwargs: Optional[Dict[str, Any]] = None,
        negative_prompt_ids: Optional[torch.Tensor] = None,
        negative_prompt_attention_mask: Optional[torch.Tensor] = None,
    ) -> LogitsProcessorList:
        """
        此函数返回一个 `LogitsProcessorList` 对象，该对象包含所有用于修改语言模型头部得分的相关 `LogitsProcessor` 实例。
        这些处理器会对模型预测的 logits 进行调整，以控制文本生成的行为，例如避免重复、控制生成长度等。

        参数:
            generation_config (GenerationConfig): 生成配置对象，包含了文本生成过程中的各种配置参数。
            input_ids_seq_length (int): 输入 ID 序列的长度。
            encoder_input_ids (torch.LongTensor): 编码器的输入 ID。
            prefix_allowed_tokens_fn (Callable[[int, torch.Tensor], List[int]]): 一个可调用对象，用于指定允许的前缀标记。
            logits_processor (Optional[LogitsProcessorList]): 可选的 logits 处理器列表。
            device (str, optional): 设备名称，如 'cuda' 或 'cpu'。默认为 None。
            model_kwargs (Optional[Dict[str, Any]], optional): 模型的其他关键字参数。默认为 None。
            negative_prompt_ids (Optional[torch.Tensor], optional): 负提示的 ID。默认为 None。
            negative_prompt_attention_mask (Optional[torch.Tensor], optional): 负提示的注意力掩码。默认为 None。

        返回:
            LogitsProcessorList: 包含所有 logits 处理器的列表。
        """
        # 实例化一个空的处理器列表
        processors = LogitsProcessorList()

        # 如果配置了引导比例且不为 1，则添加无批量分类器自由引导 logits 处理器
        if generation_config.guidance_scale is not None and generation_config.guidance_scale != 1:
            processors.append(
                UnbatchedClassifierFreeGuidanceLogitsProcessor(
                    generation_config.guidance_scale,
                    self,
                    unconditional_ids=negative_prompt_ids,
                    unconditional_attention_mask=negative_prompt_attention_mask,
                    use_cache=generation_config.use_cache,
                )
            )
        # 如果配置了序列偏差，则添加序列偏差 logits 处理器
        if generation_config.sequence_bias is not None:
            processors.append(SequenceBiasLogitsProcessor(sequence_bias=generation_config.sequence_bias))

        # 如果配置了多样性惩罚且大于 0，则添加汉明多样性 logits 处理器
        if generation_config.diversity_penalty is not None and generation_config.diversity_penalty > 0.0:
            processors.append(
                HammingDiversityLogitsProcessor(
                    diversity_penalty=generation_config.diversity_penalty,
                    num_beams=generation_config.num_beams,
                    num_beam_groups=generation_config.num_beam_groups,
                )
            )
        # 如果配置了编码器重复惩罚且不为 1，并且编码器输入 ID 的形状为二维，则添加编码器重复惩罚 logits 处理器
        if (
            generation_config.encoder_repetition_penalty is not None
            and generation_config.encoder_repetition_penalty != 1.0
        ):
            if len(encoder_input_ids.shape) == 2:
                processors.append(
                    EncoderRepetitionPenaltyLogitsProcessor(
                        penalty=generation_config.encoder_repetition_penalty,
                        encoder_input_ids=encoder_input_ids,
                    )
                )
            else:
                # 如果编码器输入 ID 形状不符合要求，发出警告
                warnings.warn(
                    "Passing `encoder_repetition_penalty` requires some form of `input_ids` to be passed to "
                    "`generate`, ignoring the argument.",
                    UserWarning,
                )
        # 如果配置了重复惩罚且不为 1，则添加重复惩罚 logits 处理器
        if generation_config.repetition_penalty is not None and generation_config.repetition_penalty != 1.0:
            processors.append(RepetitionPenaltyLogitsProcessor(penalty=generation_config.repetition_penalty))
        # 如果配置了不重复 n-gram 大小且大于 0，则添加不重复 n-gram logits 处理器
        if generation_config.no_repeat_ngram_size is not None and generation_config.no_repeat_ngram_size > 0:
            processors.append(NoRepeatNGramLogitsProcessor(generation_config.no_repeat_ngram_size))
        # 如果配置了编码器不重复 n-gram 大小且大于 0，并且编码器输入 ID 的形状为二维，则添加编码器不重复 n-gram logits 处理器
        if (
            generation_config.encoder_no_repeat_ngram_size is not None
            and generation_config.encoder_no_repeat_ngram_size > 0
        ):
            if len(encoder_input_ids.shape) == 2:
                processors.append(
                    EncoderNoRepeatNGramLogitsProcessor(
                        generation_config.encoder_no_repeat_ngram_size,
                        encoder_input_ids,
                    )
                )
            else:
                # 如果编码器输入 ID 形状不符合要求，发出警告
                warnings.warn(
                    "Passing `encoder_no_repeat_ngram_size` requires some form of `input_ids` to be passed to "
                    "`generate`, ignoring the argument.",
                    UserWarning,
                )
        # 如果配置了禁用词 ID，则添加禁用词 logits 处理器
        if generation_config.bad_words_ids is not None:
            processors.append(
                NoBadWordsLogitsProcessor(
                    generation_config.bad_words_ids,
                    generation_config._eos_token_tensor,
                )
            )
        # 如果配置了最小长度且大于 0，并且有结束标记张量，则添加最小长度 logits 处理器
        if (
            generation_config.min_length is not None
            and generation_config._eos_token_tensor is not None
            and generation_config.min_length > 0
        ):
            processors.append(
                MinLengthLogitsProcessor(
                    generation_config.min_length,
                    generation_config._eos_token_tensor,
                    device=device,
                )
            )
        # 如果配置了最小新标记数且大于 0，并且有结束标记张量，则添加最小新标记长度 logits 处理器
        if (
            generation_config.min_new_tokens is not None
            and generation_config._eos_token_tensor is not None
            and generation_config.min_new_tokens > 0
        ):
            processors.append(
                MinNewTokensLengthLogitsProcessor(
                    input_ids_seq_length,
                    generation_config.min_new_tokens,
                    generation_config._eos_token_tensor,
                    device=device,
                )
            )
        # 如果提供了前缀允许标记函数，则添加前缀约束 logits 处理器
        if prefix_allowed_tokens_fn is not None:
            processors.append(
                PrefixConstrainedLogitsProcessor(
                    prefix_allowed_tokens_fn,
                    generation_config.num_beams // generation_config.num_beam_groups,
                )
            )
        # 如果配置了强制起始标记 ID，则添加强制起始标记 logits 处理器
        if generation_config.forced_bos_token_id is not None:
            processors.append(
                ForcedBOSTokenLogitsProcessor(
                    generation_config.forced_bos_token_id,
                )
            )
        # 如果配置了强制结束标记 ID，则添加强制结束标记 logits 处理器
        if generation_config.forced_eos_token_id is not None:
            processors.append(
                ForcedEOSTokenLogitsProcessor(
                    generation_config.max_length,
                    generation_config.forced_eos_token_id,
                    device=device,
                )
            )
        # 如果配置了移除无效值，则添加移除无穷大和 NaN 值的 logits 处理器
        if generation_config.remove_invalid_values is True:
            processors.append(InfNanRemoveLogitsProcessor())
        # 如果配置了指数衰减长度惩罚，则添加指数衰减长度惩罚处理器
        if generation_config.exponential_decay_length_penalty is not None:
            processors.append(
                ExponentialDecayLengthPenalty(
                    generation_config.exponential_decay_length_penalty,
                    generation_config._eos_token_tensor,
                    input_ids_seq_length,
                )
            )
        # 如果配置了抑制标记，则添加抑制标记 logits 处理器
        if generation_config.suppress_tokens is not None:
            processors.append(
                SuppressTokensLogitsProcessor(
                    generation_config.suppress_tokens,
                    device=device,
                )
            )
        # 如果配置了起始抑制标记，则添加起始抑制标记 logits 处理器
        if generation_config.begin_suppress_tokens is not None:
            begin_index = input_ids_seq_length
            begin_index = (
                begin_index
                if (input_ids_seq_length > 1 or generation_config.forced_bos_token_id is None)
                else begin_index + 1
            )
            processors.append(
                SuppressTokensAtBeginLogitsProcessor(
                    generation_config.begin_suppress_tokens,
                    begin_index,
                    device=device,
                )
            )
        # 如果配置了强制解码器 ID，则抛出异常，提示使用 input_ids 或 decoder_input_ids 代替
        if generation_config.forced_decoder_ids is not None:
            # TODO (sanchit): move this exception to GenerationConfig.validate() when TF & FLAX are aligned with PT
            raise ValueError(
                "You have explicitly specified `forced_decoder_ids`. Please remove the `forced_decoder_ids` argument "
                "in favour of `input_ids` or `decoder_input_ids` respectively.",
            )

        # 合并自定义的 logits 处理器列表
        processors = self._merge_criteria_processor_list(processors, logits_processor)

        # 以下处理器之前被称为 `LogitsWarpers`，仅在采样策略下应用
        if generation_config.do_sample:
            # 在束搜索方法中，我们需要至少保留一个非结束标记来探索可能有更好得分的延续（即保留 len(list(generation_config._eos_token_tensor)) + 1）
            if generation_config.num_beams > 1:
                if isinstance(generation_config._eos_token_tensor, list):
                    min_tokens_to_keep = len(generation_config._eos_token_tensor) + 1
                elif isinstance(generation_config._eos_token_tensor, torch.Tensor):
                    min_tokens_to_keep = generation_config._eos_token_tensor.shape[0] + 1
                else:
                    min_tokens_to_keep = 2
            else:
                min_tokens_to_keep = 1

            # 以下思路主要借鉴自这个 PR: https://github.com/huggingface/transformers/pull/5420/files
            # 所有采样器可以在 `generation_utils_samplers.py` 中找到
            # 如果配置了温度且不为 1，则添加温度 logits 调整器
            if generation_config.temperature is not None and generation_config.temperature != 1.0:
                processors.append(TemperatureLogitsWarper(generation_config.temperature))
            # 如果配置了 top-k 采样且不为 0，则添加 top-k logits 调整器
            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)
                )
            # 如果配置了 top-p 采样且小于 1，则添加 top-p logits 调整器
            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)
                )
            # 如果配置了最小概率阈值，则添加最小概率 logits 调整器
            if generation_config.min_p is not None:
                # 在温度缩放后应用（见 https://github.com/ggerganov/llama.cpp/pull/3841#issuecomment-2073826084）
                processors.append(
                    MinPLogitsWarper(min_p=generation_config.min_p, min_tokens_to_keep=min_tokens_to_keep)
                )
            # 如果配置了典型概率采样且小于 1，则添加典型概率 logits 调整器
            if generation_config.typical_p is not None and generation_config.typical_p < 1.0:
                processors.append(
                    TypicalLogitsWarper(mass=generation_config.typical_p, min_tokens_to_keep=min_tokens_to_keep)
                )
            # 如果配置了 epsilon 截断且在 0 到 1 之间，则添加 epsilon logits 调整器
            if generation_config.epsilon_cutoff is not None and 0.0 < generation_config.epsilon_cutoff < 1.0:
                processors.append(
                    EpsilonLogitsWarper(
                        epsilon=generation_config.epsilon_cutoff, min_tokens_to_keep=min_tokens_to_keep
                    )
                )
            # 如果配置了 eta 截断且在 0 到 1 之间，则添加 eta logits 调整器
            if generation_config.eta_cutoff is not None and 0.0 < generation_config.eta_cutoff < 1.0:
                processors.append(
                    EtaLogitsWarper(
                        epsilon=generation_config.eta_cutoff, min_tokens_to_keep=min_tokens_to_keep, device=device
                    )
                )

        # 水印处理应该在所有 logits 处理完成后进行（见 #34630）
        if generation_config.watermarking_config is not None:
            processors.append(
                generation_config.watermarking_config.construct_processor(self.config.vocab_size, device)
            )

        # `LogitNormalization` 应该始终是最后一个 logit 处理器（如果存在）
        if generation_config.renormalize_logits is True:
            processors.append(LogitNormalization())
        return processors

transformers/src/transformers/generation/utils.py