Choosy Babies Need One Coach: Inducing Mode-Seeking Behavior in B aby L lama with Reverse KL Divergence (2024)

Abstract: This study presents our submission to the Strict-Small Track of the 2nd BabyLM Challenge. We use a teacher-student distillation setup with the BabyLLaMa model (Timiryasov and Tastet, 2023) as a backbone. To make the student’s learning process more focused, we replace the objective function with a reverse Kullback-Leibler divergence, known to cause mode-seeking (rather than mode-averaging) behaviour in computational learners. We further experiment with having a single teacher (instead of an ensemble of two teachers) and implement additional optimization strategies to improve the distillation process. Our experiments show that under reverse KL divergence, a single-teacher model often outperforms or matches multiple-teacher models across most tasks. Additionally, incorporating advanced optimization techniques further enhances model performance, demonstrating the effectiveness and robustness of our proposed approach. These findings support our idea that “choosy babies need one coach”.

## introduction one important feature of child language learning is its incrementality, gradually moving from simple to more complex language. when talking to a child, adults often choose to use simple words and expressions, effectively allowing the child to first focus on what's easy to learn (e.g., @xcite . in machine learning, this 'starting small' approach @xcite has informed the paradigm of curriculum learning @xcite , where models are trained using examples of increasing difficulty. in the 1st babylm challenge, organized in 2023 to stimulate training of language models on smallersized and child-appropriate data sets @xcite , curriculum learning was the most commonly used method among all submissions (e.g., @xcite @xcite . interestingly, despite its pop-ularity, curriculum learning did not yield consistent improvements over baselines @xcite . this suggests that while curriculum learning remains a valuable approach, other methods such as knowledge distillation and architectural modifications may offer additional advantages in certain contexts @xcite . in our submission to the 2nd babylm challenge, we leverage and combine some of the last year's successful approaches, while also enabling the learner to use a more selective learning strategy. more specifically, we take as starting point the baby llama model and its teacher-student knowledge distillation framework @xcite . we then experiment with changing its objective function from forward kl divergence to reverse kl divergence, inspired by @xcite ; @xcite , and implement several strategies to optimize the distillation process. unlike forward kl, which encourages the student model to approximate the full output distribution of the teacher and often leads to 'mass-covering' behavior, reverse kl focuses on high-probability outputs, helping the student to capture the teacher's main modes. this effectively results in a more selective, or 'choosy' learner. furthermore, while the original baby llama model was trained using an ensemble of two different teacher models, we demonstrate that having a single teacher is sufficient in our setup, which further speeds up the training process and leads us to observe that choosy babies only need to be trained by one coach (choobaca). ## methodology the general approach that leads to the development of choobaca is knowledge distillation. we start from baby llama as a backbone model @xcite and implement three important modifications. first, we change the stu-dent's objective function to reverse kl divergence, following @xcite ; @xcite . second, we replace baby llama's ensemble of two teachers with a single teacher, the original llama model @xcite . third, we implement several techniques to stabilize the distillation process, inspired by the minillm model of @xcite . in the remainder of this section, we unpack the general framework and each of our implemented modifications. ## results and discussion tables 2 and 3 present the evaluation results for all our model variants, as well as the two baselines adopted in the 2nd babylm challenge (the original baby llama model and ltg-bert), on the three benchmarks used in babylm (see previous section). whereas it is clear from the tables that there is no single best model, we can still observe several important patterns. our primary finding demonstrates that under reverse kl divergence (see rv models), knowledge distillation with a single-teacher model generally outperforms or matches the performance of models with two teachers. specifically, within the (super)glue benchmark (11 tasks), rv models with a single teacher outperform two-teacher rv models in 3 tasks (27%) and match their performance in 8 tasks (73%). within the other two benchmarks -ewok (11 tasks) and blimp (17 tasks), single-teacher rv models match the performance of two-teacher rv models across all tasks. these results support our hypothesis that a choosy, mode-seeking learning strategy enhances the choobaca model's ability to generalize effectively across diverse language understanding tasks under the reverse kl divergence framework. when comparing reverse kl divergence models (rv) to forward kl divergence models (fw), our results indicate that rv models achieve better or comparable performance in the majority of tasks. specifically, when trained on the 2024 data set, across the superglue and ewok benchmarks (22 tasks), rv models with a single teacher outperform fw models with two teachers in 3 tasks (14%) and match their performance in 18 tasks (81%), with fw models slightly outperforming rv models in 1 task (5%). this result highlights the effectiveness of inducing mode-seeking behavior through reverse kl divergence, as rv models focus on high-probability linguistic patterns, leading to improved generalization and performance across various benchmarks compared to the traditional fw approach with two teachers. models that incorporate additional optimization techniques (marked with a '+' character at the end) show better performance under the reverse kl divergence setting. these models outperform their non-optimized counterparts in 8 out of 22 (super)glue and ewok tasks (36%) and 7 out of 17 blimp tasks (41%). in the remaining tasks, their performance is comparable. similarly, when trained on the 2023 dataset, optimized rv models show improvements in 5 out of 22 (super)glue and ewok tasks (23%) and 4 out of 17 blimp tasks (24%), with performance remaining comparable in most other tasks. these findings suggest that the additional optimization techniques, namely step-wise loss computation and progressive training strategy, contribute to a more stable and efficient training processes, enabling the models to better capture complex linguistic structures. comparing the results for models trained on the 2023 vs. the 2024 dataset, we observe consistent patterns across the two. when trained on the 2023 dataset, rv models with a single teacher outperform two-teacher rv models in 12 out of 22 (super)glue and ewok tasks (55%) and 8 out of 17 blimp tasks (47%), with performance being comparable in 7 (super)glue and ewok tasks (32%) and 2 blimp tasks (12%). these patterns across both datasets once again support our main argument: choosy babies need one coach. notably, the performance of all models on the ewok benchmark is close to chance level, 50%. this suggests that our relatively small models might lack the capacity to effectively handle the complex ewok tasks, which are likely more demanding compared to the other benchmarks included in this study. overall, these results suggest that a selective, mode-seeking learning strategy, based on the use of reverse kl divergence with a single teacher model, enhances the choobaca model's ability to generalize effectively across diverse language understand-ing tasks. at the same time, all choobaca variants struggle with more complex tasks grounded in realworld knowledge. ## conclusion our findings support the use of reverse kullback-leibler divergence in knowledge distillation, particularly in a single-teacher setup. while it has been shown that multiple instructional sources can be advantageous @xcite , our results suggest that in a constrained setup with one small model trained on limited amounts of data, in combination with using reverse kl divergence, an ensemble of teachers may not be necessary. our single-teacher setup promotes mode-seeking behavior, resulting in a more focused and efficient learning process. it also simplifies the learning process and eliminates the need to train more than one teacher model. our choobaca model is able to efficiently generalize across diverse language understanding tasks. in future work, we plan to explore hybrid kl divergence methods, such as alternating between forward and reverse kl divergence or employing a weighted combination during training, to balance the learner's focus between dominant and minor modes. additionally, investigating layer-wise distillation -where different layers of the student model learn from different teachers -could more effectively accommodate varied distribution peaks. finally, we aim to examine dynamically averaging the outputs of multiple teachers before applying reverse kl divergence, which might smooth out the distribution and help the student model identify and prioritize significant modes without bias.