AntLM: Bridging Causal and Masked Language Models

TL;DR

AntLM introduces a unified pretraining framework that bridges Causal Language Modeling and Masked Language Modeling by alternating their objectives and corresponding attention masks. Using BabyLlama (CLM) and LTG-BERT (MLM) on the BabyLM 10M track, AntLM demonstrates macro-average improvements of about 1% and 2.2% over strong baselines, with ablations showing CLM and MLM contribute differently across tasks. The approach highlights the complementary strengths of generative and bidirectional context learning under data-scarce conditions, offering a path toward more data-efficient foundation models. Future work will scale resources and investigate additional alternating schedules to further improve generalization across evaluation tasks.

Abstract

Causal Language Modeling (CLM) and Masked Language Modeling (MLM) are two mainstream learning paradigms based on Transformer networks, specifically the Decoder-only and Encoder-only architectures. The strengths of each paradigm in downstream tasks have shown a mix of advantages and disadvantages. In the past BabyLM Challenge 2023, although the MLM paradigm achieved the best average performance, the CLM paradigm demonstrated significantly faster convergence rates. For the BabyLM Challenge 2024, we propose a novel language modeling paradigm named $\textbf{AntLM}$, which integrates both CLM and MLM to leverage the advantages of these two classic paradigms. We chose the strict-small track and conducted experiments on two foundation models: BabyLlama, representing CLM, and LTG-BERT, representing MLM. During the training process for specific foundation models, we alternate between applying CLM or MLM training objectives and causal or bidirectional attention masks. Experimental results show that combining the two pretraining objectives leverages their strengths, enhancing overall training performance. Under the same epochs, $AntLM_{BabyLlama}$ improves Macro-average by 1%, and $AntLM_{LTG-BERT}$ achieves a 2.2% increase over the baselines.

Figures (2)

• Figure 1: A diagram of AntLM$_\text{LTG-BERT}$. Based on the LTG-BERT architecture, we propose a joint MLM and CLM training objective. It is worth noting that the two objectives fully share parameters, but differ in their attention masks. The diagram also applies to AntLM$_\text{BabyLlama}$, with the difference in the architecture (e.g., positional encoding and the activation function of GLU).
• Figure 2: The phased experimental results on three datasets. The evaluation line chart for each stage of "3_CLM + 8_MLM + 2_CLM + 8_MLM + 3_CLM" on the BabyLlama model. The reason for the discontinuity in evaluation results between training phases is that we applied the evaluation method corresponding to the specific task categories at each stage of the training process.