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.

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

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

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)
Exact Match (EM)
- Description: Checks if critical keywords are present.
- Example:
  - Expected: "Paris"
  - Response: "The capital of France is Paris." (EM = 1)
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
Evaluation Repo: Knowledge Distillation Evaluation Repository

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

HenryHHHH/DistilLlama

Overview