深入解析 Loss 減少方式：mean和sum的區別及其在大語言模型中的應用（中英雙語）

深入解析 Loss 減少方式：`mean` 和 `sum` 的區別及其在大語言模型中的應用

在訓練大語言模型（Large Language Models, LLM）時，損失函數（Loss Function）的處理方式對模型的性能和優化過程有顯著影響。本文以 reduce_loss 參數為例，詳細探討 mean 和 sum 兩種方式的定義、適用場景及其對對話模型性能的潛在提升原因，并通過代碼實例加深理解。

1. 什么是 `reduce_loss`？

reduce_loss 決定了在每個 batch 中，如何對 token-level 的損失進行歸一化或累加處理。常見的選項是：

mean: 取每個 token 損失的平均值。
sum: 將每個 token 損失直接累加。

參數定義示例（在代碼中通過 dataclass 定義）：參考來源：https://github.com/allenai/open-instruct

from dataclasses import dataclass, field@dataclass
class TrainingArguments:reduce_loss: str = field(default="mean",metadata={"help": ("How to reduce loss over tokens. Options are 'mean' or 'sum'.""Using 'sum' can improve chat model performance.")},)

2. `mean` 和 `sum` 的定義

2.1 `mean` 模式

定義：將 batch 中所有 token 的損失值取平均。
公式：
$\text{Loss}_{\text{mean}} = \frac{\sum_{i=1}^{N} \text{Loss}_i}{N}$
其中 ( $N$ ) 是當前 batch 中的 token 總數。
特性：每個 token 的損失對最終的 loss 貢獻相等，損失值與 batch 中的 token 數無關。

2.2 `sum` 模式

定義：將 batch 中所有 token 的損失值直接累加。
公式：
$\text{Loss}_{\text{sum}} = \sum_{i=1}^{N} \text{Loss}_i$
特性：長序列（更多 token）的損失對總 loss 的貢獻更大，損失值直接與 token 數成正比。

3. `mean` 和 `sum` 的區別

模式	特點	優點	缺點
`mean`	損失對 token 數歸一化，獨立于 batch size。	穩定性強，適用于 token 數差異大的批次。	長序列與短序列對損失的貢獻相同，可能弱化長序列的重要性。
`sum`	損失值與 token 總數成正比，長序列貢獻更大。	在注重長序列表現的任務中效果更好（如對話生成）。	損失值隨 batch size 變化波動，需要動態調整學習率。

4. 適用場景分析

4.1 `mean`

適用任務：大多數語言建模任務，如 GPT 或 BERT 的預訓練。
適用場景：當訓練數據中序列長度差異較大時，mean 可以避免因長序列的損失值過大而導致梯度更新不均衡。

4.2 `sum`

適用任務：對長序列表現要求較高的任務，如對話生成（Chat Models）和長文本生成。
適用場景：長序列的損失占比更高，從而使優化過程更加關注全局上下文的建模。

5. 為什么 `sum` 能提升對話模型性能？

對話模型（Chat Models）的訓練中，長序列往往包含豐富的上下文信息，而短序列則可能無法體現模型的上下文理解能力。在 sum 模式下：

長序列的重要性增加：長序列的損失對總損失的貢獻更大，這促使模型更關注上下文的建模。
對全局一致性更敏感：sum 模式下，模型的優化方向更傾向于全序列的一致性，特別適合需要長距離依賴的任務。

示例：
假設一個 batch 包含以下兩個樣本：

樣本 A: 長度為 10，損失總和為 5。
樣本 B: 長度為 50，損失總和為 25。

計算損失貢獻：

mean 模式：
$\text{Loss}_{\text{mean}} = \frac{5 + 25}{10 + 50} = 0.5$
樣本 A 和 B 的貢獻權重相同。
sum 模式：
$\text{Loss}_{\text{sum}} = 5 + 25 = 30$
樣本 B 的貢獻權重顯著增加，優化更關注長序列。

6. 實戰代碼

