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:

Model	Accuracy
FLAN T5 Small	0.0227
FLAN T5 Base	0.0318
FLAN T5 Large	0.0538
Trained Student	0.0243

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

Usage

python

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.