When Babies Teach Babies: Can student knowledge sharing outperform Teacher-Guided Distillation on small datasets? (2024)

Abstract: We present our submission to the BabyLM challenge, aiming to push the boundaries of data-efficient language model pretraining. Our method builds upon deep mutual learning, introducing a student model search for diverse initialization. We address the limitation of treating students equally by formulating weighted mutual learning as a bi-level optimization problem. The inner loop learns compact students through online distillation, while the outer loop optimizes weights for better knowledge distillation from diverse students. This dynamic weighting strategy eliminates the need for a teacher model, reducing computational requirements. Our evaluations show that teacher-less methods can match or surpass teacher-supervised approaches.

## introduction the substantial computational and memory requirements of large language models pose significant challenges for deployment on intelligent edge systems, where resources are often constrained. as the demand for real-time processing and lowlatency responses increases in edge computing environments, the need for lightweight and memoryefficient models becomes critical. recent research, notably the chinchilla paper @xcite ), demonstrated that a 70b parameter model trained on 1.4 trillion tokens outperformed larger models with less data, highlighting the intricate balance between model size and training data. this massive data requirement-equivalent to over 10,000 times the words a 13-year-old encounters-is becoming a significant bottleneck. to address these challenges, several techniques have emerged such as network pruning @xcite ), quantization @xcite ), neural architecture search @xcite and knowledge distillation @xcite , @xcite al. (2020),wang et al. (2022)) in response to these challenges, the babylm challenge invites researchers to explore the limits of data-efficient language model pretraining @xcite . participants are constrained to training their models on limited text corpora: 10m and 100m word text-only tracks and a newly introduced multimodal track containing 50m words of paired text-image data, and 50m words text-only data. our paper describes our submission to the 10m and 100m text-only tracks. it builds upon the approach of weighted mutual learning @xcite introducing key modifications to enhance generalizability. our methodology focuses on distilling a roberta-base model (125m parameters) to less than half its size while maintaining performance. our main contributions include : the inner loop minimizes a combined loss to train individual student models, consisting of a supervised learning loss and a kl divergence loss that aligns each student's class posterior with others'. (b) instead of treating each student model equally, we introduce an outer loop to optimize student importance weights by minimizing the ensemble loss. this approach generally performed better than both conventional supervised learning and traditional distillation from a larger pretrained teacher. notably, our weighted mutual learning strategy can improve performance even among several large networks compared to independent learning, challenging the conventional understanding that distillation requires a larger, more powerful teacher. ## related work the vanilla distillation @xcite method consists of two stages, firstly train a large teacher model, followed by transfer of soft logits to a smaller student model. also known as offline distillation, it keeps the teacher fixed, only allowing a one-way knowledge transfer. to reduce memory consumption of training a large teacher model, @xcite proposed an online distillation framework called mutual learning where a group of student (or student) models were trained simultaneously. although, online distillation eliminated the teacher model, similar networks in online distillation may prevent the students from learning knowledge from the students zhang et al.. recent approaches have attempted to induce diversity in online distillation to improve overall performance. @xcite proposed inducing data diversity by training student models with varying image augmentations. however, this method relies heavily on data augmentations, which can be unpredictable in real-world deployment scenarios. @xcite introduced an adaptive ensemble knowledge distillation method using multiple diverse teacher models to train a student model. while this approach shows promise, it requires maintaining several teacher models, leading to increased memory usage and computational overhead. the reported accuracy improvements are also relatively modest, typically ranging from 0.5% to 1% across benchmarks. our approach closely resembles to that of zhang et al.. they present a diversity induced weight mutual learning approach for distillation. they introduce diversity by assigning varying pruning ratios to different student models. although this method reduces memory consumption, the manual assignment of pruning ratios may not generalize well across different architectures and tasks. the reported performance gains are limited, with improvements of less than 0.5% on most benchmarks. as shown by @xcite , while pruning induces sparsity within networks and can reduce computational complexity (measured in flops), the relationship between pruning percentage and actual model size reduction is not always linear. moreover, in zhang et al., we observe a performance drop when pruning beyond 30%, indicating a trade-off between model compression and accuracy. ## results this section provides the results of the empirical evaluation of dwml. first, we compare our method to baselines, then we compare our method with other distillation methods and then we perform an ablation study of different dwml variations. ## conclusion in this paper, we introduced diversity induced weighted mutual learning (dwml) as an alternative to teacher-supervised knowledge distillation. while our approach showed modest improvements over the roberta-base baseline, it was the simpler self-distillation method that achieved the strongest performance. our ablation studies on our approach (dwml) revealed that two-peer configurations offered optimal efficiency, a balanced loss function (α = 0.5) was crucial, and model performance correlated strongly with both dynamically learned importance weights and model size. regarding computational efficiency, while dwml showed the lowest average gpu utilization, it required longer training times