以下是一個完整的訓練腳本，展示如何在 Hugging Face 的 transformers 框架中使用 reduce_loss 參數。

from transformers import AutoModelForCausalLM, AutoTokenizer
from datasets import load_dataset
from torch.utils.data import DataLoader
import torch# 模型和數據集
model_name = "meta-llama/Llama-3.1-8B"
dataset_name = "allenai/tulu-3-sft-mixture"model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False)dataset = load_dataset(dataset_name)
tokenized_dataset = dataset.map(lambda x: tokenizer(x['text'], truncation=True, padding="max_length"), batched=True)
train_loader = DataLoader(tokenized_dataset["train"], batch_size=2, shuffle=True)# 訓練設置
reduce_loss = "sum"  # 改為 "mean" 可對比效果
optimizer = torch.optim.AdamW(model.parameters(), lr=5e-6)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)# 訓練循環
for epoch in range(2):for batch in train_loader:inputs = torch.tensor(batch["input_ids"]).to(device)labels = inputs.clone()outputs = model(inputs, labels=labels)if reduce_loss == "sum":loss = outputs.loss.sum()else:  # 默認 "mean"loss = outputs.loss.mean()loss.backward()optimizer.step()optimizer.zero_grad()print(f"Epoch: {epoch}, Loss: {loss.item()}")

7. 注意事項與優化建議

動態調整學習率：
- 使用 sum 時，由于損失值放大，建議適配學習率，如降低到 mean 模式的 ( $1/ N$ )。
- 配合學習率調度器（如 linear）優化訓練。
對長短序列的平衡：
- 若長序列權重過大導致模型性能退化，可結合 curriculum learning 或混合訓練策略（如對長短序列按比例采樣）。
性能評估：
- 在驗證集上，關注長序列和短序列的生成性能對比。

8. 總結

reduce_loss 的選擇對模型性能有直接影響：

mean 更通用，適合大多數語言建模任務。
sum 在對話生成等長序列敏感任務中表現更優。

希望本文能為 LLM 研究人員提供思路和參考，在具體任務中靈活選擇合適的損失歸一化方式，從而提升模型性能。

Understanding the Difference Between `mean` and `sum` Loss Reduction in LLM Training

When training large language models (LLMs), the way token-level loss is reduced across a batch can significantly impact optimization and model performance. This article delves into the reduce_loss parameter, exploring the differences between mean and sum reduction modes, their definitions, use cases, and why sum might improve the performance of chat-oriented models. Practical code examples are also provided for clarity.

1. What is `reduce_loss`?

The reduce_loss parameter determines how the token-level loss values in a batch are aggregated. The two most common options are:

mean: Averages the loss over all tokens in a batch.
sum: Sums the loss of all tokens in a batch.

Example definition (from the codebase using Python dataclass):

from dataclasses import dataclass, field@dataclass
class TrainingArguments:reduce_loss: str = field(default="mean",metadata={"help": ("How to reduce loss over tokens. Options are 'mean' or 'sum'.""Using 'sum' can improve chat model performance.")},)

2. Definitions of `mean` and `sum`

2.1 `mean`

Definition: Averages the loss across all tokens in a batch.
Formula:
$\text{Loss}_{\text{mean}} = \frac{\sum_{i=1}^{N} \text{Loss}_i}{N}$
where ( $N$ ) is the total number of tokens in the batch.
Characteristics: The contribution of each token to the final loss is normalized, making the loss independent of the batch’s token count.

2.2 `sum`

Definition: Sums up the loss across all tokens in a batch.
Formula:
$\text{Loss}_{\text{sum}} = \sum_{i=1}^{N} \text{Loss}_i$
Characteristics: The total loss is proportional to the number of tokens, giving longer sequences more weight in the optimization process.

3. Key Differences Between `mean` and `sum`

