Source-linked AI summary
UniMax: Fairer and more Effective Language Sampling for Large-Scale Multilingual Pretraining
Hyung Won Chung, Noah Constant, Xavier Garcia, Adam Roberts, Yi Tay, Sharan Narang, Orhan Firat
TL;DR
Multilingual pretraining must choose language proportions despite severe data imbalance and limited evidence across model scales. The paper introduces UniMax, which seeks uniform language coverage while capping corpus repeats, and evaluates it across multilingual benchmarks and scales. UniMax outperforms temperature sampling, with benefits persisting as scale increases, while the study releases a refreshed multilingual corpus and pretrained checkpoints.
Problem
The open problem is how to balance pretraining languages with highly unequal data availability, especially across model scales.
Method
UniMax allocates training to make language coverage more uniform while explicitly capping repeats of each language corpus, and evaluates sampling strategies across benchmarks and model scales.
Results
UniMax outperforms temperature sampling across multilingual benchmarks, with benefits persisting as model scale increases.
Takeaways & Limitations
UniMax provides language distributions that approach uniformity without harmful repetition and performs well across benchmarks and scales up to 13 billion parameters.
Takeaways & Limitations
The study covers a single pretraining paradigm, leaving encoder-only, decoder-only, parallel-data, and language-specific-parameter settings for future work.
Abstract
from arXiv · showhide
Pretrained multilingual large language models have typically used heuristic temperature-based sampling to balance between different languages. However previous work has not systematically evaluated the efficacy of different pretraining language distributions across model scales. In this paper, we propose a new sampling method, UniMax, that delivers more uniform coverage of head languages while mitigating overfitting on tail languages by explicitly capping the number of repeats over each language's corpus. We perform an extensive series of ablations testing a range of sampling strategies on a suite of multilingual benchmarks, while varying model scale. We find that UniMax outperforms standard temperature-based sampling, and the benefits persist as scale increases. As part of our contribution, we release: (i) an improved and refreshed mC4 multilingual corpus consisting of 29 trillion characters across 107 languages, and (ii) a suite of pretrained umT5 model checkpoints trained with UniMax sampling.
1 INTRODUCTION
Multilingual pretraining must balance languages with vastly different data availability, but temperature sampling can excessively repeat low-resource data. UniMax addresses this by allocating training more uniformly while explicitly limiting repeats.
- English has roughly 9.7 trillion characters in mC4, over 92,000 times the size of Yoruba.
- Temperature sampling can repeat the lowest-resource languages excessively when tuned to balance higher-resource languages.The problem is illustrated using τ = 3.33 and a trillion-token budget.
- UniMax allocates a fixed training-token budget to underrepresented datasets using a maximum-repeat limit, then prioritizes linguistic utility with the remaining budget.
- UniMax controls repeats for every language, countering overfitting on low-resource languages without reprioritizing higher-resource languages.
- The paper evaluates sampling strategies across multilingual benchmarks and model scales, and releases a refreshed 29-trillion-character corpus across 107 languages plus UniMax-trained checkpoints.
2 RELATED WORK
Prior work identifies capacity competition and language-balancing challenges, but systematic comparisons of sampling strategies at larger model scales remain limited. This study extends evaluation beyond earlier scale and task settings while focusing on broad multilingual pretraining.
- Multilingual models face a capacity bottleneck because many languages compete for limited model capacity.
- This study systematically explores balancing strategies above 1 billion parameters and targets general-purpose encoder-decoder models.
- Prior translation studies often assumed English was always the source or target language, which may favor distributions covering more English.
- The paper focuses on approximately one hundred languages occurring with some frequency in CommonCrawl, while noting that other methods may be needed for thousands of languages.
3 SAMPLING METHODS
Temperature sampling smooths empirical language frequencies through a temperature parameter, but may not simultaneously balance higher- and mid-resource languages while limiting tail-language repetition. UniMax instead seeks uniform allocation under a per-language epoch cap.
- Temperature-based sampling: Temperature sampling raises each language’s empirical probability to the inverse temperature and renormalizes the result.
- Temperature-based sampling: Higher temperature produces a flatter distribution, but no temperature is guaranteed to balance high- and mid-resource languages while avoiding tail-language overfitting.
- Related balancing assumptions: Translation balancing studies may favor English coverage because English is typically assumed to be a source or target language.
- UNIMAX sampling: UniMax allocates a character budget as uniformly as possible while limiting every language to N epochs, processing languages from smallest to largest corpus.
- Corpus-size measurement: The method uses character counts to measure sub-corpus size, avoiding dependence on a pre-existing tokenizer for massively multilingual data.
4 EXPERIMENTS
The experiments use cleaner multilingual data, dedicated vocabularies, diverse evaluation tasks, and an encoder-decoder training setup. Vocabulary analyses and benchmark design are intended to assess language balancing across resource levels without privileging English.
- Corpus construction: Only 40% of documents in the original Marathi mC4 bucket were well-formed Marathi, motivating stricter language filtering.
- Corpus construction: Filtering mC4 with a 0.95 language-ID confidence threshold removes 6.1% of documents but only 5.1% of characters.
- Vocabulary: Higher temperature, smaller UniMax budgets, and higher max-epoch thresholds produce vocabularies with more uniform cross-language capacity allocation.
- Evaluation tasks: Evaluation tasks are selected for linguistic diversity, balanced language coverage, limited English bias, and relevance to real language-technology use.
- Evaluation tasks: TyDi QA evaluates nine typologically diverse languages using in-language multitask fine-tuning and per-language validation-selected checkpoints.
- Evaluation tasks: WMT21 evaluation covers translation among six languages in all 30 combinations, using 1,000 fine-tuning examples per language pair to reduce English bias.
- Training setup: The models follow mT5 with an encoder-decoder Transformer and span-corruption pretraining on 101 languages plus six Latin-script variants.
5 RESULTS
Across pretraining-loss and downstream evaluations, UNIMAX reduces the imbalance between high- and low-resource languages and outperforms temperature-based baselines across model scales. Its advantages persist on TyDi QA and WMT21, while longer training and repeat controls clarify when overfitting emerges.
- Pretraining loss: UNIMAX closes the loss gap between high- and low-resource languages without showing overfitting, whereas excessive upsampling produces scale-dependent overfitting.Overfitting is conspicuous at XXL scale and emerges even for Large models after about 300,000 steps under longer training.
- TyDi QA: UNIMAX consistently outperforms the other sampling strategies on average TyDi QA across the evaluated model scales.The comparison covers three model sizes; UNIMAX also outperforms τ = 3.33 on both higher- and lower-resource languages, with only the stated high-resource exception against τ = 1 at large scales.
- Per-language analysis: Allocating 47.7% of training to English with τ = 1 yields only 1% better performance than UNIMAX, which allocates 1% to English.UNIMAX also outperforms τ = 3.33 on Swahili despite seeing fewer Swahili examples during training.
- WMT21: UNIMAX outperforms temperature sampling on WMT21 across all model sizes, with benefits spread across the majority of language pairs.The improvement is more pronounced when the target language is non-English, and additional benchmarks show similar trends.
- Further ablations: At a 4× larger character budget, UNIMAX still outperforms temperature sampling, while disallowing repeats entirely gives the best TyDi QA performance among N ∈ {1, 5, 10}.The effect of the max-epoch setting is small, and its optimal value likely depends on the character budget.
6 UMT5 MODELS
The authors train umT5 models with an updated mC4 corpus and UniMax sampling, closely matching mT5’s setup for comparison. umT5 outperforms mT5 on most tasks across model sizes, especially at the largest size.
- The authors train a suite of umT5 models over a trillion tokens using UniMax sampling.The final models also use an updated mC4 corpus covering documents crawled through August 2022.
- The refreshed mC4 corpus raises language-detection confidence from 0.7 to 0.95 and softens bad-word filtering.The changes aim to improve detection accuracy, reduce documents with little or no natural language, and ensure minimal exposure to any term.
- umT5 outperforms mT5 on most tasks across all model sizes, particularly at the largest size.The comparison closely follows mT5’s training setup and evaluation tasks; an appendix isolates the data-refresh effect.
7 CONCLUSION
The conclusion presents UniMax as a near-uniform sampling strategy that avoids harmful repetition and performs well across benchmarks and model scales. The study also identifies limits in its evaluation scope and in treating languages as the sole unit of normalization.
- UniMax approaches uniform language coverage while avoiding harmful repetition and performs well across benchmarks and model scales up to 13 billion parameters.The authors expect its repetition guarantees to remain relevant at larger scales, where overfitting and memorization may become more severe.
- The study focuses on a single pretraining paradigm, leaving encoder-only, decoder-only, parallel-data, and language-specific-parameter settings for future evaluation.The authors explicitly identify these settings as requiring further testing.
- Uniform language distributions are presented as potentially more equitable, but the paper leaves open whether language is the right normalization unit.The authors suggest exploring demographic utility, which treats each speaker as equal when balancing languages.
- The released Large model is excluded because umT5-Large showed pretraining instability and underperformed umT5-Base on all metrics.
- Future sampling strategies should consider both linguistic and demographic utility to support stronger-performing and more equitable pretrained models.
A VOCABULARY ANALYSIS
The vocabulary analysis shows that sampling choices reshape script allocation, token lengths, and language representation. The mT5 distribution also exhibits large disparities in representation ratios, especially disadvantaging populous non-European languages.
- Higher temperature, smaller UniMax budgets, and higher max-epoch thresholds allocate vocabulary capacity more uniformly across languages.These settings reduce tokens from higher-resource scripts and increase tokens from lower-resource languages.
- UniMax vocabularies resemble τ = 3.33, while larger pretraining budgets shift allocation toward high-resource languages and larger max-epoch thresholds shift it toward low-resource languages.
- Higher temperature, smaller UniMax budgets, and higher UniMax max-epoch thresholds produce shorter tokens.When a few high-resource languages dominate the corpus, the optimizer instead allocates capacity to rare words from those languages.
- Compared with mT5, the authors’ vocabularies contain longer tokens, more whitespace-marked tokens, and fewer punctuation-containing tokens.The authors attribute these differences partly to mT5’s whitespace-splitting pre-tokenizer versus training on raw text.
- mT5 language representation ratios range from 7× underrepresented to over 900× overrepresented, with European languages generally more represented.The ratios compare each language’s training-use rate with its native-speaker rate, using Wikipedia speaker counts.
- The five most underrepresented mT5 languages are Chinese, Punjabi, Yoruba, Bengali, and Hindi, all Asian or African languages with at least 50 million native speakers.
C ADDITIONAL TRAINING DETAILS
The models use mT5-like architectures with a 256,000-subword vocabulary and byte-level fallback, while training and fine-tuning use specified Adafactor schedules and loss normalization.
- The models share mT5 architectures except that relative position embeddings are not shared across layers.All models use 256,000 subwords and byte-level fallback for unknown tokens.
- Pretraining uses Adafactor with learning rate 0.01 for the first 10,000 steps, followed by inverse square root decay.Fine-tuning uses Adafactor with a constant learning rate of 5e−5.
- The effective loss normalization uses the number of real target tokens rather than mT5’s loss normalization factor.
- The Adafactor states use an unfactorized second moment and momentum, unlike the T5 and mT5 studies.
D ADDITIONAL BENCHMARKS
Across three additional multilingual benchmarks, UNIMAX performs best overall, with temperature sampling at τ = 3.33 a close second.
- UNIMAX performs best overall across XNLI zero-shot, XNLI translate-train, and XQuAD.
E ABLATION ON MC4 REFRESH
The mC4 refresh was evaluated separately from UNIMAX sampling, showing that the refresh alone provides only a small XNLI gain while offering a potentially useful up-to-date corpus. The evaluation also documents corpus coverage and benchmark layouts used in the ablation.
- Refresh ablation: The refresh ablation compares two 100,000-step UNIMAX models that differ only in using original versus refreshed mC4 data.
- Corpus coverage: Languages absent from mT5 training have a representation ratio of zero; Odia is given as an example despite having 33 million native speakers.
- Evaluation setup: XNLI reports average per-language accuracy, while XQuAD reports average per-language exact match and F1.
- Refresh ablation: +0.1 XNLI gain from the data refresh alone was small, compared with UNIMAX's +1.0 gain over mT5.The authors attribute the larger improvement primarily to improved language sampling.