flan-t5-small-gsm8k-77m-gkd-jsd-lr3e4-b16-u4k-lam05-beta05

This model is a google/flan-t5-small student trained on openai/gsm8k with a local Generalized Knowledge Distillation (GKD) setup based on the paper On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes.

Source code:

• GitHub: https://github.com/Mantissagithub/opd/tree/main/t5_gsm8k

Training setup

• Student: google/flan-t5-small
• Teacher: google/flan-t5-large
• Dataset: openai/gsm8k
• Student size: about 77M parameters
• Divergence: jsd
• Beta: 0.5
• Lambda / student data fraction: 0.5
• Learning rate: 3e-4
• Micro batch size: 1
• Gradient accumulation: 16
• Effective batch size: 16
• Optimizer updates: 4000
• Max input length: 512
• Max output length: 320
• Generation max new tokens: 192
• Dropout: 0.05

Objective

The student is trained with a batch-level GKD loop:

• Supervised branch: train on GSM8K gold answer strings.
• On-policy branch: sample answers from the current student and distill the teacher signal on those sequences.
• Teacher forcing is used over the target sequence for both teacher and student.

Evaluation

Exact final-answer accuracy on the GSM8K test split:

The trained student improves slightly over the raw flan-t5-small baseline, but remains well below the large teacher.

Usage

from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

repo_id = "Pradheep1647/flan-t5-small-gsm8k-77m-gkd-jsd-lr3e4-b16-u4k-lam05-beta05"
tokenizer = AutoTokenizer.from_pretrained(repo_id)
model = AutoModelForSeq2SeqLM.from_pretrained(repo_id)

prompt = "question: Janet’s ducks lay 16 eggs per day. She eats three for breakfast every morning. How many eggs does she sell?\nanswer:"
inputs = tokenizer(prompt, return_tensors="pt")
outputs = model.generate(**inputs, max_new_tokens=128)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

Notes

• This is a research checkpoint, not a production-tuned math model.
• GSM8K remains difficult for a flan-t5-small student at this scale.