Improving Language Models Trained on Translated Data with Continual Pre-Training and Dictionary Learning Analysis

TL;DR

This work tackles the challenge of training language models for low-resource languages when data must be sourced via machine translation. It translates TinyStories into Arabic using NLLB-3B, trains small to mid-sized models on the translated data, and evaluates them with a GPT-4 judge, revealing MT-induced linguistic and cultural biases. To mitigate these issues, the authors synthesize a small high-quality Arabic story dataset with Command R+ and perform continual pre-training on this data, analyzed through Sparse Auto-Encoders and Token Set Enrichment Analysis to track shifts in representations and bias. The results show that continual pre-training with a modest amount of high-quality data reduces MT pitfalls and biases, providing a cost-effective, interpretable strategy for improving small Arabic language models and potentially other low-resource languages.

Abstract

Training LLMs for low-resource languages usually utilizes data augmentation from English using machine translation (MT). This, however, brings a number of challenges to LLM training: there are large costs attached to translating and curating huge amounts of content with high-end machine translation solutions; the translated content carries over cultural biases; and if the translation is not faithful and accurate, data quality degrades causing issues in the trained model. In this work, we investigate the role of translation and synthetic data in training language models. We translate TinyStories, a dataset of 2.2M short stories for 3-4 year old children, from English to Arabic using the open NLLB-3B MT model. We train a number of story generation models of size 1M-33M parameters using this data. We identify a number of quality and task-specific issues in the resulting models. To rectify these issues, we further pre-train the models with a small dataset of synthesized high-quality Arabic stories generated by a capable LLM, representing 1% of the original training data. We show, using GPT-4 as a judge and Dictionary Learning Analysis from mechanistic interpretability, that the suggested approach is a practical means to resolve some of the machine translation pitfalls. We illustrate the improvements through case studies of linguistic and cultural bias issues.

Figures (13)

• Figure 1: The proposed TinyStories Arabic dataset is formed by translating 2M tiny stories from English to Arabic using NLLB-3B and synthesizing 20K Arabic tiny stories using Command R+ LLM. The former data is used to pre-train small language models (SLMs) with different architectures. The latter is used for continual pre-training. The models are qualitatively and quantitatively evaluated using a GPT judge. Further we train Sparse Auto-Encoder (SAE) on a selected SLM to analyze the model behavior.
• Figure 2: Scatter plot of feature enrichment scores for Arabic and English names in 2L-33M-ar.
• Figure 3: Scatter plot of feature enrichment scores for Arabic and English names for 2L-33M-ar-CP model after further pre-training.
• Figure 4: Manhattan plots of Enrichment Scores in (a) base and (b) further pre-trained models.
• Figure 5: Dashboard of feature #3114 from SAEs trained on the last MLP layer of continually trained model 2L-33M-ar-CP. The Feature #3114 shows that cultural bias was corrected after continual pre-training.