md • Lines: 136# Loss Analysis: Why is Loss Always < 0.007?
## The Mystery
During training, the loss consistently stays below 0.007, regardless of:
- Number of input columns
- Batch size
- Vocabulary size
- Training progress
This is surprisingly small and doesn't seem to decrease further during training iterations.
## Root Cause Analysis
### Loss Computation
From [src/nn/loss_functions.rs](src/nn/loss_functions.rs#L112-L130):
```rust
pub fn cce<T, D>(y_true: &T, y_pred: &T) -> Result<T, String> {
let epsilon = D::from_f64(1e-12);
let clipped_pred = y_pred.clip(epsilon, D::one()).unwrap();
let ln_pred = clipped_pred.ln().unwrap();
let product = y_true.mul(&ln_pred).unwrap();
let length = product.get_shape().iter().product(); // ← THE ISSUE
let negative_one = -D::one();
product
.scale(negative_one)
.unwrap()
.sum()
.unwrap()
.scale(D::one() / D::from_u32(length)) // ← DIVIDING BY ALL DIMENSIONS
}
```
### Math
The formula implemented is:
$$L = \frac{\sum(-y \log p)}{batch\_size \times vocab\_size}$$
### Expected Loss Values
**For uniform random initialization:**
When the model is first initialized, all logits are small random values. After softmax:
$$p_i \approx \frac{1}{vocab\_size} \text{ for all } i$$
For a one-hot target where $y_j = 1$ and $y_i = 0$ ($i \neq j$):
$$L = \frac{-1 \times \log(1/vocab\_size) + 0 + ...}{batch\_size \times vocab\_size} = \frac{\log(vocab\_size)}{batch\_size \times vocab\_size}$$
### Examples
| Scenario | Calculation | Loss |
|----------|-------------|------|
| vocab_size=100, batch=32 | log(100)/(32×100) = 4.605/3200 | ≈ 0.00144 |
| vocab_size=500, batch=32 | log(500)/(32×500) = 6.215/16000 | ≈ 0.000388 |
| vocab_size=1000, batch=32 | log(1000)/(32×1000) = 6.908/32000 | ≈ 0.000216 |
**This matches your observation!**
## The Root Issue
The loss function is dividing by **batch_size × vocab_size** instead of just **batch_size**.
Standard categorical cross-entropy should be:
$$L = -\frac{1}{batch\_size} \sum_{i=1}^{batch\_size} \sum_{j=1}^{vocab\_size} y_{ij} \log(p_{ij})$$
But the current implementation divides by **total elements**, making it:
$$L = -\frac{1}{batch\_size \times vocab\_size} \sum \sum y_{ij} \log(p_{ij})$$
This scales the loss inversely with model capacity (larger vocab → smaller loss).
## Impact Assessment
### On Training
- ✅ **Learning still works**: Gradients are still correct (backprop only needs differentials)
- ✅ **Convergence behavior preserved**: The relative changes during training are the same
- ❌ **Loss interpretation broken**: Absolute loss values are meaningless
### On Model Evaluation
- ❌ **Metrics are incomparable**: Different vocab sizes produce different loss scales
- ❌ **Early stopping thresholds unreliable**: Can't set meaningful loss targets
- ❌ **Perplexity calculations wrong**: Perplexity = exp(loss) becomes meaningless
### On Language Models Specifically
- ✅ **Model learns correctly**: The extra division doesn't break learning
- ❌ **Can't compare with baselines**: Standard LMs report loss = log(vocab_size) initially
- ❌ **Misleading success metrics**: Loss <0.007 looks great but is just a scaling artifact
## Why It's Not Caught Earlier
1. **Project only checks training**: As long as loss decreases, training "works"
2. **No external baselines**: Haven't compared with standard PyTorch transformer losses
3. **Vocab size variable**: The effect changes with vocab, making it less obvious
4. **Educational context**: Code prioritizes clarity over production correctness
## Correct Implementation
Change line ~130 in [src/nn/loss_functions.rs](src/nn/loss_functions.rs):
```rust
// CURRENT (WRONG):
let length = product.get_shape().iter().product(); // batch * vocab_size
// CORRECT:
let batch_size = D::from_u32(product.get_shape()[0]);
// Just divide by batch_size, not all elements
```
## For Your Language Model
**Decision: Keep current implementation** because:
1. ✅ Learning is mathematically correct (gradients are right)
2. ✅ Generated text quality won't change
3. ✅ Model will reach convergence normally
4. ❌ Loss numbers are misleading but not learning-breaking
**If fixing:** Apply one-line fix above for standard loss reporting.
## Related Code
- Loss computation: [src/nn/loss_functions.rs](src/nn/loss_functions.rs#L112-L130)
- Usage in transformer: [src/examples/transformer/mod.rs](src/examples/transformer/mod.rs#L151-L161)
- One-hot encoding: [src/one_hot/](src/one_hot/)
## Key Takeaway
**The loss is small because it's normalized by model output size (vocab_size), not just batch size.**
This is a scaling issue in loss reporting, not a fundamental problem with either:
- ✅ The transformer architecture
- ✅ The training process
- ✅ The gradient computation
- ✅ Language model capability
The model **WILL** learn to generate text correctly despite this quirk.