File size: 4,937 Bytes
1a6a8c3
 
67c47f1
1a6a8c3
67c47f1
 
 
1a6a8c3
67c47f1
1a6a8c3
67c47f1
 
 
 
 
 
 
1a6a8c3
67c47f1
 
 
1a6a8c3
67c47f1
1a6a8c3
67c47f1
 
 
 
 
 
 
 
 
 
1a6a8c3
67c47f1
1a6a8c3
67c47f1
1a6a8c3
67c47f1
1a6a8c3
67c47f1
 
 
 
 
 
 
 
 
 
 
 
1a6a8c3
67c47f1
1a6a8c3
67c47f1
 
1a6a8c3
67c47f1
 
1a6a8c3
67c47f1
1a6a8c3
67c47f1
 
 
 
1a6a8c3
67c47f1
 
 
 
 
1a6a8c3
67c47f1
 
 
 
 
1a6a8c3
67c47f1
1a6a8c3
67c47f1
 
 
 
 
1a6a8c3
67c47f1
1a6a8c3
67c47f1
 
 
1a6a8c3
67c47f1
1a6a8c3
67c47f1
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
---

### 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.

### 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 = \( \frac{\vec{A} \cdot \vec{B}}{||\vec{A}|| \, ||\vec{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 = \( \frac{LCS(R, C)}{\text{Length of } C} \)
     - Recall = \( \frac{LCS(R, C)}{\text{Length of } R} \)

### Model Evaluation Summary

| Model Name      | Duration (s) | Emissions Rate | Avg. EM | Avg. Cosine Similarity | Avg. ROUGE Score |
|-----------------|--------------|----------------|---------|------------------------|------------------|
| LLaMA-2-7B-HF   | 18215.61     | 1.01e-05       | 0.715   | 0.7257                 | 0.0821           |
| baby-llama-58m  | 57.20        | 4.79e-08       | 0.025   | 0.6556                 | 0.0097           |
| DistilLlama     | 77.12        | 1.01e-05       | 0.02    | 0.6623                 | 0.0115           |

### Acknowledgments

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

### 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*