Reproduce “Matryoshka Model Learning for Improved Elastic Student Models” on OpenResearch
Turn your favorite coding agents into research agents that review literature, develop hypotheses, and run experiments.
Universal app for Apple silicon and Intel. Requires macOS 11 or later.
About this paper
Chetan Verma, Aditya Srinivas Timmaraju, Cho-Jui Hsieh, Suyash Damle, Ngot Bui, Yang Zhang, Wen Chen, Xin Liu and 2 more
Matryoshka Teaching Assistant (MatTA) is an elastic distillation framework from Google and Google DeepMind that co-trains a nested Teaching Assistant and Student model using online distillation and second-order optimization. This framework generates multiple servable models from a single training run, delivering an 8.05% relative improvement in AucLoss for a production relevance model, which resulted in a 20% improvement in a key live metric, and over 24% relative accuracy improvement on SAT Math for GPT-2 Medium.
Key insights
- MatTA effectively bridges the "relatability gap" in knowledge distillation by introducing an M-nested Teaching Assistant (TA) that is more structurally similar to the student model.
- The combination of Matryoshka architecture with second-order optimizers like Shampoo results in super-additive performance gains, demonstrating their synergistic benefits for complex parameter structures.
- A single training run can yield a spectrum of high-quality, elastic student models, enabling dynamic deployment choices while significantly reducing retraining costs.
Summary by alphaXiv.
Tip: change arxiv.org to autoarxiv.org in any arXiv abstract link to open the paper here.