Implementation for soft offline distillation using saved top-k teacher logits#3382

Merged

copybara-service[bot] merged 9 commits intomainfrom

ajkv/offline-distillation-soft

Mar 20, 2026

Commits on Mar 19, 2026

Added train script for offline distillation training
ajkv-google
committed
updated code formatting and style
ajkv-google
committed
updated iterator to ensure weight updates when training student model
ajkv-google
committed
moved cmd args into the distillation config to make command easier to use
ajkv-google
committed
removed the need for hardcoding arrayrecord file and read directly from config
ajkv-google
committed
Removed redundant offline_distillation flag and relied on offline_data_dir to know when to run offfline vs online distillation
ajkv-google
committed
added unit tests to make sure offline and online distillation is loading the correct models
ajkv-google
committed
updated spacing and formatting
ajkv-google
committed
updated tests to add gradient accumulation flags in config
ajkv-google
committed