Reduction Mode	Characteristics	Advantages	Disadvantages
`mean`	Normalizes the loss by token count.	Stable and robust for datasets with variable-length sequences.	Long sequences are underweighted relative to short ones.
`sum`	Loss scales with the number of tokens.	Places greater emphasis on longer sequences, improving performance in tasks requiring context modeling.	Loss values vary with batch size, necessitating dynamic learning rate adjustment.

4. Use Cases for `mean` and `sum`

4.1 `mean`

Best Suited For: Pretraining or general language modeling tasks like GPT or BERT.
Scenario: When the dataset contains sequences of widely varying lengths, mean ensures that longer sequences do not disproportionately influence gradient updates.

4.2 `sum`

Best Suited For: Tasks requiring high performance on long sequences, such as dialogue generation or document-level text generation.
Scenario: Encourages the model to prioritize sequences with richer contexts, as their loss contributes more to the overall optimization.

5. Why Does `sum` Improve Chat Model Performance?

In chat-oriented models, sequences are typically longer and require the model to understand and generate coherent responses over extended contexts. Using sum mode:

Enhances Long Sequence Weighting: Longer sequences contribute more to the total loss, emphasizing the importance of context modeling.
Encourages Global Consistency: By assigning more weight to longer contexts, the model better captures dependencies across the entire sequence.
Balances Token Importance: Since chat models are often evaluated on dialogue-level coherence, sum ensures that tokens from the context and the response are proportionally weighted.

Example:
Consider a batch with two samples:

Sample A: Sequence length = 10, loss = 5.
Sample B: Sequence length = 50, loss = 25.

Loss calculations:

mean mode:
$\text{Loss}_{\text{mean}} = \frac{5 + 25}{10 + 50} = 0.5$
Both samples contribute equally to the loss.
sum mode:
$\text{Loss}_{\text{sum}} = 5 + 25 = 30$
Sample B contributes much more to the total loss, focusing the optimization on longer contexts.

6. Practical Implementation

Here’s a practical training script that demonstrates the use of reduce_loss in both modes.

from transformers import AutoModelForCausalLM, AutoTokenizer
from datasets import load_dataset
from torch.utils.data import DataLoader
import torch# Model and dataset
model_name = "meta-llama/Llama-3.1-8B"
dataset_name = "allenai/tulu-3-sft-mixture"model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False)dataset = load_dataset(dataset_name)
tokenized_dataset = dataset.map(lambda x: tokenizer(x['text'], truncation=True, padding="max_length"), batched=True)
train_loader = DataLoader(tokenized_dataset["train"], batch_size=2, shuffle=True)# Training setup
reduce_loss = "sum"  # Change to "mean" to compare effects
optimizer = torch.optim.AdamW(model.parameters(), lr=5e-6)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)# Training loop
for epoch in range(2):for batch in train_loader:inputs = torch.tensor(batch["input_ids"]).to(device)labels = inputs.clone()outputs = model(inputs, labels=labels)if reduce_loss == "sum":loss = outputs.loss.sum()else:  # Default: "mean"loss = outputs.loss.mean()loss.backward()optimizer.step()optimizer.zero_grad()print(f"Epoch: {epoch}, Loss: {loss.item()}")

7. Practical Considerations

Learning Rate Adjustment:
- When using sum, the loss magnitude increases with batch size, so you may need to adjust the learning rate (e.g., scale it down by ( $1/ N$ )).
Balancing Long and Short Sequences:
- Overweighting long sequences can sometimes harm generalization. Using curriculum learning or sampling strategies (e.g., proportional sampling) can help mitigate this.
Validation:
- Evaluate model performance on both short and long sequences to confirm improvements in the intended metrics.

8. Conclusion

The choice between mean and sum loss reduction modes depends on the specific task and dataset:

Use mean for general-purpose language modeling tasks where sequence lengths vary significantly.
Use sum for tasks that prioritize long-sequence performance, such as chat models or long-text generation.

Understanding and experimenting with these settings can lead to better-optimized models, particularly in the nuanced field of LLM fine-tuning.