CrossEntropyLoss has per-token weight that is different from original semantics?
Author: fzyzcjy
Pull request: no
Labels: name: bug

### System Info

Hi thanks for the library! Suppose we have two samples `[a,b,c,d,e,f,g]` and `[x,y]`. Then, the naive CrossEntropy will give: 

```
loss = 1/2( 1/7(loss(a) + loss(b) + ... + loss(g)) + 1/2(loss(x) + loss(y)) )
```

However, it seems that Transformers flatten everything and then compute cross entropy. Therefore, ths loss is indeed as if we have a sequence like `[a,b,c,d,e,f,g,x,y]`. And it will be

```
loss = 1/9(loss(a) + ... + loss(y))
```

Surely the two are different. So I wonder is this a bug, i.e. will this cause trouble? For example, suppose sample 1 is super long while sample 2 is super short, then the two ways will give very different weights for sample 2, possibly resulting in problems. But since huggingface transformers is so popular and is widely used, I wonder maybe it is OK to have that.

### Who can help?

@ArthurZucker 

### Information

- [ ] The official example scripts
- [ ] My own modified scripts

### Tasks

- [X] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

(see above)

### Expected behavior

(see above)

Comments:
I think the most common practice is using token level average. For example, the original gpt-2 training loss:
https://github.com/openai/gpt-2/blob/9b63575ef42771a015060c964af2c3da4cf7c8ab/src/model.py#L170


And response from Openai-o1:
When training large language models (LLMs), the choice of how to average the loss function can impact the learning dynamics and the model's performance. Specifically, deciding whether to average the loss over each **token** or over each **sentence** (or sequence) can lead to different outcomes.

**Averaging over Tokens:**

- **Equal Contribution of Tokens:** Averaging the loss over each token means that every token in the dataset contributes equally to the overall loss. This approach aligns with the nature of language modeling, where the goal is to predict the next token given the previous context.

- **Handling Variable-Length Sequences:** Text data often consists of sequences of varying lengths. By averaging over tokens, longer sentences naturally contribute more to the total loss because they contain more tokens. This allows the model to learn more from longer sequences, which may carry richer contextual information.

- **Stable Gradients:** Averaging over a larger number of tokens can lead to more stable gradient estimates during optimization, which can improve convergence during training.

**Averaging over Sentences:**

- **Equal Contribution of Sentences:** Averaging the loss over each sentence means that each sentence contributes equally to the total loss, regardless of its length. This can cause shorter sentences to have a disproportionately higher per-token influence compared to longer sentences.

- **Bias Towards Shorter Sentences:** Since shorter sentences have fewer tokens, averaging over sentences can amplify their impact on the loss. This might bias the model towards optimizing for shorter sequences at the expense of longer ones.

- **Inconsistent Token-Level Learning:** Language models typically operate at the token level, predicting one token at a time. Averaging loss at the sentence level might misalign with the token-level prediction objective, potentially hindering the model's ability to learn accurate token probabilities.

**Recommendation:**

It is generally better to average the loss over each **token** when training LLMs. This approach ensures that:

- **Tokens are Treated Equally:** Each token, regardless of the sentence it belongs to, contributes equally to the training objective.

- **Natural Emphasis on Longer Contexts:** Longer sentences, which provide more contextual information, have a proportionally larger impact on the total loss, aiding the model in learning from complex patterns.

- **Alignment with Modeling Objectives:** Since LLMs predict tokens sequentially, using a token-level loss aligns with the model's prediction task and encourages learning at the appropriate granularity.

**Conclusion:**

Averaging the loss over tokens is the preferred method when training large language models. It provides a balanced and effective way of updating the model parameters, ensuring that learning occurs uniformly across all tokens and sequences in the dataset.

@techkang That looks reasonable, thank you!

Closing for now, feel free to reopen if you have any other questions!