In this new Machine Learning era dominated by LLMs, knowledge Distillation is going to be at the forefront of LLMOps. For widespread adoption and further development of generative ML, we first need to make those models more manageable to deploy and fine-tune.
Just to put some numbers on how unmanageable it can be, SOTA models these days have about ~500B parameters and that represents at least ~1TB of GPU memory to operate with specialized infrastructure. That's a minimum of ~$60,000 - $100,000 per year per deployed model just for inference servers. And that doesn't include fine-tuning nor typical elastic load balancing costs for reliability best practices. Not impossible, but somewhat a high barrier to entry for most businesses.
I always felt that knowledge distillation was a silent hero in this era of transformer-type language models. There are tons of distilled BERT-like models on HuggingFace for example. The concept behind distillation is actually pretty simple. Let's assume you have a large pre-trained model. In the context of LLMs, it could be pre-trained with self-supervised learning and fine-tuned in a RLHF fashion for example. Your pre-trained model now becomes the all-knowing teacher for a smaller student model. If we call the teacher model T and the student model S, we want to learn the parameters for S such that
T(x) = y_t ~ y_s = S(x)
For some data x, we want the predictions y_t and y_s by T and S to be as close to each other. To train a student model, we simply pass the training data through the teacher and student and update the student's parameters by minimizing the loss function l(y_t, y_s) and back-propagating its gradient. Typically we use cross-entropy for the loss function. To train the student model, think about typical supervised learning where the training data is the same or similar to the teacher's training data, but the ground truth label for the student model is the output prediction of the teacher model. You can read more about it in this survey: https://arxiv.org/pdf/2006.05525.pdf.
With the advent of prompt engineering, we now understand better how to extract the right piece of knowledge from LLMs. Techniques like Chain-of-Thought (CoT) greatly improved LLMs performance on few-shot learning tasks. The guys at Google just published an article (https://arxiv.org/pdf/2305.02301.pdf) utilizing CoT to improve the distillation process. The idea is to have the student LLM predicting the rationales for the predictions alongside the predictions and minimizing a loss function between the teacher rationale and the student rationale. Basically by forcing the LLM to explain its predictions, they were able to beat all the distillation SOTA. For example, they outperformed a 540B parameters PaLM model with a 770M parameters T5 model after distillation! I think this paper will have a huge impact in the coming year!
----
Find more similar content in my newsletter: https://TheAiEdge.io
Next ML engineering Masterclass starting July 29th: https://MasterClass.TheAiEdge.io
#machinelearning #datascience #artificialintelligence
Just to put some numbers on how unmanageable it can be, SOTA models these days have about ~500B parameters and that represents at least ~1TB of GPU memory to operate with specialized infrastructure. That's a minimum of ~$60,000 - $100,000 per year per deployed model just for inference servers. And that doesn't include fine-tuning nor typical elastic load balancing costs for reliability best practices. Not impossible, but somewhat a high barrier to entry for most businesses.
I always felt that knowledge distillation was a silent hero in this era of transformer-type language models. There are tons of distilled BERT-like models on HuggingFace for example. The concept behind distillation is actually pretty simple. Let's assume you have a large pre-trained model. In the context of LLMs, it could be pre-trained with self-supervised learning and fine-tuned in a RLHF fashion for example. Your pre-trained model now becomes the all-knowing teacher for a smaller student model. If we call the teacher model T and the student model S, we want to learn the parameters for S such that
T(x) = y_t ~ y_s = S(x)
For some data x, we want the predictions y_t and y_s by T and S to be as close to each other. To train a student model, we simply pass the training data through the teacher and student and update the student's parameters by minimizing the loss function l(y_t, y_s) and back-propagating its gradient. Typically we use cross-entropy for the loss function. To train the student model, think about typical supervised learning where the training data is the same or similar to the teacher's training data, but the ground truth label for the student model is the output prediction of the teacher model. You can read more about it in this survey: https://arxiv.org/pdf/2006.05525.pdf.
With the advent of prompt engineering, we now understand better how to extract the right piece of knowledge from LLMs. Techniques like Chain-of-Thought (CoT) greatly improved LLMs performance on few-shot learning tasks. The guys at Google just published an article (https://arxiv.org/pdf/2305.02301.pdf) utilizing CoT to improve the distillation process. The idea is to have the student LLM predicting the rationales for the predictions alongside the predictions and minimizing a loss function between the teacher rationale and the student rationale. Basically by forcing the LLM to explain its predictions, they were able to beat all the distillation SOTA. For example, they outperformed a 540B parameters PaLM model with a 770M parameters T5 model after distillation! I think this paper will have a huge impact in the coming year!
----
Find more similar content in my newsletter: https://TheAiEdge.io
Next ML engineering Masterclass starting July 29th: https://MasterClass.TheAiEdge.io
#machinelearning #datascience #artificialintelligence