如何在模型训练时避免计算 Padding Token 的 Loss

如何在模型训练时避免计算 Padding Token 的 Loss

在自然语言处理（NLP）任务中，模型输入的句子长度通常不同，因此需要通过 Padding 来对齐句子的长度，以便构建批量数据（Batch）。然而，Padding 的位置不包含有效的语义信息，不应该参与 Loss 的计算。为此，我们可以通过特定的方法屏蔽 Padding Token 的 Loss。本文将详细介绍实现原理和代码实现，帮助你更好地理解这一过程。

---

1. 什么是 Padding？为什么需要屏蔽它的 Loss？

1.1 什么是 Padding？

在 NLP 中，处理批量数据时，模型需要输入固定长度的序列。为了对齐序列长度，短于最大长度的句子需要在末尾补全填充标记（Padding Token）。例如：

输入句子1: ["I", "love", "coding"]
输入句子2: ["Coding", "is", "fun", "and", "challenging"]

补齐后（长度=6）:
句子1: ["I", "love", "coding", "[PAD]", "[PAD]", "[PAD]"]
句子2: ["Coding", "is", "fun", "and", "challenging", "[PAD]"]

这里的 [PAD] 是填充值，用于对齐句子长度。

1.2 为什么需要屏蔽 Padding Token 的 Loss？

填充值 [PAD] 不包含有效的语义信息。如果不屏蔽其 Loss，模型会被迫学习无意义的填充值，从而浪费计算资源，并可能影响模型的最终性能。因此，我们需要在计算 Loss 时忽略填充位置对应的 Token。

---

2. 如何实现屏蔽 Padding Token 的 Loss

实现屏蔽的关键在于两点：

1. 利用 attention_mask：
   attention_mask 是一个二进制向量，标记了输入序列中哪些位置是有效的 Token，哪些位置是填充值。例如：

   input_ids:        [1, 2, 3, 0, 0]  # 0 是填充值
   attention_mask:   [1, 1, 1, 0, 0]  # 1 表示有效，0 表示填充值
   

   模型通过 attention_mask 确保填充值不会影响计算过程。

2. 利用 labels 的特殊标记 -100：
   在 labels 中，填充位置的 Token 被设置为 -100，因为 Hugging Face 的 CrossEntropyLoss 损失函数支持通过 ignore_index 参数忽略特定的 Token。例如：

   labels:           [1, 2, 3, -100, -100]  # -100 表示填充值
   

   当 labels 的某些位置被标记为 -100，这些位置的 Loss 不会被纳入计算。

---

3. 代码实现：如何屏蔽 Padding Token 的 Loss

以下代码展示了如何在模型训练中屏蔽填充值的 Loss。

请注意，下面的-100是默认值，不用特别指定。

3.1 数据处理阶段：设置 label_pad_token_id=-100

在批量数据处理中，使用 Hugging Face 提供的 DataCollatorForSeq2Seq 工具，设置 label_pad_token_id=-100，自动将填充值的 labels 标记为 -100。

from transformers import DataCollatorForSeq2Seq

collate_fn = DataCollatorForSeq2Seq(
    tokenizer=tokenizer,  # 分词器
    model=model,          # 模型
    padding="longest",    # 对齐到当前 batch 中最长的句子
    label_pad_token_id=-100  # 确保填充值的标签被标记为 -100
)

• 参数解析：
  • padding="longest"：对齐批量数据中最长的句子。
  • label_pad_token_id=-100：填充值的标签被设置为 -100，确保在计算 Loss 时被忽略。

3.2 模型训练阶段：忽略 -100 标签

Hugging Face 的 CrossEntropyLoss 默认支持 ignore_index=-100，会自动忽略 labels 中等于 -100 的位置。因此，无需额外修改代码，直接训练即可。

from torch.nn import CrossEntropyLoss

# Hugging Face 模型会自动调用 CrossEntropyLoss
loss_fct = CrossEntropyLoss(ignore_index=-100)

# 示例计算
logits = model(input_ids).logits  # 模型输出
loss = loss_fct(logits.view(-1, logits.size(-1)), labels.view(-1))  # 计算损失

---

4. 完整流程分析

4.1 数据处理

在数据预处理中，attention_mask 标记填充值的位置，而 labels 中填充值的标签被标记为 -100：

输入数据:
input_ids:        [1, 2, 3, 0, 0]  # 填充值 0
attention_mask:   [1, 1, 1, 0, 0]  # 1 是有效位置
labels:           [1, 2, 3, -100, -100]  # 填充值标记为 -100

4.2 批处理阶段

通过 DataCollatorForSeq2Seq，自动对齐序列长度，并设置填充值的 labels 为 -100。

4.3 Loss 计算

• attention_mask 确保填充值不会被模型注意（Attention）。
• labels 中 -100 确保填充值的 Loss 被忽略。

---

5. 总结

在 NLP 模型训练中，通过以下方式屏蔽 Padding Token 的 Loss：

1. 使用 attention_mask：
   在模型的输入中，attention_mask 标记填充值的位置，确保模型不会关注无意义的 Padding。

2. 设置 label_pad_token_id=-100：
   在批处理阶段，使用 DataCollatorForSeq2Seq 将填充值对应的 labels 标记为 -100。

3. 使用 CrossEntropyLoss 的 ignore_index 参数：
   在 Loss 计算时，CrossEntropyLoss 会自动忽略 labels 中等于 -100 的位置。

通过这些机制，模型的训练更加高效，避免了无意义的计算浪费，最终提升了模型的性能和泛化能力。

补充

写一段代码，使用MODEL_NAME=google/gemma-2-2b 模型的tokenizer，输出label_pad_token_id是多少

from transformers import AutoTokenizer, DataCollatorForSeq2Seq

# Define the model name
MODEL_NAME = "google/gemma-2-2b"

# Load the tokenizer
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

# Initialize the DataCollatorForSeq2Seq
collate_fn = DataCollatorForSeq2Seq(
    tokenizer=tokenizer,
    model=None,  # Model is not needed to determine label_pad_token_id
    padding="longest"
)

# Output the label_pad_token_id
label_pad_token_id = collate_fn.label_pad_token_id
print(f"label_pad_token_id: {label_pad_token_id}")

根据 transformers 文档，如果没有显式指定 label_pad_token_id，它通常会默认使用 -100，因为这是 CrossEntropyLoss 的 ignore_index 默认值。

后记

2024年12月3日14点04分于上海，在GPT4o大模型辅助下完成。