transformers safetensors llama text-generation knowledge-distillation causal-lm openwebtext wikitext transfer-learning en dataset:openwebtext dataset:wikitext arxiv:2308.02019 license:apache-2.0 autotrain_compatible text-generation-inference endpoints_compatible region:us
---
base_model: meta-llama/LLaMA-2-7B
datasets:
- openwebtext
- wikitext
language: en
library_name: transformers
license: apache-2.0
metrics:
- cosine-similarity
- exact-match
- rouge
model_name: DistilLLaMA
tags:
- text-generation
- knowledge-distillation
- llama
- causal-lm
- openwebtext
- wikitext
- transfer-learning
parameter_count: 80M
---


### Overview

This model is a distilled version of LLaMA 2, containing approximately 80 million parameters. It was trained using a mix of OpenWebText and WikiText Raw V1 datasets. Knowledge distillation was employed to transfer knowledge from a larger "teacher" model—Meta’s 7B LLaMA 2—to help this smaller model mimic the behavior of the teacher.

### Model Architecture

The architecture is based on LLaMA 2, with the following parameters:
| Parameter               | Value |
|-------------------------|-------|
| Hidden Dimension        | 512   |
| Intermediate Dimension  | 1536  |
| Max Positional Embeddings | 128   |
| Attention Heads         | 8     |
| Transformer Layers      | 16    |


### Training Process

During each training step, the input data \( X \) is fed to both the teacher and student models. The student model calculates output logits and loss with the true labels, while the teacher model only generates logits. The total loss combines task-specific loss and distillation loss:

```python
def distillation_loss(student_logits, teacher_logits, temperature=2.0):
    return F.kl_div(
        F.log_softmax(student_logits / temperature, dim=-1),
        F.softmax(teacher_logits / temperature, dim=-1),
        reduction='batchmean'
    ) * (temperature ** 2)

# Loss Calculation
loss = (alpha * distill_loss) + ((1 - alpha) * task_loss)
```

### Training Configuration

- **Batch Size**: 64
- **Max Sequence Length**: 128
- **Epochs**: 2
- **Log Interval**: 3000
- **Learning Rate**: 3e-4
- **Warmup Steps**: 4000
- **Accumulation Steps**: 8
- **Load Model**: True
- **Temperature**: 2.0
- **Alpha**: 0.3

### Performance Comparison

The model’s performance is evaluated on 200 queries created in-house. For more details, visit the [GitHub repository](https://github.com/svarnim1805/Knowledge-Distillation/tree/main).

#### Example Queries

| **Query**                                    | **Keyword**       |
|----------------------------------------------|-------------------|
| The capital of France is                     | Paris             |
| The author of *To Kill a Mockingbird* is     | Harper Lee        |
| The planet known as the Red Planet is        | Mars              |
| The smallest country in the world is         | Vatican City      |
| The *Mona Lisa* was painted by               | Leonardo da Vinci |
| The tallest mountain in the world is         | Mount Everest     |
| The host country of the 2016 Summer Olympics | Brazil            |
| The currency of Japan is                     | Yen               |
| The "Father of Computers" is                 | Charles Babbage   |
| The largest ocean on Earth is                | Pacific Ocean     |

### Text Generation Examples

1. **Input**: *The capital of France is*
   - **Output**: "The capital of France is located in the southern province of Lyon, France. The capital is the main hub of the French capital, La Caillion, and the main hub of the French capital. A major part of the French capital is in the southern city of Paris."

2. **Input**: *The tallest mountain in the world is*
   - **Output**: "The tallest mountain in the world is Mount Everest, which is located in the northwest corner of the village of Kahuna. The mountain is about 1,000 feet (3,000 m) above sea level."

### Evaluation Metrics

1. **Cosine Similarity using Word Embeddings**
   - **Description**: Measures semantic similarity by mapping words/phrases to vectors.
   - **Equation**: Cosine Similarity = ( A • B ) / ( ||A|| ||B|| )
   - **Example**: "The dog chased the cat." vs. "A canine pursued a feline." (High similarity)

2. **Exact Match (EM)**
   - **Description**: Checks if critical keywords are present.
   - **Example**: 
     - Expected: "Paris"
     - Response: "The capital of France is Paris." (EM = 1)

3. **ROUGE Score**
   - **Description**: Measures the overlap of the longest common subsequences between reference and response texts.
   - **Equation**:
     - Precision = Precision = LCS(R, C) / Length of C
     - Recall = Recall = LCS(R, C) / Length of R

### Model Evaluation Summary

| Model Name      | Duration (s) | Emissions (kgCO₂e) | Avg. EM | Avg. Cosine Similarity | Avg. ROUGE Score |
|-----------------|--------------|--------------------|---------|------------------------|------------------|
| LLaMA-2-7B-HF   | 18215.61     | 1.84e-01           | 0.715   | 0.7257                 | 0.0821           |
| baby-llama-58m  | 57.20        | 2.73e-06           | 0.025   | 0.6556                 | 0.0097           |
| DistilLlama     | 77.12        | 7.79e-04           | 0.02    | 0.6623                 | 0.0115           |

*Note: CodeCarbon was used to track carbon emission. Allocated 80GB memory, 32 cores, Intel(R) Xeon(R) Gold 6448H for the evaluation*

### Acknowledgments

- **University of Melbourne**
- **AGL Energy**
- **My teammates**: Svarnim and Mohit

### GitHub Repositories

- **Training Repo**: [DistilLlama Training Repository](https://github.com/HenryHuang2/DistilLlama)
- **Evaluation Repo**: [Knowledge Distillation Evaluation Repository](https://github.com/svarnim1805/Knowledge-Distillation)

### Reference

@misc{timiryasov2023babyllamaknowledgedistillation,
      title={Baby Llama: knowledge distillation from an ensemble of teachers trained on a small dataset with no performance penalty}, 
      author={Inar Timiryasov and Jean-Loup Tastet},
      year={2023},
      eprint={2308.02019},
      archivePrefix={arXiv},
      primaryClass={cs.CL},
      url={https://arxiv.org/abs/2308.02019}, 
}

*Note: The repository will be updated as training progresses. Last update 2024-10-23*