LLM-KD: Knowledge Distillation of LLMs for Low-Resource Machine Translation

This repository contains the code used in the experiments for the paper:

"Knowledge Distillation of Large Language Models for Low-Resource Machine Translation: Sequence-Level Rewards vs. Token-Level Imitation"
Under review — Natural Language Processing, Special Issue on Language Models for Low-Resource Languages

---

Overview

This project investigates different strategies for transferring the translation capabilities of large teacher LLMs to smaller, more efficient student models, with a focus on low-resource machine translation. Four training paradigms are compared:

Script	Method	Description
train_sft.py	Supervised Fine-Tuning (SFT)	Fine-tunes a student model on teacher-generated synthetic translations (sequence-level knowledge distillation)
train_grpo.py	GRPO (Group Relative Policy Optimization)	Trains a student model using sequence-level rewards (chrF score) via reinforcement learning, without a teacher model at inference time
train_gkd.py	GKD (Generalized Knowledge Distillation)	On-policy token-level distillation that minimises a divergence between the student and teacher token distributions during training
train_gold.py	GOLD (General Online Logit Distillation)	An extension of GKD that supports student/teacher pairs with different tokenizers, supporting pure word-level distillation, on-policy distillation, or a hybrid of both

All scripts expect a synthetic parallel dataset pre-generated by a teacher LLM, stored in JSONL format under synthetic/<src_lang>-<tgt_lang>.json.

---

Requirements

Install the required Python packages:

pip install torch transformers datasets trl sacrebleu

---

Usage

SFT and GRPO

python train_<method>.py <src_lang> <tgt_lang> <model_name>

GKD and GOLD

python train_<method>.py <src_lang> <tgt_lang> <teacher_model> <student_model>

The GOLD Trainer GOLD supports:

• Word-Level Knowledge Distillation (WLKD): lmbda=0.0
• On-policy distillation: lmbda=1.0
• Hybrid (ULD): intermediate lmbda values (default 0.5) The lmbda and loss type (use_uld_loss, uld_use_hybrid_loss) can be modified directly in the script.

---

Output

Each script saves the fine-tuned model to a directory named after the training setup:

<src_lang>-<tgt_lang>-<model_name>-<METHOD>/