DistilLlama

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:

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:

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

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

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