DEPT: Decoupled Embeddings for Pre-training Language Models
Authors: Alex Iacob
Lorenzo Sani
Meghdad Kurmanji
William F. Shen
Xinchi Qiu
Dongqi Cai
Yan Gao
Nicholas Donald Lane
Venue: ICLR 2025 Oral (2025)
Keywords: Decentralized Training
Federated Learning
Multi-domain Training
Multilingual Training
Primary area: foundation or frontier models, including LLMs
TL;DR: We propose DEPT, a pre-training framework that decouples embedding layers from the transformer body, enabling robust training on heterogeneous data, improving generalization, and reducing memory footprint.

Abstract: Language Model pre-training uses broad data mixtures to enhance performance across domains and languages. However, training on such heterogeneous text corpora requires extensive and expensive efforts. Since these data sources vary significantly in lexical, syntactic, and semantic aspects, they cause negative interference or the ``curse of multilinguality''. To address these challenges we propose a communication-efficient pre-training framework, DEPT. Our method decouples embeddings from the transformer body while simultaneously training the latter on multiple data sources without requiring a shared vocabulary. DEPT can: (1) train robustly and effectively under significant data heterogeneity, (2) minimize token embedding parameters to only what the data source vocabulary requires, while cutting communication costs in direct proportion to both the communication frequency and the reduction in parameters, (3) enhance transformer body plasticity and generalization, improving both average perplexity (up to 20%) and downstream task performance, and (4) enable training with custom optimized vocabularies per data source. We demonstrate DEPT's potential via the first vocabulary-agnostic federated pre-training of billion-scale models, reducing communication costs by orders of magnitude and embedding memory by 4-5x.

Published as a conference paper at ICLR 2025
DEPT: DECOUPLED EMBEDDINGS FOR PRE-TRAINING LANGUAGE MODELS
Alex Iacob +,1,2,
Lorenzo Sa
Meghdad Kurmanji'
William F. Shen 1,4
Xinchi Qiu
Dongqi C,
Yan Gao 1.2, *0
Nicholas D. Lane 1.2,
ABSTRACT
Language Model pre-training uses broad data mixtures enhance performance across domains and languages. However, training on such heterogeneous text corpora requires extensive and expensive efforts. Since these data sources vary significantly in lexical, syntactic, and semantic aspects, they cause negative inter- ference or the "curse of multilinguality", To address these challenges we propose pre- pre-training framework, DEPT. Our method decouples embeddings from the transformer body while simultaneouslyy training the latter on multiple data sources without requiring shared vocabulary. DEPT (1) train robustly and effectively under significant data heterogeneity, (2) minimize token embedding parameters to only what the data source vocabulary requires, while cutting communication costs in direct proportion to both the communication fre- quency and the reduction in parameters, (3) enhance transformer body plasticity and generalization, improving both average perplexity (up to 20%) and down- stream task performance, and (4) enable training with custom optimized vocabu- laries per data source. We demonstrate DEPT's potential via the first vocabulary- agnostic federated pre-training of billion-scale models, reducing communication costs by orders of magnitude and embedding memory by 4 - 5x.
communication-efficii
INTRODUCTION
Language models (LMs) rely on sizable pre-training datasets to generalize across tasks RRafford
al. 2019 Brown al. 2020), and languages (Pires et al. 2019 Artetxe et al. 2020 Zhao al.,
2024 More data boosts generalization and language acquisition (Hoffmann et al. 2022). However,
scaling data creates heterogeneous mix of data sources different domains and languages that challenges LMs. Issues like Negative interference (Wang al.] 2020). where diverse sources com- pete for capacity, and the Curse of Multilinguality (Conneau et al.) 2020), where adding languages yields diminishing returns, especially on low-resoures languages (Magueresse et al. 2020), persist.
Existing methods for pre-training on heterogencous data are costly and complex. Multilingual mod- els like BERT (Devlin et al. 2019), XLM (Conneau et al. 2020), and mT5 (Xue et al. 2021) require temperature-tuning of language sampling ratios for each model-tokenizer pair, involving expensive model selection optimize perplexity (Conneau al. 2020). Large Language Models (LLMs) such as LLaMA handle heterogeneous data with intensive sanguage-specifie heuristics and model- based filters" (Dubey al. 2024). However, these methods still face challenges such as vocabulary dilution (Rust et al. 2021) and sub-optimal cross-lingual/domain performance (Chang et al. 2023a).
This paper proposes communication-efficii pre-training pipeline to address heterogeneous data challenges. Observing that custom vocabularies boost performance across languages (Rust et al., )021 and domains (McLeish et al. 2024), we propose partially fully decoupling the embed- ding space from transformer bodies. This approach optimizes embeddings for specific data sources while the transformer learns abstract representations. We introduce Decoupled Embeddings for Pre- Training (DEP T) in three variants, GLOB, TRIM, and SPEC (see Fig. each increasingly leveraging
"Corresponding author: Alex Iacoh aa 300cam. uc.uk Equal contribution; 'University of Cam- bridge; FFlower Labs; 'Beijing University of Posts and Telecommumications.

---

Published as a conference paper at ICLR 2025
RB
abc
Figure I: Pipeline for DEPT variants: TRIM (top-right), GLOB (bottom-left), SPEC (bottom--ighh with the STANDARD approach (top-left). The numbered pipeline steps proceed follows: (1) text corpora are processed into a vocabulary and tokenizer (global for STANDARD, GLOB, and TRIM; global personalized for SPEC); (2) corpora tokenized into a pre-tokenized dataset; (3) WORKERS train the model on their pre-tokenized data; (4) partial training results are collected; (5) results are aggregated; (6) the new model is to WORKERS. Steps 3-6 repeat convergence.
specialized representations to allow pre-training with distinct domains//anguages, embedding matri- ces, and vocabularies. For example, our SPEC variant scales the vocabulary size linearly with the number of data sources without increasing memory requirements.
DEPT enables pre-training on heterogeneous data sources with unique vocabularies and linguistic features. In the DEPT pipeline, data sources are isolated as silos, akin to clients in cross-silo Feder- ated Learning (FL) (McMahan et al. (2017)). DEPT trains on each silo and aggregates contributions like FL clients. This work examines whether an LM can converge data mixtures without shared (1) output vocabulary, (2) embedding matrices, or (3) tokenization.
In summary, work brings the following scientific contributions:
1. DEPT offers a solution to train an effective transformer body without shared global embed- dings, avoiding the time, electricity, and carbon-intensive HPO tuning.
2. DEPT reduces the memory requirements of models by V(())
where [D& is the average data source's vocabulary size, |V] the global vocabulary size, and dmedaa the embed- ding dimension. For multilingual models, this can save up to 80% of the embedding- matrix size, reducing 409M parameters for our cillion-scale multilingual model.
3. DEPT-based transformer bodies show better generalization. achieving lower validation per- plexities, with improvements upward 15.3 20% average perplexity. DEPT models also excel in model plasticity, quickly adapting to new languages/domains. Finally. DEPT improves downstream fine-tuning performance on Natural Language Understanding tasks. in distributed settings, reducing communication costs
4. DEPT is communication-effiicc. compared to standard distributed data parallelism (Zhao et al., 2023) proportionally to its communication frequency. Compared
coamumiiation-efffceen SGD (Stich) 2019), it obtains further reductions proportional to the size of the model embeddings. Additionally, DEPT enables vocabulary-agnostic federated pre-training for the first time.
2 DECOUPLED EMBEDDINGS FOR PRE-TRAINING (DEPT)
Prior work attributes the Curse of Multilinguality to capacity contention, vocabulary dilution (Con- neau et al.) 2020), and suboptimal tokenization (Rust et al. (02T). These issues affect embed- (Xu et al. 2024). For instance,
dings-even even though the transformer body while English may need 150 000 tokens (Tao et al. 2024), multilingual models allocate 250 000 tokens across hundreds of languages, leading to dilution, contention, and under-represennaaion (Ma-) gueresse et al. 2020). We propose decoupling embeddings during training to enable custom param- eters that reduce contention and vocabularies that avoid dilution and suboptimal tokenization.
cewaabulrr-indepeeddde
We argue that training the transformer body without shared embeddings is feasible. Our intuition is based on evidence that: (a) transformers adapt to new languages by re-learning embeddings (Artetxe et al. 2020); (b) syntactic similarity matters more than subword sharing for performance (Pires et al.) 2019) ; and (c) periodically re- re-initializing embeddings enhances plasticity (Chen et al.) (023). This suggests that transformer body performance is partly embedding--inependdee allowing decoupling.
2

---

Published as a conference paper at ICLR 2025
Algorithm Decoupled Embedding for Pre-Training (DEPT) variants: GLOB TR
SPEC
Require: St set K data sources, T: number rounds
Require: 0o: initial transformer blocks, Φo, $20: optional token/positional embeddings
Require: (D=) source-specific datasets, [V& } source-specific vocabularies
Require: InnerOPT: inner optimizer, OuterCPI: outer optimizer, e.g., AdamW and FedAvg
1: for each update round 1, .......'' do
2: Randomly select subset S; S of data sources for round
3: for each data source A E St in parallel do
4:
0., ob w < InnerOP (0,- O 1 D2)
b
GLOB: Global embeddings
5:
@1-1 1v Trim(0--1.1) V.) 8t delie the + InnerOPI -1-1 11-1, Da)
TRIM:
Trim global token embeddings
P TRIM
7:
0t oit wh InnerOPT(00- vi D2)
SPEC: specialized embeddings
8:
Ad, + 0º 0:-1
Compute parameter update
9:
AØi oi Φ1-1
GLOB:
Compute global token embedding update
AΦc v. + 01l12 01-1
TRIM:
Compute Trimmed embeddings update
11:
Aufi + wi 1/--1
.OB
TR MM global positional embedding update
12:
OuteroPI(@I-I,, Aet 1
Apply the updates for the transformer body
13:
Φ OuterOPI((r-11 , (NES.)
GLOB:
Apply token updates
14:
OuterOPT((t-111 1/ )NES.)
TRIM:
Apply token updates
15:
de < OuterOPT (1. keS.)
GLOB TRIM:
Apply position updates
16:
return 0r, dr, Uhr
Our method, DEPT, achieves this decoupling by: (1) tokenizing data sources independently. using global or custom vocabulary; (2) randomly initializing LM parameters; and (3) training iteratively over random source subsets (see Section 2. This contrasts with standard pre-training, which uses shared embeddings and draws random samples from distribution of all sources.
2.1
METHOP
Akin federated and meta-learning, DEFT optimizes global parameter 9 (the transformer body) along with optional embeddings Φ, & across data sources S. It trains iteratively by selecting a subset S, each round For each data source (k St), DEPT independently performs inner-loop optimization (InnerOPT, e.g., SGD) and then aggregates the transformer bodies using an outer-loop optimizer (OuterOPT, c.g., FedAvg). We present three variants for managing o and vb, offering progressively stronger specialization, and compare them in Section 2.4
CLOB
Shared Embeddings: Based on FedAvg-like methods, GLOB sends a global transformer and embeddings each data source, which then trains locally. The updated models aggregated via OuterOPT making GLOB suitable for federated and centralized settings.
TRIM
Partially Each data source gets global transformer and embeddings but trims the token embeddings its local vocabulary VE. reducing the input/output space. During OuterOPT aggregation, trimmed embeddings are projected to the global vocabulary.
SPEC Fully-decoupled: Each data source gets global transformer and, when first sampled, randomly initializes specialized token/position embeddings. These remain local (never aggregated), supporting any vocabulary, including those from specialized tokenizers.
DEPT replaces the standard pre-training pipeline (Fig. [ for broad pre-training before adapta- tion (Dubey al. (2024). Algorithm runs in parallel, scales with hardware, and reduces com- munication. Reduced communication makes ideal for low-bandwidth settings like cross-silo FL.
2.2 TRIMMED EMBEDDING AGGREGATION (TRIM)
For data source k. trimmed embeddings de E R [Va]/aammdel are derived from global ones @ E R xdmnddl as Φ1. = Igo. where [V] is the global vocabulary size, [& the source-specific size, and dmadel the embedding dimension. The indicator function
3

---

Published as a conference paper at ICLR 2025
[) j-th token in V corresponds to thei-th local token in Vis] selects tokens from Φ. Af- ter InnerOPT we create Φx [dllixdmade, using zero-padding for tokens in V { Ve, and use I ERIVAIA to project dh back, dy. I oz oz. Aggregation (OuterOPT) then applied to {ou with zero-padding ignored to avoid interference between tokens not shared across sources.
2.3 POSITIONAL EMBEDDING SPECIALIZATION (SPEC)
Unlike other variants, SPEC specializes both token embeddings & and positional embeddings vh, as evidence shows syntactic order-dependent properties matter more than subword sharing (Pires al.) 2019). Thus, SPRC is agnostic to vocabulary and sequence length, enabling federated learning without shared tokenization. Without positional specialization, SPEC resembles TRIM, but with the embedding matrix split across sources and disjoint vocabularies {Ve k= such that V U Vk.
2.4 VARIANT CHARACTERISTICS
Table 1: Memory and communication costs of DEPT, where: M the number of model parame- ters; [V] is the global vocabulary size; [& is the mean data source vocabulary size; dmodel the embedding dimension; Noccl N/T is the number of local steps done per iteration for a total num- ber steps N; C is the sequence length. GLOB reduces comms by only communicating every Niocal steps while TRIM also reduces embedding size. SPEC brings further reductions over TRIM by not sharing token or position embeddings. The standard baseline is assumed be distributed training with per-step synchronization. Concrete numbers for our models (see Table are shown in Table
Method
Memory Cost
Per-step Comms Cost
Vocab Agnostic
STD
((M)
O(M)
GLOB
((M)
0
x
o( OMM- MM (1V] Didanuiel
TRIM
O(M - (1V) - (.TDdadd))
Nincal
x
0 M- MM-((2)+-C)dmmmdd
SPEC
O(M ([V] Pa(immdd))
Niocal
In most scenarios, practitioners can deploy any of our proposals, obtaining reduced communication and memory costs as shown in Table However, some settings are appropriate for given variant.
GLOB resembles a standard pre-training pipeline. Although does not explicitly decouple embed- dings from the transformer, they decouple over the course of inner-loop iteration since only local tokens influence them. As a communication-effcian costs compared distributed algorithms such DDP (Li et al., (020) or FSDP (Rajbhandari al.) 2020), which synchronize gradients at every step. However, constructing a global vocabulary re- quires sufficient knowledge of the dataset and may risk vocabulary dilution and capacity contention.
form of SGD, GLOB reduces communication
TRIM shares the same assumptions as GLOB and can be deployed similarly. It further reduces mem- ory requirements for embeddings match the data source's needs (dmode) VR), also lowering communication costs. These savings are substantial for multilingual models with large vocabular- iesiushio et al. 2023). for instance, mT5 and mBART (Xue al. 202T, Lewis al., 2020) allocate 40% 80% of parameters to embeddings. Since our models use ried weights (Inan et al. 2017), TRIM their output space, unlike GLOB, bringing a slight impact to perplexity.
SPEC enables pre-training across data sources without shared vocabulary, providing TRIM's ben- efits plus local specialization. Communication costs are minimized by transferring only the trans- former body to the outer optimizer and decoupling embeddings, enabling vocabulary-agnose ing. This makes SPEC ideal for training a transformer body with unknown or private data. To en- able inference, SPEC requires global embedding matrix. While several methods exist (Section 6. and Appendix F). we use the straightforward approach of multi malti-phase adaptive pre-training (Gu- Turangan et al. 2020), continued pre- pre-training with randomly initialized matrix. This approach follows other techniques for enhancing model capabilities, e.g., long-context pre- traininii stages (De- vlin al. 2019 Dubey et al. 2024) and domain adaptation (Gururangan al. 2020).
train-
4

---

Published as a conference paper at ICLR 2025
3 EXPERIMENTAL DESIGN
We propose DEPT as an efficient alternative to standard pre-training to address the Curse of Mul- filinguality and Negative interference. In this section, we conduct experiments to evaluate DEPT's performance, focusing on the following research questions:
RQ1 Does DEPT allow us to increase the number of training tokens from heterogeneous data?
RQ2 Does DEPT improve efficiency, in terms of memory and communication costs?
RQ3 Does DEPT improve zero-shot generalization
out-of-distribution data?
RQ4 Does DEPT improve model plasticity when learning new distributions?
3.1
EXPERIMENTAL SETUP
For experiments, train decoder-only transformers currently the most relevant architec- tures Bures-ranging from 125M to 1.3B parameters with 12 to 24 blocks (Tables 2 and (}. We use parameter averaging (McMahan et al. 2017a Stich (019)
OuterOpt. optimizer, and AdamW (Loshchilov & Hutter 2019) for InnerOpt. Full experimental details on our architecture,
training hyperparameters (Tables and 8 dataset, and baseline implementation are in Appendix A
3.2 MULTI-DOMAIN AND MULTILINGUAL METHODOLOGY
To evaluate DEPT on multi-domain data, we use The Pile (Gao et al. 2021 ). which in- cludes 22 subsets. We select 16 non-copyrighted subsets our K data sources in Algorithm T GitHub (GH), DeepMind Mathematics (DM). Wikipedia (WK), Common Crawl (CC), PubMed Ab-
stracts (PA), PubMed Central (PC), USPTO Backgrounds (UB), NIH Exporter (NH), FreeLaw (FL), Enron Emails (EE), EuroParl (EP), Stack Exchange (SE), Philosophy Papers (PP), ArXiv (AX), Project Gutenberg (GU), and Hacker News (HN). Ubuntu IRC (UI) the diistribiittiit dataset.
For multilingual data, we use MC4 (Xue et al. 2021) with mix of high, medium, and low-resource languages: English (EN), Italian (IT), and Chinese (ZH) as high-resource; Serbian (SR) and Malay (MS) as medium-resource; and Swahili (SW), Urdu (UR), and Latin (I.A) as low-resource. Following (Rust et al. 202T), we train unigram SentenceP ece (Kudo Richardson, 2018) tokenizers with a 50 257 vocabulary per data source. SPEC variants with optimized per-source vocabularies have the OPT suffix; otherwise, they use a global vocabulary with specialized embeddings.
3.3 BASELINES
We compare DEPT with standard pre-training methods from prior works (Conneau et al. 2020). General distributed SGD methods (Li et al. 2020; Rajbhandari et al. 2020). which synchronize gradients each step and sample from all data sources simultaneously. are labeled STD. For multilingual data, we apply temperature-weightee sampling (Devlin et al. 2019) with т = 0.3, denoted as STD (7 0.3), as well uniform, STD (T 0), and proportional, STD (r 1), sampling For multi-domain data, we use uniform and proportional sampling. Given our data sources random sampling (Algorithm[]) [ baselines with uniform sampling are closest to DEPT.
Additionally, we compare against the "pre-training with active forgetting" (ACT) method (Chen al. 2023), which enhances plasticity and generalization by periodically randomly resetting em- beddings. While Chen et al. (2023) transfer monolingual models between languages, we only utilize their pre-training phase due our different settings. Like SPEC, ACT does not produce fully trained embedding matrix and we employ the same multi- multi-phase adaptive pre-training to create new embedding matrix from a random initialization. Despite this similarity, SPEC significantly more compute efficient than ACT, as it avoids extensive retraining of embeddings. Full details for how we implemented and adapted ACT can be found in Appendix A.1.3
3.4 METRICS
The key characteristics for multi-domain and multilingual pre-training are model generalization and plasticity. Generalization refers to the model's ability to perform well on distributto
0.3 was tuned and found effective in Devlin al. 2019;; Conneau al. (202)) Xue al. 2021
5

---

Published as a conference paper at ICLR 2025
(OOD) data, whether in-domain out-of-domain. We assess in-domain generalization by evaluat- ing the perplexity of a model on the test set of each training data source, while OOD generalization is evaluated with unseen datasets. Furthermore,
evaluate DEP T's efficacy in building foundation models through downstream tasks: Natural Language Inference via MNLI (Williams et al. 2018),
Question Answering via RACE (Lai et al.; 2017), Sentence Similarity via STSB (Cer et al. 2017
and Sentence Classification via SST 2 Socher et al. 2013) Since we use decoder-only models below the model-size threshold for in-context learning abilities (Brown et al. 2020), we follow Rad- ford et al. (2118) for fine-tuning. The evaluation metrics are accuracy (MNLI, RACE, SST-2) and Pearson correlation (STSB). The full details are in Appendix[E]
Plasticity refers to the model's ability to quickly and effectively adapt to new domain, either to reach target performance with minimal steps or to achieve the highest possible performance. We evaluate the plasticity of DEPT models by training them on new data, such as different domain or language, as well as the most heterogeneous subset of the training data, determined by the size of its local vocabulary within the shared global vocabulary (see Appendix A.2
We assess training robustness and stability using the L2 norm of model parameters and activations. Model divergence in LLMs, noted by the OPT (Zhang et al. (022) and PaLM (Chowdhery al.) 2023) teams, correlates with rapid increases in activation norms, a trend also observed in vision transformers (Dehghani al. 2023). While more common at large scales, this issue arise in smaller transformers depending on learning rate suitability ( (Wortsman et al. 2004), which, like hatch size, is influenced by the gradient noise scale for a given data distribution (McCandlish al.. 2018). Notably, all performance comparisons use optimized baseline hyperparameters (see Appendix A).
3.5 CONTINUED PRE-TRAINING AND EVALUATION
Once pre-training is complete, some methods, including SPEC and ACT, lack global embed- ding, while others. such as STANDARD pre-training, GLOB, and TRIM, include one. For ACT and SPEC (see Section 3.5). we enable global (shared) embedding through multi-phase adaptive pre- training (Gururangan et al. 2020). This involves broad DEPT pre-training (Algorithm [ followed by continued pre-training on another 15-19% of the total steps on non-private dataset using a ran- domly initialized embedding matrix with a global vocabulary tailored to the specific corpus. For this phase, we use the tokenizer of Black et al. (2022) for English data and Xue et al. (202I) for multilingual data. These extra steps are applied to all models for fair comparison. While random initialization reveals the quality of the transformer body for all DEPT variants, we are also concerned with the independent effectiveness of GLOB and TRIM in building high-quality global embeddings compared to STANDARD methods. We perform the same 15-19% extra steps for this comparison, starting from pre-trained embeddings.
Unlike pre-training, this stage requires a sampling strategy. Since The Pile is curated for propor- tional sampling (Gao et al. 2021), we use it for multi-domain continued pre-training, while uniform sampling is applied to multilingual data to support low-resource languages.
4 RESULTS
Our results show that DEPT improves transformer body generalization (Tables 3 and 4) enhanc- ing robustness (Fig. 2). plasticity (Fig. 3 and downstream performance (Table while bringing
communication and memory costs reduction (Table
4.1 DEPT Is ROBUST To DATA HETEROGENEITY (RQ1)
Our experiments demonstrate DEPT's robustness to multilingual and multi-domain data heterogene- ity. As shown in Zig.Z DEPT resists activation divergence and model norm increases, which can halt perplexity improvements or cause divergence (Zhang et al. 2022: Chowdhery et al. 2023; Worts- the baselines, models trained with all DEPT variants maintain lower activation norms due to the regularization effects of Out (Algorithm T.. Learning rates for baselines are reduced for later comparisons ensure convergence.
man et al.) 2024). When using the same local hyperparameters

---

Published as a conference paper at ICLR 2025
Activationn (U30
Norm Model ente
STD (f
(F- DEPT-AVG
STD (r -1)
2.5
DEPT-AYG
0.0
6:00
8000
1000
Sequential Steps
Sequential Steqe
(a) The Pilc pre-train, activation norms, 24-block
(b) The Pilc pre-train, parameter norms, 24-block
Figure 2: Activations and model norms of STANDARD (STD) training versus DEPT (avg + min/max) for a 350M model trained with identical local hyperparameters-ppa prior to adjusting STD (7 = 0) and STD (T = (uniform and proportional sampling) to lower learning rate. The OuterOpt DEPT
introduces regularization effects due to noise- injection (Lin et al. 2020). meta-learning (Nichol
al. 2018) characteristics, which constrain these sources (Zhang et al. 2022) of model divergence.
Table 2: Practical memory and communication costs for DEPT, where the total number of steps is N = Mocaal with the total number of iterations, and V& as the average vocabulary size across data sources. Standard pre-training requires a full in-memory embedding matrix for the global vo- cabulary while synchronizing gradients every step rather than every Niocal steps. All DEPT variants yield communication savings, with GLOB as the baseline. TRIM provides additional savings pro- portional to the gap between global and local vocabulary sizes, while SPEC further reduces costs by never communicating embeddings. For the full comparison, see Table
Type
ABlocks
Method
Niows
T
[V. 4
[V x dmnddl
My (4)
Per-step Comms Cust (4)
Multilnnuul
5 10'
250
192M
278M (1x)
278M (1x)
Multilinuull
G10B
278M (1x)
0.58M (0.002x)
Multiinguul
12
TRIM
500
10
216 135 d: 27 160
166M
262M (0.92x)
0.5M (0.002x)
Multilingual
SPBC
216
160M
252M (0.922x
0.17M (0.0006
Multilingual
SPEC-OPT
500
10
50/257 #: 0
38.6M
125M (0.15%)
0.17M (0.0006 x)
Multilingual (1BB
24
STD SPEC-OFT
7 10'
250112
512 2M
1,718 (1x)
1,71B (1x)
Multilingual (1B)
502257
(0.76x)
2.4M (0.001x)
4.2 DEPT IMPROVES TRAINING EFFICIENCY (RQ2)
Tables 1 and 2 show that DEPT significantly reduces average GPU memory and per-step communi- cation costs compared to DDP. The 500x memory cost reduction from GLOB matches that of Local SGD, as it synchronizes gradients only every Niocal steps, allowing GPUs to operate independently in between. TRIM further improves memory and communication costs by reducing vocabulary size, shrinking the global embedding matrix by 8% to 32% for multilingual data and by 2% to 78% for The Pile, with the largest reduction (78%) achieved for the mathematics subset (see Ap- pendix A.2 for precise vocab sizes). SPEC eliminates embedding-related communication, reducing costs by an additional 13% to 30% for multi-domain data and 34% for multilingual data. Finally, DEPT enables efficient training of billion-scale models (Fig. 4 multilingual data, achieving 714x reduction in communication costs (Table and 24% reduction in memory costs.
4.3 DEPT IMPROVES ZERO-SHOT GENERALIZATION (RQ3)
We show that DEPT variants significantly enhance transformer body generalization, outperforming STANDARD pre-training and active-forgetting (ACT) in: (a) perplexity on pre-training validation data, (b) perplexity on OOD validation data, and (c) downstream fine-tuning on MNLI, RACE, STSB. As detailed in Section 3.5, DEPT serves as the first stage of a multiphase adaptive pre- training pipeline, followed by continued pre-training on a non-private dataset. With pre-training data coalesced as in STANDARD training, Our results reflect performance after this phase is applied to baselines as well, ensuring embeddings process the same number of tokens. To gauge tokenizer effectiveness on dataset, we report the unigram cross-entropy (UNIGRAM-CE) of the unigram model defined by the token frequencies. with higher values indicating a harder-to-mode distribu- tion (Tao et al. 20e4) Appendix 1 l. Overall, DEPT variants win 82.2% 62 of our main comparisons across The Pile, MC4 and downstream tasks, producing generalizable and perfor- mant transformer bodies.
7

---

Published as a conference paper at ICLR 2025
4.3.1 TRANSFORMER BODY GENERALIZATION
Table 3: Validation perplexity (L) for 24-block models trained on The Pile after continued pre- training with proportional sampling from randomly-initializzn embeddings shows that DEPT improves performance across all data sources, outperforming baselines by 15.3% on average. SPEC-OPT, using an optimized vocabulary, outperforms GLOB on high UNTGRAM- CF sources.
Name
DM
EN
EF
F1.
GH
CC
PA
SR
PP
WK
AX (7.7)
UB
PC
NH
GU
HN
UI-OOD
AVG
(UNGGRMF-CR
(6.9)
(7.9)
(10)
(7:8)
(7.9)
(1.
(5.2)
(7.1)
(9.1)
(8.2)
(7.8)
(8)
(8.1)
(7.7)
(7-7)
(10)
(8.1)
SVD (+
6.5
44.8
83.5
30.9
8.1
79.6
46.6
25.4
126.6
68.2
14.3
34.1
22.3
58.9
76.3
65.2
163.6
56
20.6
56.2
16.8
39.1
11
39.3
GLOB
26.7
38.2
17.3
5-4
47.7
25.7
14.7
08.3
32.7
9.9
32.2
46.5
94.8
TRIM
4.5
27.3
19.5
18.5
5.6
51.2
27.8
15.4
71.8
35.1
10.3
21.7
14.8
35.1
49.1
42.2
96.77
21..
SPEC
4.8
26.7
36.8
18.2
5.5
50.1
27.1
15.1
600.1
34.2
10.1
21.1
14.5
34.3
48.5
41.7
97.6
32.7
SPEC-OPY
4.7
25.9
36
17.5
5.4
45.3
26.1
14.7
66.
32.8
9.9
20.4
14.1
32.5
47.3
40.5
88.6
81.2
Min Imp (%)
3.7
10.6
20.2
10.1
7.4
8.9
10.3
8.4
11.5
10.3
8.6
8.2
10.6
0.9
11
36.5
10.6
Table 4: Validation perplexity (L) for 12-block models trained on MC4 using continued pre-training with uniform sampling from randomly-initfalized embeddings. DEPT improves transformer per- formance across all languages, averaging a 17.3% gain for pre-train data and 20.8% on OOD sources. SPEC outperforms GLCB on high UNIGRAM OOD data.
In-Distribution
Out-of-Distribution
Name
ZH
UR
MS
IT
SR
LA
EN
SW
Avg (In-D) (9.3)
EL.
HI
DE
Avg (OOD)
(UNIGRAM-CE)
(9.8)
(10.5)
(9.2)
(7.7)
(10.5)
(9)
(7.5)
(10)
(1/1.4)
(13.9)
(9.7)
(12.6)
STD (r 0)
154.8
38.2
96.8
83.8
73.3
63
112.7
62.8
85.7
5660.8
4600.3
1339.2
3866.8
(r 0.3)
129.5
56.3
1083.7
4219.2
1076.3
3097.1
STD (r 1)
84.6
26.8
61.8
55.1
17.1
11.1
77.6
12.4
5-19
33:10.3
25111.
672.5
2175.8
96.1
28.8
71.3
60.4
62.3
44.9
86.6
46.3
00.7
2450.2
2412.5
715.9
1850.5
GLOB
67.7
22.4
53.7
46
38.6 6
33.9
65.4
35.2
45.4
2308.3 3
1676.5
559.5
1514.7
TRIM
67.7
22.8
55.2
47.5
33.7
35.1
67.2
36.3
46.4
2547.7
1911
567.4
1675.4
SPBC
69.5
23
55.4
47.8
40.3
34.7
68.1
36.3
46.9
2232.1
1578.8
544.7
1451.9
Min Tmp (%)
17.8
14
14.5
13.4
14.6
14.6
12.2
14.3
14.4
20.8
15.6
10.8
Max Imp (%)
20
16.4
17.1
16.6
18.1
17.4
15.7
16.9
17.3
8.9
34.6
19
20.8
Tables 3 and E present results where embedding matrices are initialized randomly. DEPT variants significantly outperform all baselines across validation sets for multilingual and multi-domain data sources, including high- and low-resource subsets. Min and max improvements, shown in the last two rows of the tables, compare the worst and best DEPT variants to the best-performing baseline. The best DEPT variant achieves average performance improvement of 17.3% MC4 and 15.3% on The Pile, while even the worst variant shows improvements of 14.4% and 9.7%. respectively. DEPT wins 100% JF H comparisons for The Pile and MC4, respectively. For OOD data, DEPT variants outperform by 10-20% on average for MC4 and 1.5-10.5% on The Pile, despite the high UNIGRAM-CS CE OOD data, which makes it more difficult. This demonstrates that DEPT produces superior transformer bodies with better generalization. Notably, TRIM performs compa- rably to GLOB despite significant reductions parameter counts and communication costs during pre-training, suggesting that out-of- mistakes do not drastically impact performance, For downstream tasks, however, TRIM surpasses GLOB (Table SPEC performs similarly to GLOB and TRIM, even without sharing token embeddings across data sources. The SPEC-OP CPT variant, trained with unique vocabularies and parameters for each The Pile data source, outperforms GLOB on datasets with high UNIGRAM-CE those dissimilar to natural language, such as multilingual EP, math math-heavy DM, code-based GH, and the high-UNIGRAM- tently outperforms OOD datasets with high UNIGRAM- CE. These results hold across model sizes (see Table 12). and across sampling techniques (Table 10
dataset UI. For MC4, SPEC consis-
4.3.2 PRE-TRAINED EMBEDDING GENERALIZATION
Tables and represent cases where the global embedding is initialized using the final global embedding obtained during pre-training, applicable only to the GLOB and TRIM variants. For The Pile (). both variants outperform their standard pre-training counterparts, achieving a 5.5% 12 comparisons. Two of the lost comparisons, the
improvement in average accuracy and winning I7 small subsets EN and EP, are instead won when using uniform sampling (Table TT).
8

---

Published as a conference paper at ICLR 2025
Table 5: Validation perplexity (4) for 24-block models trained on The Pile with continued pre- training using proportional sampling from pre-trained embeddings. DEPT wins 70% 12 comparisons with GLOB consistently outperforming TRIM. In Table B DEPT wins the remaining 5 due to its superior transformer body. Likewise, the EN and EP comparisons are won when using uniform sampling (Table TT) as embeddings become more refined on these smaller datasets.
17
Name
DM (6.9)
GH
SF
WK
GL.
HN
UL-0DD
AVG
(UNGGRAM.C))
(7:9)
(7.8)
(7.0)
(7:9)
(8.2)
(7.7)
(9.1)
(5.2)
(7.7)
(7.8)
(8.1)
(7.7)
(7.7
(10)
13
57.2
90.9
STD
4.5
21.9
12.3
19.7
11.6
24.5
16.2
35.4
31.7
54.3
22.8
GLOB
4.5
18.2
17.9
11.2
22.4
20.6
85.6
28.3
61.2
7AIM
4.6
20.5
13.0
20.2
12
40.9
25.1
8.7
10.6
26.7
38
32.9
56.8
23.7
Min Tmp (%)
-3
18.7
-6.9
3.0
-35
-27
2.7
-34
30.1
-26
2.0
-27
-2
0.00
4.5
-6.4
-18.
Max Imp (%)
-1.2
20.6
-3
0.9
-0.8
6.8
3.4
1.4
8.4
0.9
11
12.3
2.6
8.4
-5
5.5
Table 6: Validation perplexity (L) for 12-block models trained on MC4 using continued pre-training with uniform sampling from pre-trained embeddings. DEPT achieves a 6.4% improvement in average perplexity for in-distribution data but slightly underperforms for OOD data, winning 50% of OOD comparisons. In Table 4 DEPT wins the remaining cases of in-distribution and 33% 3 due to a better transformer body.
In- Hn-Distribution
Out-of-Distribution
Name (UNIGRAM--CE)
711
UR
M8
IT
SR
LA
EN
SW
Avg (In-D)
EI.
HI
DE
Avg (OOD)
(9.8)
(10.5)
(9.2)
(7-7)
(10.5)
(9)
(7.5)
(10)
(9.3)
(14.4)
(13.9)
(9.7)
(12.6)
STD (r ()
57.8
21
46.5
40
33.6
29.4
57.5
30.3
39.5
1698..
1365.7
385.5
1150
STD (r 0.3)
45.5
20.6
41.5
31.7
29.3
46.1
31.1
34.6
1419.4
1087.6
321.9
943
STD (7 1)
44.4
23.9
44.3
25.2
36.5
33.4
38.3
36.4
35.3
1583.6
1299.5
285.5
1056.2
GLOB
15.5
30.1
39.6
39
29.7
32.4
1737.3
823.4
335.1
965.3
TRIM
41.9
16.2
31.3
41.3
40.8
30.8
42
25.6
33.7
1725
855.2
345.6
975.3
Imp (%)
21.1
24.7
28.7
22.44
21.1
Max Imp (%)
9.7
24.4
27.6
-57.4
22.8
1.2
-5.8
18.7
6.4
-215
24.3
-1744
-2.4
Furthermore, DEPT consistently outperforms when starting from random embeddings due to its superior transformer body. Thus, argue that differences in performance com- pared to results in Section 4.3.] are primar- ily driven by variations in embedding sam- pling ratios. For MC4 (Table σ). DEPT wins comparisons for in-distribution data and 3 for OOD data, providing disproportionate ben- efits for the low-resource UR and sw languages. These languages have very high UNI GRAM-CE values, indicating that the global shared tok-
Table 7: The performance on downstream tasks (t), following continued pre-training, shows that DEPT models achieve 3% 7.5% relative im- provements over the baselines, with TRIM deliv- ering the best results. DEPT consistently outper- forms baselines. For the full results see Table 21
Random Inil
Name
RACE (ACC)
(ACI)
STSB (PC)
SST2 (ACC)
0.66
STD (
0.48
0.88
0.73
0.81
ACT
0.45
0.68
0.73
0.80
GLOB8
0.51
0.72
0.78
0.83
enizer, trained with weeigattrre-eiighted sam- pling, underserve them. Switching to propor- tional sampling during continued pre-training improves performance high-resource lan-
TRIN
SPEC
0.52
0.58
0.71
0.71
0.79
0.78
0.83
0,81
0.85
Imp
2.9%
4.6%
5.9%%
-0.7%
Max Imp (%)
5.8%
0.1%
7.5%
4.1%
guages, winning EN. Similarly to The Pile, the other comparisons are all won when starting from random embeddings. Thus, while DEPT may benefit the transformer body, care must be taken to design an appropriate continued pre-training pipeline to effectively fine-tune the embeddings.
4.3.3 DOWNSTREAM GENERALIZATION
Table 7 presents the downstream performance of 24-block DEPT models pre-trained and contin- ued pre-trained (with uniform sampling) on The Pile. DEPT models consistently outperform the baselines, regardless of initialization, with TRIM achieving the best results and SPEC matching GLOB in wins. Despite occasional losses to GLOB in language modeling, we speculate that the re- stricted vocabulary of TRIM forces it to adapt to language shifts, improving generalization, akin to ACT's 's ie-iiiiaaiizaiinn but more effective. While ACT performs better on downstream tasks than on language modeling (Chen et al. 2023), it is outperformed by DEPT. DEPT leverages inherent aggre- gation noise to develop robust parameters without artificial re- re-initialization, ensuring that parameter updates are not discarded and avoiding the waste of compute cycles.

---

Published as a conference paper at ICLR 2025
4.4 DEPT IMPROVES MODEL PLASTICITY (RQ4)
Finally, we investigate how plastic DEPT models are in adapting to either new data source or to the most heterogeneous subset of the pre-training set. Figure 3 the perplexity adaptation plots when starting from a random initialization on the full pre-training set (serving as a baseline), the data source with the smallest vocabulary (SW), or new languages (HI, DE). DEPT variants are always the fastest to adapt to each data source and provide the lowest final perplexity; for the full pre-training set, we use perplexity taken over all language validation sets.
230
STD B-3)
SED
STD
Purplecky
150
ACI DEPT-AYG
Porplectty
ACT DEPTAYG
1000
Sequrttial Stugu
Sequenttaal Steps
(a) MC4-FULL, 12-block.
(b) HI, 12-block.
Figure 3: Adaptation curves starting from randomly initialized matrix. DEPT variants are al- ways stable their convergence, reaching the lowest perplexity for the full dataset and the out-of- distribution language (HT). It is also always the fastest to adapt, full results available in Figure
5 RELATED WORK
Large language models (LLMs) exhibit cross-lingual alignment due to "incidental bilingualism" (Briakou et al. 2023) and cross- cross-lingual data sharing (Choenni et al. (2023). Expanding multilingual
data during pre-training can enhance language diversity (Scao et al. | (022) but often results in uneven
performance due to data imbalance and low-resource degradation (Ding et al.] [222]] Lai et al.)
2023). Supervised parallel data (c.g., XLM (Conncau & Lample| 2019). PaLM2 (Anil et al. (023)). Knowledge Transfer (Zhang et al. 2023 Wang et al. 2003), and Domain Adaptation (Huang al.
2024) face challenges in low-resource settings (Chang et al. 2023b Li et al. 2024), with risks like training instability and catastrophic forgetting (Kirkpatrick et al. (017). This motivates novel pipeline, focusing on language heterogeneity, generalization, and plasticity. Vocabulary construction crucial in multilingual pre-training. Techniques include tokenization with temperature setting (Devlin et al., (019) and language-clustered vocabularies (Chung et al. (020),, though the latter requires predefined clusters. Active forgetting (Chen et al., 2003). a related approach, enhances model plasticity by periodically re-initializing embeddings, casing adaptation to new languages.
CONCLUSION
We investigated pre-training Language Models (LMs) under data heterogeneity, proposing effi cient and robust pipeline, DEPT, which supports training under diverse data sources while mitigating Negative Interference and the Curse of Multilinguality. The core DEPT is decoupling the embed- ding space from the transformer body during pre-training, offered in three variants with varying degrees of separation. Experiments showed that DEPT (1) allows training across heterogeneous data efficiently, (2) reduces the memory footpring of token embedding matrices by 4 - 5x, (3) improves model generalization and plasticity with lower perplexity on validation and out-of-distribution test datasets, and (4) supports custom vocabularies per data source, enabling vocabulary agnostic feder- ated pre-tree pre-training, which we have tested up to billion-scal scale models and intend to push further.
6.1 LIMITATIONS FUTURE WORK
DEPT offers pre-training framework intended to precede further adaptation fine-tuning. How- ever, DEPT models require a final global embedding for practical use. The GLOB and TRIM variants provide this the end of pre-training, while SPEC does not, suggesting future work on embed- ding generation methods, such as zero-shot embedding transfer (Mosin et al. 2023), vocabulary matching (Xu et al. 2024) and model stitching (Moschella et al. 2023).
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
All costs for the computation used for this work was funded by Flower Labs, and the research conducted by team of researchers from Flower Labs and The University of Cambridge. Support for university- baased researchers came from a variety of sources, but in particular, the following funding organizations are acknowledged: the European Research Council (REDIAL), the Royal Academy of Engineering (DANTE), and the Ministry of Education of Romania through the Credit and Scholarship Agency.
REFERENCES
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Ale- man, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. Gpt-4 technical report. arXiv preprint vrxiv:2303.08774, 2023.
Rohan Anil, Andrew M Dai, Orhan Firat, Melvin Johnson, Dmitry Lepikhin, Alexandre Passos, Siamak Shakeri, Emanuel Taropa, Paige Bailey, Zhifeng Chen, al. Palm 2 technical report. arXiv preprint arXiv:2305. 10403, 2023.
Mikel Artetxe, Sebastian Ruder, and Dani Yogatama. On the cross-lingual transferability of mono- lingual representations. AnAC pp. 4623 4637. Association for Computational Linguistics, 2020.
Daniel J. Beutel, Taner Topal, Akhil Mathur, Xinchi Qiu, Javier Fernandez-Marques Yan Gao, Lorenzo Sani, Kwing Hei Li, Titouan Parcollet, Pedro Porto Buarque de Gusmão, and Nicholas D. Lane. Flower: A friendly federated learning research framework. CoRR, abs/2007. 2022.
Stella Biderman, Hailey Schoelkopf, Quentin Gregory Anthony, Herbie Bradley, Kyle O' O'Brien, Eric Hallahan, Mohammad Aflah Khan, Shivanshu Purohit, USVSN Sai Prashanth, Edward Raff, Aviya Skowron, Lintang Sutawika, and Oskar van der Wal. Pythia: A suite for analyzing large language models across training and scaling. In ICML, volume 202 of Proceedings of Machine Learning Research, pp. 2397-2430. PMLR, 2023.
Sid Black, Stella Biderman, Eric Hallahan, Quentin Anthony, Leo Gao, Laurence Golding, Horace He, Connor Leahy, Kyle McDonell, Jason Phang. Michael Pieler, USVSN Sai Prashanth, Shivan- shu Purohit, Laria Reynolds, Jonathan Tow, Ben Wang, and Samuel Weinbach. Gpt-neox- 20b: An open-source autoregressive language model, 2022. URL http:ttpsa///axi.ooooo////////// 06745
Cody Blakeney, Mansheej Paul, Brett W. Larsen, Sean Owen, and Jonathan Frankle. Does your data spark joy? performance gains from domain upsampling the end of training. In First Conference on Language Modeling, 2024. URL htttps://openreviiee
Eleftheria Briakou, Colin Cherry, and George Foster. Searching for needles in haystack: On the role of incidental bilingualism palm's translation capability. arXiv preprint arXiv:2305.10266 2023.
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners, 2020.
Daniel M. Cer, Mona T. Diab, Eneko Agirre. lñigo Lopez Gazpio, and Lucia Specia. Semeval-2017 task 1: Semantic textual similarity multilingual and cross-lingual focused evaluation. CoRR, abs/1708. 708.00055, 2017.
Tyler A. Chang, Catherine Arnett, Zhuowen Tu, and Benjamin K. Bergen. When is multilinguality a curse? language modeling for 250 high- and low-resource languages. CoRR, abs/2311. 1.02005, 2023a.
11

---

Published as a conference paper at ICLR 2025
Tyler A Chang, Catherine Arnett, Zhuowen Tu, and Benjamin K Bergen. When is multilin- guality a curse? language modeling for 250 high-and low-resource languages. arXiv preprint arXiv:2311.09205. 2023b.
Zachary Charles, Nicole Mitchell, Krishna Pillutla, Michael Reneer, and Zachary Garrett. Towards federated foundation models: Scalable dataset pipelines for group-structured learning. In Alice Oh, Tristan Naumann, Amir Globerson, Kate Saenko, Moritz Hardt, and Sergey Levine (eds.), Advances in Neural Information Processing Systems 36: Annual Conference on Neural Informa- tion Processing Systems 2023, NeurIPS 2023, New Orleans, LA, USA, December 10 16, 2023, 2023.
Yihong Chen, Kelly Marchisio, Roberta Raileanu, David Ifeoluwa Adelani, Pontus Lars Erik Saito Stenetorp. Sebastian Riedel, and Mikel Artetxe. Improving language plasticity via pretraining with active forgetting. In NeurlPS, 2023.
Rochelle Choenni, Dan Garrette, and Ekaterina Shutova. How do languages influence each other? studying cross-lingual data sharing during IIm fine-tuning. arXiv preprini htlv:2305.13880 2023.
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, Parker Schuh, Kensen Shi, Sasha Tsvyashchenko, Joshua Maynez, Abhishek Rao, Parker Barnes, Yi Tay, Noam Shazeer, Vinodkumar Prabhakaran, Emily Reif, Nan Du, Ben Hutchinson, Reiner Pope, James Bradbury, Jacob Austin, Michael Isard. Guy Gur-Ari, Pengcheng Yin, Toju Duke, Anselm Lev- skaya, Sanjay Ghemawat, Sunipa Dev, Henryk Michalewski. Xavier Garcia, Vedant Misra, Kevin Robinson, Liam Fedus, Denny Zhou, Daphne Ippolito. David Luan, Hyeontaek Lim, Barret Zoph, Alexander Spiridonov, Ryan Sepassi, David Dohan, Shivani Agrawal, Mark Omernick, Andrew M. Dai, Thanumalayan Sankaranarayana Pillai, Marie Pellat, Aitor Lewkowycz, Erica Moreira, Rewon Child, Oleksandr Polozov, Katherine Lee, Zongwei Zhou, Xuezhi Wang. Bren- nan Saeta, Mark Diaz, Orhan Firat, Michele Catasta, Jason Wei, Kathy Heils-Hellstemm, Douglas Eck, Jeff Dean, Slav Petrov, and Noah Fiedel. Palm: Scaling language modeling with pathways. J. Mach. Learn. Res., 14:240:1-240:: 133, 2023.
Hyung Won Chung, Dan Garrette, Kiat Chuan Tan, and Jason Riesa. Improving multilingual models with language-clustered vocabularies. In EMNLP (1), pp. 4536-4546. Association for Computa- tional Linguistics, 2020.
Alexis Conneau and Guillaume Lample. Cross-lingual language model pretraining. Advances in neural information processing systems, 32, 2019.
Alexis Conneau, Kartikay Khandelwal, Naman Goyal, Vishrav Chaudhary, Guillaume Wenzek, Francisco Guzmán, Edouard Grave, Myle Oit, Luke Zettlemoyer, and Veselin Stoyanov. Un- supervised cross- representation learning at scale. In ACL, pp. 1440-8457 Association for Computational Linguistics, 2020.
Databricks. mosaic research, 2024. URL
mosaic
Mostafa Dehghani, Josip Djolonga, Basil Mustafa, Piotr Padlewski, Jonathan Heek, Justin Gilmer, Andreas Peter Steiner, Mathilde Caron, Robert Geirhos, Ibrahim Alabdulmohsin, Rodolphe Jenatton, Lucas Beyer, Michael Tschannen, Anurag Arab, Xiao Wang, Carlos Riquelme Ruiz, Matthias Minderer, Joan Puigcerver, Utku Evci, Manoj Kumar, Sjoerd van Steenkiste, Gamaleldin Fathy Elsayed, Aravindh Mahendran, Fisher Yu, Avital Oliver, Fantine Huot, Jasmijn Bastings, Mark Collier, Alexey A. Gritsenko, Vighnesh Birodkar, Cristina Nader Vasconcelos, Yi Tay. Thomas Mensink, Alexander Kolesnikov, Filip Pavetic, Dustin Tran, Thomas Kipf. Mario Lucic, Xiaohua Zhai, Daniel Keysers, Jeremiah J. Harmsen, and Neil Houlsby. Scaling vision transformers to billion parameters. In ICML., volume 202 of Proceedings of Machine Learning Research, pp. 7480-7512. PMLR, 2023.
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: pre-training of deep bidirectional transformers for language understanding. In Jill Burstein, Christy Doran, and Thamar Solorio (eds.), Proceedings of the 2019 Conference of the North American Chapter of the
12

---

Published as a conference paper at ICLR 2025
Association for Computational Linguisties: Human Language Technologies, NAACL-HLT 2019, Minneapolis, MN. USA, June 2-7, 2019, Volume (Long and Short Papers), pp. 4171-4186. Association for Computational Linguisties, 2019.
Bosheng Ding, Chengwei Qin, Ruochen Zhao, Tianze Luo, Xinze Li, Guizhen Chen, Wenhan Xia, Junjie Hu, Anh Tuan Luu, and Shafiq Joty. Data augmentation using Imms: Data perspectives, learning paradigms and challenges. arXiv preprint arXiv:243.00299. 2024.
Arthur Douillard, Qixuang Feng, Andrei A. Rusu, Rachita Chhaparia, Yani Donchev, Adhiguna Kuncoro, Marc' Aurelio Ranzato, Arthur Szlam, and Jiajun Shen. Diloco: Distributed low- communication training of language models. CoRR, abs/2311 1111110011,,,, 2023.
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, Anirudh Goyal, Anthony Hartshorn, Aobo Yang, Archi Mitra, Archie Sravankumar, Artem Korenev, Arthur Hinsvark, Arun Rao, Aston Zhang, Aurélien Rodriguez. Austen Gregerson, Ava Spataru, Baptiste Rozière, Bethany Biron, Binh Tang, Bobbie Chern, Charlotte Caucheteux, Chaya Nayak, Chloe Bi, Chris Marra, Chris McConnell, Christian Keller, Christophe Touret, Chunyang Wu, Corinne Wong, Cristian Canton Ferrer, Cyrus Nikolaidis, Damien Allonsius, Daniel Song, Danielle Pintz, Danny Livshits, David Esiobu, Dhruv Choudhary, Dhruv Mahajan, Diego Garcia-Olano, Diego Perino, Dieuwke Hupkes, Egor Lakomkin, Ehab AlBadawy, Elina Lobanova, Emily Dinan, Eric Michael Smith, Filip Radenovic, Frank Zhang, Gabriel Synnaeve, Gabrielle Lee, Georgia Lewis Anderson, Graeme Nail, Grégoire Mialon, Guan Pang, Guillem Cucurell, Hailey Nguyen, Hannah Korevaar, Hu Xu, Hugo Touvron, Iliyan Zarov, Imanol Arrieta Ibarra, Isabel M. Kloumano, Ishan Misra, Ivan Evtimov, Jade Copet, Jaewon Lee, Jan Geffert, Jana Vranes, Jason Park, Jay Mahadeokar, Jeet Shah, Jelmer van der Linde, Jennifer Billock, Jenny Hong, Jenya Lee, Jeremy Fu, Jianfeng Chi, Jianyu Huang, Jiawen Liu, Jie Wang. Jiecao Yu, Joanna Bitton, Joe Spisak, Jongsoo Park, Joseph Rocca, Joshua Johnstun, Joshua Saxe, Junteng Jia, Kalyan Vasuden Alwala, Kartikeya Upasani, Kate Plawiak, Ke Li, Kenneth Heafield, Kevin Stone, and al. The llama 3 herd of models. CoRR, abs/2407.21783. 2024.
Fahim Faisal, Yinkai Wang, and Antonios Anastasopoulos. Dataset geography: Mapping language data language users. In Smaranda Muresan, Preslav Nakov, and Aline Villavicencio (eds.), Proceedings of the 60th Annual Meeting the Association for Computational Linguistics (Volume 1: Long Papers), ACL 2022, Dublin, Ireland, May 22-27, 2022, pp. 3381-3411. Association for Computational Linguistics, 2022.
Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, Shawn Presser, and Connor Leahy. The pile: An 800gb dataset of diverse text for language modeling. CoRR, abs/2101.0000 2021.
Michael M Grynbaum and Ryan Mac. The times sues openai and microsoft over a.i. use of copy- righted work, Dec 2023.
Suchin Gururangan, Ana Marasovic, Swabha Swayamdipta, Kyle Lo, Iz Beltagy, Doug Downey, and Noah A. Smith. Don't I stop pretraining: Adapt language models to domains and tasks. In ACL, pp. 8342-8360. Association for Computational Linguistics, 2020.
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, Tom Hen- nigan, Eric Noland, Katie Millican, George van den Driessche, Bogdan Damoc, Aurelia Guy, Simon Osindero, Karen Simonyan, Erich Elsen, Jack W. Rae, Oriol Vinyals, and Laurent Sifre. Training compute-optimal large language models. CoRR, abs/2203. 15556, 2022.
Kaiyu Huang, Fengran Mo, Hongliang Li, You Li, Yuanchi Zhang, Weijian Yi, Yulong Mao, Jinchen Liu, Yuzhuang Xu, Jinan Xu, et al. A survey on large language models with multilingualism: Recent advances and new frontiers. arXiv preprint arXiv:2405.10936. 2024.
Hakan Inan, Khashayar Khosravi, and Richard Socher. Tying word vectors and word classifiers: A loss framework for language modeling. In ICLR (Poster). OpenReview.net, 2017.
13

---

Published as a conference paper at ICLR 2025
James Kirkpatrick, Razvan Pascanu, Neil Rabinowitz, Joel Veness, Guillaume Desjardins, Andrei A Rusu, Kieran Milan, John Quan, Tiago Ramalho, Agnieszka Grabska Barwinska, et al. Overcom- ing catastrophic forgetting in neural networks. Proceedings of the national academy of sciences, 114(13):352133526 2017.
Taku Kudo and John Richardson. Sentencepiece: A simple and language independent subword tokenizer and detokenizer for neural text processing. In EMNLP (Demonstration)) pp- 66-71. Association for Computational Linguistics, 2018.
Guokun Lai, Qizhe Xie, Hanxiao Liu, Yiming Yang, and Eduard H. Hovy. RACE: large-scale reading comprehension dataset from examinations. In EMNLP, pp. 785-794. Association for Computational Linguistics, 2017.
Viet Dac Lai, Nghia Trung Ngo, Amir Pouran Ben Veysch, Hieu Man, Franck Demoncourt, Trung Bui, and Thien Huu Nguyen. Chatgpt beyond english: Towards a comprehensive evaluation of large language models in multilingual learning. arXiv preprint ar arXiv:2304.0500 2023.
Mike Lewis, Yinhan Liu, Naman Goyal, Marjan Ghazvininejad, Abdelrahman Mohamed, Omer Levy, Veselin Stoyanov, and Luke Zettlemoyer. BART: denoising sequenne-to-ssqquenee pre- training for natural language generation, translation, and comprehension. In ACL, pp. 7871-7880. Association for Computational Linguisties, 2020.
Shen Li, Yanli Zhao, Rohan Varma, Omkar Salpekar, Pieter Noordhuis, Teng Li, Adam Paszke, Jeff Smith, Brian Vaughan, Pritam Damania, and Soumith Chintala. Pytorch distributed: Experiences on accelerating data parallel training. Proc. VLDB Endow,, 13(12):005533018 aug 2020. ISSN 2150-8097.
Zihao Li, Yucheng Shi, Zirui Liu, Fan Yang, Ninghao Liu, and Mengnan Du. Quantifying multilin- gual performance of large language models across languages. arXiv preprint arXiv:2404.11555 2024.
Tao Lin, Sebastian U. Stich, Kumar Kshitij Patel, and Martin Jaggi. Don't use large mini-batches, use local SGD. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net, 2020.
Lin Long, Rui Wang, Ruixuan Xiao, Junbo Zhao, Xiao Ding, Gang Chen, and Haobo Wang. On llms-driven synthetic data generation, curation, and evaluation: survey, arXiv preprint 10X1:22406.15126, 2024.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net, 2019.
Alexandre Magueresse, Vincent Carles, and Evan Heetderks. Low-resource languages: A review of past work and future challenges. CoRR, fass2006.07264/ 2020.
Sam McCandlish, Jared Kaplan, Dario Amodei, and OpenAl Dota Team. An empirical model of large-batch training. CoRR, abs/1812.06162. 2018.
Sean McLeish, Arpit Bansal, Alex Stein, Neel Jain, John Kirchenbauer, Brian R. Bartoldson, Bhavya Kailkhura, Abhinav Bhatele, Jonas Geiping, Avi Schwarzschild, and Tom Goldstein. Transform- ers can do arithmetic with the right embeddings. CoRR, abs/2405. 17399, 2024.
Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communication-effcieenn learning of deep networks from decentralized data. In Artificial intelli- gence and statistics. PMLR, 2017a.
Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communication-efficeent learning of deep networks from decentralized data. In Artificial intelli- gence and statistics, pp. 1273-1282. PMLR, 2017b.
AI Meta. Introducing meta llama 3: The most capable openly available Ilm date. Meta AI, 2024.
14

---

Published as a conference paper at ICLR 2025
Luca Moschella, Valentino Maiorca, Marco Fumero, Antonio Norelli, Francesco Locatello, and Emanuele Rodolà. Relative representations enable zero-shot latent space communication. In ICLR. OpenReview.net, 2023.
Vladislav Mosin, Igor Samenko, Borislav Kozlovskii, Alexey Tikhonov, and Ivan P Yamshchikov. Fine-tuning transformers: Vocabulary transfer. Artificial Intelligence, 317:103860 2023.
Alex Nichol, Joshua Achiam, and John Schulman. On first-order meta-learning algorithms. CoRR, hts//803.02999) 2018.
Nous Research. DisTrO, 2024. URL htt.esss/otttuu.cc/////eeeeeeeea blob/main/a A Preliminary Report on DisTrO.pdF
OpenAl, Dec 2023. URL https://openaa:: com.
Jose Javier Gonzalez Ortiz, Jonathan Frankle, Mike Rabbat, Ari S. Morcos, and Nicolas Ballas.
Trade-offs of local SGD at scale: An empirical study. CoRR, abs/21 10.00133, 2021.
Sahil Patel and Stephanie Palazzolo. OpenAl offers publishers as little as million a year the information, Jan 2024.
Telmo Pires, Eva Schlinger, and Dan Garrette. How multilingual is multilingual bert? In ACL (I), pp. 4996-5001. Association for Computational Linguistics, 2019.
Ofir Press, Noah Smith, and Mike Lewis. Train short, test long: Attention with linear biases enables input length extrapolation. In International Conference Learning Representations, 2022.
Alec Radford, Karthik Narasimhan, Tim Salimans, and lly Sutskever. Improving language un-
derstanding by generative pre-training. OpenAl blog, 2018 URL httpt::ttp//eeeeetee com/
blog/language unsupervised/
Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners, 2019.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. J. Mach. Learn. Res., 10:1011:140:6, 2020.
Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. Zero: memory optimizations toward training trillion parameter models. Christine Cuicchi, Irene Qualters, and William T. Kramer (eds.). Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, SC 2020, Virtual Event / Atlanta, Georgia, USA, November 9-19, 2020. pp. 20. IEEE/ACM, 2020.
François Remy, Pieter Delobelle, Hayastan Avetisyan, Alfiya Khabibullina, Miryam de Lhoneux, and Thomas Demeester. Trans- trann-tokenization and lingull vocabulary transfers: Language adaptation of LLMS for low-resource NLP. In First Conference Language Modeling, 2024. URL htt.tpss://ppeneettttet wwtpeeewwwwwweeeeeeeeeeeeeeeeeee /fot/fformmt d d=sBxvoDhvao
Phillip Rust, Jonas Pfeiffer. Ivan Vulic, Sebastian Ruder, and Iryna Gurevych. How good is your tokenizer? on the monolingual performance of multilingual language models. In ACI/IJCNLP (1), pp. 3118-311100 11-3133. Association for Computational Linguistics, 2021.
Lorenzo Sani, Alex Iacob, Zeyu Cao, Bill Marino, Yan Gao, Tomas Paulik, Wanru Zhao, William F. Shen, Preslav Aleksandrov, Xinchi Qiu, and Nicholas D. Lane. The future of large language model pre-training is federated, 2024.
Teven Le Scao, Angela Fan, Christopher Akiki, Ellie Pavlick, Suzana Ilie, Daniel Hesslow, Roman Castagné, Alexandra Sasha Luccioni, François Yvon, Matthias Gallé, Jonathan Tow, Alexan- der M. Rush, Stella Biderman, Albert Webson, Pawan Sasanka Ammanamanchi, Thomas Wang, Benoît Sagot, Niklas Muennighoff, Albert Villanova del Moral, Olatunji Ruwase, Rachel Baw- den, Stas Bekman, Angelina AiMillln--aaoon Iz Beltagy, Huu Nguyen, Lucile Saulnier, Samson Tan, Pedro Ortiz Suarez, Victor Sanh, Hugo Laurençon, Yacine Jernite. Julien Launay, Margaret Mitchell, Colin Raffel, Aaron Gokaslan, Adi Simhi, Aitor Soroa, Alham Fikri Aji, Amit Alfassy,
15

---

Association for Published as a conference paper at ICLR 2025
Anna Rogers, Ariel Kreisberg Nitzav, Canwen Xu, Chenghao Mou, Chris Emezue, Christopher Klamm, Colin Leong, Daniel van Strien, David Ifeoluwa Adelani, and et al. BLOOM: A 176b- parameter open-access multilingual language model. CoRR, abs/2211.0000 1.0221111100000 1000 2022.
Richard Socher, Alex Perelygin, Jean Wu, Jason Chuang, Christopher D. Manning, Andrew Ng. and Christopher Potts. Recursive deep models for semantic compositionality over a sentiment treebank. In EMNLP, pp. 1631-1642. ACL, 2013.
Sebastian U. Stich. Local SGD converges fast and communicates little. In 7th International Con- ference on Leurning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. Open- Review.net, 2019.
Chaofan Tao. Qian Liu, Longxu Dou, Niklas Muennighoff, Zhongwei Wan, Ping Luo, Min Lin, and Ngai Wong. Scaling laws with vocabulary: Larger models deserve larger vocabularies. CoRR, abs/2407.13623 2024.
Asahi Ushio, Yi Zhou, and José Camacho-Collados. Efficient multilingual language model com- pression through vocabulary trimming. In EMNLP (Findings), pp. 14725-14739, Association for Computational Linguistics, 2023.
(Findiniin
Ahmet Üstün, Viraat Aryabumi. Zheng Xin Yong, Wei- Yin Ko, Daniel D souza, Gbemileke Onilude, Neel Bhandari. Shivalika Singh, Hui-Lee Ovi, Amr Kayid, Freddie Vargus, Phil Blunsom, Shayne Longpre, Niklas Muennighoff, Marzieh Fadaee, Julia Kreutzer, and Sara Hooker. Aya model: An instruction finetuned open-access multilingual language model. In ACL (1), pp. 15894-15939. Association for Computational Linguisties, 2024.
Pablo Villalobos, Jaime Sevilla, Lennart Heim, Tamay Besiroglu, Marius Hobbhahn, and Anson Ho. Will we run out of data? analysis of the limits of scaling datasets in machine learning. CoRR, abs/221 1. 04325, 2022.
Guan Wang, Sijie Cheng, Xianyuan Zhan, Xiangang Li, Sen Song, and Yang Liu. Openchat: Ad- vancing open-source language models with mixed-qualin data, arXiv preprint vrXiv:2309.12233 2023.
Zirui Wang, Zachary Lipton, and Yulia Tsvetkov. On negative interference in multilingual mod- els: Findings and A meta-learning treatment. In EMNLP (/), pp. Computational Linguisties, 2020.
Adina Williams, Nikita Nangia, and Samuel R. Bowman. A broad-coverage challenge corpus for sentence understanding through inference. In NAACL-HLT, pp. 1112-1122. Association for Com- putational Linguistics, 2018.
Herbert Woisetschläger, Alexander Erhen, Bill Marino, Shiqiang Wang. Nicholas D. Lane, Ruben Mayer, and Hans-Amo Jacobsen. Federated learning priorities under the european union artificial intelligence act. CoRR, abs/2402.059968, 2024.
Mitchell Wortsman, Peter Liu, Lechao Xiao, Katic E. Everett, Alexander A. Alemi, Ben Adlam, John D. Co-Reyes, Izzeddin Gur, Abhishek Kumar, Roman Novak, Jeffrey Pennington, Jascha Sohl-Dickstein, Kelvin Xu, Jaehoon Lee, Justin Gilmer, and Simon Kornblith. Small-scale prox- ies for large-scale transformer training instabilities. In ICLR. OpenReview.net, 2024.
Yangyifan Xu, Jinliang Lu, and Jiajun Zhang. Bridging the gap between different vocabularies for LLM ensemble. In NAACL-HLT, pp. 7140-7152. Association for Computational Linguistics, 2024.
Linting Xue, Noah Constant, Adam Roberts, Mihir Kale, Rami AI-Rfou, Aditya Siddhant, Aditya Barua, and Colin Raffel. mt5: A massively multilingual pre-trained text-to-text transformer. In Kristina Toutanova, Anna Rumshisky, Luke Zettlemoyer, Dilek Hakkani-Tür, Iz Beltagy, Steven Bethard, Ryan Cotterell, Tanmoy Chakraborty, and Yichao Zhou (eds.), Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, NAACL-HLT 2021, Online, June 6-11, 2021, pp. 483 498. As- sociation for Computational Linguistics, 2021.
16

---

Published as a conference paper at ICLR 2025
Hao Yu, Rong Jin, and Sen Yang. On the linear speedup analysis of communication efficient mentum SGD for distributed non-convex optimization. ICML, volume 97 Proceedings of Machine Learning Research, pp. 7184-7193. PMLR, 2019.
Shaolei Zhang, Qingkai Fang, Zhuocheng Zhang, Zhengrui Ma, Yan Zhou, Langlin Huang, Mengyu Bu, Shangtong Gui, Yunji Chen, Xilin Chen, et al. Bayling: Bridging cross-lingual alignment and instruction following through interactive translation for large language models. arXiv preprint 17X9::2306.10668, 2023.
Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christo- pher Dewan, Mona T. Diab, Xian Li, Victoria Lin, Todor Mihaylov, Myle Ott, Sam Shleifer, Kurt Shuster, Daniel Simig, Punit Singh Koura, Anjali Sridhar, Tianlu Wang, and Luke Zettle- moyer. OPT: open pre-trained transformer language models. CoRR, abs/2205.0100 20220000006,,, 2022.
Wanru Zhao, Yihong Chen, Royson Lee, Xinchi Qiu, Yan Gao, Hongxiang Fan, and Nicholas Don- ald Lane. Breaking physical and linguistic borders: Multilingual federated prompt tuning for low-resource languages. In The Twelfth International Conference on Learning Representations, 2024.
Yanli Zhao, Andrew Gu, Rohan Varma, Liang Luo, Chien-Chin Huang, Min Xu, Less Wright, Hamid Shojanazeri, Myle Ott, Sam Shleifer, Alban Desmaison, Can Balioglu, Pritam Damania, Bernard Nguyen, Geeta Chauhan, Yuchen Hao, Ajit Mathews, and Shen Li. Pytorch FSDP: experiences on scaling fully sharded data parallel. Proc. VLDB Endow. 16(12):3848-33)) -8600, 2023.
17

---

Published as a conference paper at ICLR 2025
A
EXPERIMENTAL DETAILS
A.1
MODEL ARCHITECTURES AND HYPERPARAMETERS
Table 8 presents the vocabulary-agnostic hyperparameters of our decoder-only models, while Ta- ble 9 details vocabulary sizes, DEPT-specific parameters, memory costs, and communication costs. Standard pre-training pipeline parameters were chosen based on the recommendations of Hottmann et al. (202)) and MosaicML, except for the billion-scale model, where we aligned with the recent state-of of tft (SOTA) for English federated pre- pre-training by |Sani et al. (2024). We always use gradient clipping norm of 1 and ALiBi (Press et al., 2022) positional embeddings.
During continued pre-training, for models initialized randomly, we begin with 7mnax and decay over Ner learning steps, allowing quick embedding matrix learning without requiring another full train- ing pass, as common in language rewiring (Artetxe et al., 2020). When using pre-initialized models, we start from 7/max 2 since both the model and embeddings are reasonably well-trained.
Importantly, the only parameter changed between DEPT models and baselines is the learning rate 7max We use the same learning rate to contrast convergence properties for comparisons in Fig, 2 We tune the baselines' learning rate for later comparisons to ensure they perform the same number of training steps, selecting the best checkpoint for a baseline across all experiments. Except for tuning the learning-rate, DEPT models always use the same hyperparameters as the baselines during local training.
Table 8: Architectural details and vocabulary-independes hyperparameters of our models. The number of transformer blocks is denoted by #Blocks, the number of attention heads by #Heads, and the expansion ratio refers to the ratio of the hidden dimension in the feedforward layers. The total number of model parameters is MM, the vocabulary size is (V]. and the model embedding dimension is dmodel- We train standard decoder-only transformers whose body ranges in size from 86.4M to 1.2B independent of embeddings. As see in Table the size of the embedding matrix the model size drastically. Our batch size is [B] while [Sc|/|S] is our sampling ratios for the various data sources. The 31, .3 pair are AdamW parameters while the S. tuple represents the parameters of the cosine scheduler that we use, including the decay alpha the decay period 1/max. and the total number of sequential steps N. Finally, we show the number of continued pre-training steps NeT that we use, representing 15% of total steps for the 298M model and 19.3% for the 86. 4M model. All of our models use a sequence length of 2048. We followed the hyperparameters of [Sani al.) (2024) for the billion-scale federated pre-training. We report the tuned 7mmax. for each baseline according STD(r=0) (T.( STD(r=1)
change
to Appendix A.1.2 7/max
7/max
1/max
we find that the embedding resting allows
ACT to use the same 7mma as DEPT. T.
Tye
JM...
&Boal
Eag,
[8]
(8.. 3))
Soloe, Boaer N)
Nay
Waac
Malt-domaan
12
706
S6.IM
256
4/10
[30. 0.60
(00
60 18
(9")
1.x 10
6.6. 10
46 10 2
Matti-denati
3024
tM
4/16
[30 036)
[E
24 10
Matidlagual
2548
1.28
512
(591 (366
(18"1,
1i Tx kai)
LAx 10
We had to select a particular sampling ratio for the continued pre-training using the full pre-training set rather than a single language domain. Due its high heterogeneity, default to uniform sampling for MC4 in these cases. In contrast, for The P Pile, we preferred proportional sampling as the dataset is entirely in English and has already had its data sources upsampled/downsampled based on usefulness. We also provide results using the alternative sampling policy in Appendix B
A. SOFTWARE AND HARDWARE
Our software is based on the MosaicML. composer (Databricks) 2024) library for LLM pre-training and the open-source Flower (Beutel et al. 2022) framework for federated learning. Crucially, we heavily rely the MosaicML hyperparameters and infrastructure for our InnerOPT, making no changes to it after our embedding matrix manipulation from Algorithm [ has been performed. For the standard baselines, we ran them completely unmodified version of the MosaicML code- base (beyond using our data), which has been independently verified by thousands of users and used to submit accepted conference publications (Blakeney et al. 2024).
18

---

Published as a conference paper at ICLR 2025
Table 9: Practical memory and communication costs for DEPT, where the total number of steps is N = Noccll with T the total number of iterations, and V& as the average vocabulary size across data sources. Standard pre-training requires a full in-memory embedding matrix for the global vo- cabulary while synchronizing gradients every step rather than every Nocal steps. All DEPT variants yield communication savings, with GLOB as the baseline. TRIM provides additional savings propor- tional to the gap between global and local vocabulary sizes, while SPEC further reduces costs with or without optimized vocabularies by never communicating the token or positional matrices.
Type
#Blocks
Method
Niocat
T
[V 1
V.l x desoda
Myg (4)
Per-step Comms Cost (4)
Multilingual
10"
192M
275M (1x)
(1x)
Multilingual
12
GLOB
500
10
250 112
192M
278M (1x)
0.56M (0.002x)
Multilinuual
12
TRIM
500
10
216 135
1: 27 160
166M
252M (0.92x)
0.5M (0.002x]
Multilingual
12
SPDC
500
10
216
160
166M
252M (0.92x)
0.17M (0.0006
Multiliguaa
12
SPRC-OPT
500
10
301257
38.6M
125M (0.45%)
0.17M (0.0006X) X)
Multilingual-I
24
STD
x 10'
1
250 112
512-2M
1,718 (1x)
1,718 (1x)
Multilingual-B
SPEC-OPT
500
14
501257
102.9M
1.38 (0.76x)
2.4M (0.00T
Multi-domain
12
STD
5 x 10'
202557
38.6M
125M (1x)
125M (1x)
Multi-domain
12
GLOB
500
10
257
38.GM
125M (1x)
0.35M (0,002X)
Multi-domain
500
45 554544 9
35M
1211 (0.97x)
W22-0 (0.002x)
Multi-donain
12
SPEC
500
10
15 551
= 9-002
35M
121M (0.97x)
WLLO (0.001x)
Multi-domain
21
STD
13.5 10"
257
51. 11-LM
350M (1X)
350M (1x)
Multi-domain
GLOB
500
27
350M (13)
6.7M (0.002x
Multi-domain
24
TRIM
500
27
45 554 94822
46.6M
346, 2M (0.97X)
OR90 (0.002x)
Multi-doman 24
3PEC
500
27
45 554 =: 9462
46.6M
345.2M (0.97x)
0.6M (0.002x)
In terms of hardware, the low communication properties of DEPT allowed us run experiments via mixture of loaned resources from separate cloud providers. Over the course of our experimentation, used various machines equipped with either 1 H100 A100 GPU in the USA. Canada, and Europe, which turned out to be more cost-effective. We rented machines with 4-8 H100 GPUs for the centralized baselines since could not Distributed Data Parallelism techniques over low- bandwidth internet connections. When the standard training baseline has sufficiently low learning rate converge, the difference in training time is driven by three factors.
First, the throughput achieved by individual workers: for GLOB, this should be identical to standard pre-training as the model in memory remains unchanged. For TRIM and SPEC, the reduced memory requirements may allow increasing the device micro-batch size in certain scenarios (but not the global batch size, which heavily influences optimization properties). This depends heavily on the hardware; for example, in DeepMind Mathematics workloads, TRIM or SPEC can double the device micro-batch size, and similarly for SPEC-OPT in the case of multilingual data.
Second, the communication topology significantly impacts wall clock time. For instance, in 10 Gbps bandwidth connection using Ring AllReduce for aggregation across workers, DEPT can reduce training time by 33% for a billion parameter model. In cases with very fast connection, such as InfiniBand, the training time difference is primarily determined by throughput differences.
Third, the number of local data sources and the number of available workers impact the total training time, for DEPT we always scale the number of workers to match the number of data sources exactly.
A.1.2 HYPERPARAMETER TUNING METHODOLOGY
Given that MosaicML. provides hyperparameter-tune models on the C4 (Raffel et al. 2020) dataset, we use their learning rate schedule and number of training steps starting point. In the case of DEPT, we find that we can always use the MosaicML parameters since the OuterOpt application of DEPT acts as regulariser via coise-injection (Lin et al. 2020) and meta-learning effects (Nichol et al. 2018). This makes DEPT models highly unlikely to diverge, even under extreme data hetero- geneity and without a shared input or output space. In the case of standard training baselines, we gradually lower the learning rate, starting from the one reported in Table
We begin with the maximum learning rate Mmxx and systematically reduce on coarse grid in
intervals of 0.5 x 10-5;
11 7max 0.5k x 10 1: E K).
19

---

Published as a conference paper at ICLR 2025
where & represents the step index, and K is chosen such that 1j >> 0 at the final step. Given that the length of the cosine cycle is directly extrapolated from known scaling laws on the number of tokens that the model needs to train on for compute- compute-optimality (Hoffmann al. 2022). approximately 20 tokens per parameter, we stop as early as we find a learning rate that can complete the entire co- sine schedule. Then, we choose the best-performing checkpoint, according to validation perplexity, across all experiments. We report these values in Table
This hyperparameter search does not cover all possible relevant parameters; given enough resources, we would also tune the gradient clipping norm. Furthermore, we could tune the batch size using the empirical model of large-batch training proposed by McCandlish al. (2018). Given that the appro- priate learning rate depends on the chosen batch size and the desired target loss, such an optimization would require hundreds of experiments across all baselines find an optimal configuration.
A.1.3 ADAPTING ACTIVE FORGETTING
To implement the active forgetting baseline (Chen et al. 2023) ACT, we had to adapt the method- ology to decoder-only models, which train with far fewer steps. To achieve this, we use a forgetting frequeney of 500 steps, equal to DEP T's Miocll. We also use cosine scheduler for the body with the same parameters as shown in Table S across the 500 steps using the same scheduler but setting V 500. Finally, we selected the checkpoint with the lowest validation perplexity for continued pre-training in a forgetting cycle.
however, we schedule the embedding matrix independently
A.2 DATA SOURCES
We quantify the lexical heterogeneity of a dataset based on lexical similarity between data sources. A simple similarity measure is the size of the intersection of subwords between vocabularies. The smaller the intersection, the more dissimilar the vocabularies, and thus, the more challenging it becomes to train a shared tokenizer effectively across different domains or languages. For this section,
the size of local vocabulary subset of the global vocabulary proxy, with smaller local vocabulary indicating that global tokenization does not serve a particular data source well.
Our default global tokenizer for multilingual data is that proposed by Xue et al. (2021). with V 250 112.0 tokens. Owing to its diverse pre-training, the mT5 (Xue et al. [202]] tok- enizer is a robust default choice, employed in recent works such as project Aya (Ustin et al., 2024]. However, its coverage of hundreds of languages does come with many shortcomings relating to the capacity allocated to each language. To showcase these challenges, we care- fully selected languages from distinct families in the MC4 subset, including English (EN), Italian (IT), Serbian (SR), Swahili (SW), Urdu (UR), Latin (LA), Chinese (ZH), and Malay (MS). The corresponding vocabulary sizes of our languages are as follows: (247 720, 211 332, 208 391, 170 984, 188 002, 220 757, 240 566). Among these, Swahili (SW) is the most heterogeneous, as determined by its small subset of 170 984 tokens.
Our global tokenizer for English data posed by Black et al. (2022) with V = 257 tokens. We selected The Pile e as our multi- domain dataset for several reasons. The Pile is a diverse, large-scale dataset specifically de- signed for training large language models (LLMs). Its diversity spans domains such as scientific papers, news, books, and web content, providing comprehensive foundation for capturing var- ied linguistic patterns. Among the various subsets of The Pile. DM Mathematics stands out as the most heterogeneous. This subset contains only 11, 090 tokens from the global vocabu- lary, significantly fewer than other subsets. Here are the sizes of other subsets in terms of their unique tokens from the global vocabulary: (49 362, 49 783, 46 766, 49 469, 49 700, 47 865, 48 720} (11 090, 44 249, 42 957, 44 432, 49 992, 49 841, 47 687, 49 961, 46 825). While this indicates much lower heterogeneity than in multilingual settings, vocabulary choice may still impact highly special- ized model capabilities such as mathematical reasoning.
trained The Pile (Gao et al. 2021 and pro-
A.2.1 TOKENIZATION CONSIDERATION
One the major challenges when representing multiple data sources with a single tokenizer is vocabulary dilution. To maximize coverage, tokenizer that aims cover multiple languages or
20

---

Published as a conference paper at ICLR 2025
domains often needs adopt many short subwords. This increases the tokenizer fertility (i.e., the number of tokens produced per unit of text) (Rust et al. 2021) and also raises the overall description length the total number of tokens required represent the same data. This trade-off negatively affects the compression ratio, as the same amount of information requires more tokens, reducing the model's s sample efficiency (Tao et al. 2024). When non-uniform sampling ratios used dur- ing pre-training, high-resource languages tend to have better fertility than low-resource languages. This means high-resource languages are better represented the vocabulary, and their tokens are more likely to be shared across the model's parameters, improving their performance. In contrast, low-resource languages suffer from poor fertility, where their unique vocabulary tokens are under- represented, leading to worse performance. For example, twahili. (SW) and Urdu (UR) are low-resource languages that face these challenges. Our SPEC method allows us to avoid many of these issues by providing an optimized vocabulary to data source at the cost of losing shared vocabulary and updating several embedding matrices. An alternative approach is to cluster vocab- ularies (Chung al. 2020) to obtain subword sharing between more relevant languages. However, this requires that participating data sources are known in advance, do not change significantly, and that the appropriate number of clusters is also known.
To account for the effectiveness of a tokenizer on a given language, we report unigram cross- coss-eytropy in experiments, which represents how effective simple unigram model based the tokenizer is on that data source as proxy for the effectiveness of the tokenization. If the unigram cross-entropy high on a given data source, is likely underserved by the tokenization. Thus, all improvements brought about by using a more complex language model must consider this baseline. It can also be used compute unigram-norman normalized cross-entropy perplexity, language modeling performance metric that is comparable across different vocabulary sizes (Tao et al. 2024).
B ADDITIONAL RESULTS
B.
LARGER MODEL
500
EN
SR
400
LA
SW
Perplexity
UR
300
MS
ZH
IT
200
100
0
2
6
8
10
12
14
Federated Round
Figure 4: Convergence plot of our 1.3 billion model trained in a vocabulary agnostic federated fashion. For the initial rounds, we sample 4 data sources of 8; after seeing most of the clients, we reduce the number to 2. We make sure only to introduce EN later into the experiment.
Figure provides further insights into the performance of DEPT on a larger-scale experiment with 1.3 iillionpaarmmtter model. In this setting, the model is trained in vocabulary agnossic,, fed- erated fashion with dynamic client subsampling. During the initial rounds, out of data sources sampled, which is reduced after most clients have been processed. Importantly, EN is in-
21

---

Published as a conference paper at ICLR 2025
troduced later in the training process to evaluate the model'ss cross-lingual transfer capabilities to this high-resource language. The plot illustrates that the transformer body, enabled by DEPT, effec- tively transfers knowledge across languages and domains, allowing newly introduced previously stale data sources to converge to perplexity levels similar to their peers within one or two sampling rounds. This experiment underscores the feasibility and scalability of using DEPT for collaborative large-scale language model pre-training, even under extreme client subsampling and without prior knowledge of the underlying data distribution.
B.2 PLASTICITY
230
200
SID
Puplosity
ACT
Perpectts
ACT
DEPILAVG
30-
1000
500
Sequential
Sequestial
(a) MC4--FFFJ FUI.,, 12-block.
(h) 12-hlock.
0.50
Perplosks
1.0
Porpleity
15
NEPT:
DEPT Avi
100-
200
600
502
1000
1200
200
1000
1200
Sequeutial Stege
Sequential Steps
(c) HI, 12-block.
(d) DE, 12-block.
Figure 5: Adaptation curves starting from randomly initialized matrix. DEPT is always stable in its convergence, reaching the lowest perplexity for the pre-training distribution (MC4-FULL), for the lowest-resource languages the distribution (SW), and for the two out-of-distributtin lan- guages (HI, DE). It is also always the fastest to adapt.
The results presented in Figure demonstrate the robustness and adaptability of DEPT across var- ious settings, completing the plot shown Fig. 3 Specifically, DEPT consistently achieves the lowest perplexity across all scenarios: (1) the full pre-training distribution (MC4-FULL)) (2) the lowest-resource language within the dataset (SW), and (3) two out-rf-distributton languages (HI and DE). Furthermore, DEPT is not only effective in reaching convergence but also does so at faster rate compared to other approaches. These results showcase its utility in a wide range of multilin- gual and domain-adaptive pre-training tasks; for example, if new client were to be introduced in a federated setting, they show that the DEPT trained model could quickly adapt to its data distribu- tion. Alternatively, multi-phase adaptive pre-training represents distinet advantage in terms of data efficiency.
B.3 IID DATA PERFORMANCE
In the case of IID data (represented by random sharding of the C4 dataset), Fig. ssowss that DEPT performs similarly to standard pre-training with the benefit of lower activation norms, indicating the potential for longer and more training.
22

---

Published as a conference paper at ICLR 2025
0.6
STD (r
NEPPT-NN
200
Actioation (1.2)
0.4
Perpleatt
00
0.1
0.0
3000
Sequestial Stops
(a) C4, 12-block, Perplexity
(b) C4, 12-block, Activations
Figure 6: Perplexity (a) and activations (b) curves for DEPT versus uniform sampling the IID C4 dataset. DEPT models, outside temporary spikes caused by OuterOpt, perform similarly to standard pre-training regarding training perplexity, However, seen from the activations, still provides greater training stability with the potential of extending pre-training.
B.4 ONE-SHOT GENERALIZATION
Table 10: Validation perplexity (U) for our 24-block models trained on The Pile when using continued pre-training with uniform sampling starting from randomly-initialized embeddings. DEPT provides a better transformer body for all datasets, outperforming baselines by 17.5% on average.
Name
DM
EN
EP
FL
CC
PA
EN
PP
WK
AX (7.7)
UB
PC
GU
HN
Ul-0OD
AYG
(NNGRAMAE))
(6.9)
(7.9)
(10)
(7.8)
(7.9)
(7.9)
(8.2)
(7.7)
(9.1)
(8.2)
(7.8)
(5)
(8.1)
(7.7)
(7.7)
(10)
(8.1)
STD
4.7
17.5
19.7
22.3
64.2
30.7
18.3
48.5
41.6
13
23.9
19
32.1
50.5
42
72.9
31.1
STD (+
5.1
24.1
27.1
31.4
9.2
86.6
43.4
24.6
64.4
58.4
16.5
32.7
25.1
44.7
66.7
55.8
141
44.5
26.4
16.3
21
27.2
SPIC
14.5
16.2
18.8
55.5
25.8
16
41.1
35.
11.5
20.5
16.5
27.2
43
35.7
63.5
26.8
SPXC-OPT
4.6
15.2
16.9
19.4
6.4
67.1
26.5
16.4
42.5
35.9
11.0
21
16.5
27
44.9
57.1
60.4
27.1
Min Imp (%)
2.5
13.5
14.3
12.9
8.4
11:2
13.9
10.4
12.5
13.7
8.2
12.3
11.5
13.4
11
11.8
9.2
12.5
Max Imp (%)
4.8
19.1
18
19.3
16.5
19
15.3
17.1
18.3
13.5
17.6
16
18.6
17.1
17.6
19.4
17.5
B.5 TRANSFORMER BODY GENERALIZATION
Table 10 shows the performance of DEPT the The Pile dataset with 24-block model trained from randomly initialized embeddings. Here, DEPT outperforms all baselines across subsets, with average improvements of 17.5%.
B.6 PRE-TRAINED EMBEDDING MATRIX GENERALIZATION
Table 11: Validation perplexity (L) for our 24-block models trained on The Pile when performing continued pre-training with uniform sampling starting from a pre-trained embedding matrix. DEPT wins 10 out of 17 comparisons with TRIM always outperforming GLOB, When comparing against Tables[]and TO, we can observe that DEPT wins the complementary comparisons when starting from random embeddings or when using proportional sampling with the pre-trained embedding matrices. This indicates that baselines always have a worse transformer body, with sampling ratios heavily impacting the effectivvnnes of embeddings for given dataset.
Nause (UNURASS44))
DM
ER.
RP
F1.
GH
CC
PA
SE:
PP
WK
AX
UB
PC
NH
GU
HN
UL-00D
AVG
(6.0)
(7.0)
(10)
(7.5)
(7.9)
(7.9)
(8:2)
(7.7
(9.1)
(8.2)
(7.7)
(7.8)
(6)
(5.1)
6.7)
(7.7)
(10)
(5.1)
STD (r 0)
4.3
11.1
13.4
15
5.5
41.4
20.4
13.2
54
27.1
10.1
16.0
13.6
21.8
35.2
20.1
51.6
21.6
STD (r D
4.3
12.6
15.8
13.5
4.7
38.8
10.2
11.7
36.8
24.7
8.8
16
12
21.7
34.2
26.0
43.1
20.4
GLOB
4.5
12.1
13.8
15.6
5.5
41.4
19.4
129
35.5
25.6
10.1
15.6
13.5
10.8
38.1
20.2
58.2
21.7
2MIM
4.8
10
11.3
14.5
4.9
37.3
18.2
11.8
80.2
23
9.8
15.3
12.8
19
32.9
26.7
47.6
19.4
Min Imp (%)
-0.5
-8.6
-29
-1566
11.8
-888
-0.8
-96
1.2
-3.5
114
2
10.2
8.1
-1..5
-2.2
33.1
-6.1
Max Imap 1%)
-1.2
10.1
15.4
-95
0.1
3.8
5.1
-0.3
11.1
7.1
11.9
42
-6
12.1
3.9
6.6
-104
4.8
When continuing pre-training with pre-trained embedding matrices, shown in Table II DEPT secures 10 out of 17 wins, with TRIM consistently outperforming GLOB. A comparison with Ta- a 5 reveals that DEPT also outperforms in other scenarios, whether starting from random
bles 10
23

---

Published as a conference paper at ICLR 2025
Table 12: Validation perplexity (L) for our 12-block models trained on The Pile when performing embedding matrix. DEPT can train a
continued pre pre-training starting from superior transformer body, outperforming all baselines across all subsets by up to 28%.
Name
N11
G1I
UB
EL
EE
EP
WK
CC
SE
PC
PP
DM
AX
GU
I1N
UI-0OD
AVG
(5.3)
(7.8)
(7.8)
(7:5)
(8.2)
(7.7)
(9:1)
(7.7)
(7-7)
(8.1)
STD
63.0
12.1
46.1
45.2
36.3
81.8
115.9
83.2
32.4
91.8
91.2
196.
02.2
41.5
65.1
30.5
AC2
126.7
79.9
82.1
06.3
121.2
147.7
191.6
50.4
01.22
180.1
351.8
150.8
123.9
377.8
114.9
GLOB
41.5
32.0
12.1
25.0
08.6
23.8
134.6
7AIM
43.3
31.2
30.7
24.6
29.4
66.2
82.5
23.4
22.6
62.7
64.
53.2
40.8
SPSC
42.1
8.7
40.6
30.3
20.8
8
28.0
54.8
57.0
21.9
23.8
07.5
5.2
10.8
09. 1
57.1
121.2
43.2
Min Amp (%)
24
19
25
28
24
23
27
22
20
20
21
21
18
19
20
26
24
Max Tmp (%)
24
24
29
27
29
20
155
27
24
25
26
27
22
25
26
41
28
Table 13: Validation perplexity (L) for our 12-block models trained on The Pile when performing continued pre-training starting from a pre-trained embedding matrix. DEPT performs worse than for the 24-block trained The Pi le and than for our MC4 models. However, when considering Table [2] we can observe wins all comparisons by wide margins when starting from a randomly initialized embedding matrix, indicating that this gap is driven by the embedding space being fitted to the high-resource languages despite the baselines having a worse transformer body.
Name (UNGRAMA4))
RR
PC
AX
UL-000
(5.1)
(7.9)
(8.2)
(7.8)
(7.8)
(7.9)
(10)
(5.2)
(7.9)
(3.7)
(8)
(9.1)
(6.10
(7.7)
(7.7)
(1.7)
(10)
(8.1)
5tD (y
90.5
23.3
18.1
19.2
13.6
STD (1 15
52.7
29.5
28.4
21.8
20.
38.4
67.9
17.0
17.2
56.8
12.0
51.2
02.8
81.5
31.9
GLOB T0.MM
29.8
22.8
23.7
200
19.6
13.7
29.5
29.2
22.4
23.0
19.4
40.8
60.6
17.4
19.2
46.7
6.1
13.4
81.1
$1.1
Min Tmsp (%)
-7.0
-1.2
-11.2
-67
-68
-188
-6.5
-14.3
-92.7
-37
Imp (5)
4.00
1.0
43
-.9
-29
-62
4.7
-26
-114
-.2
-11.5
-2.
-186
emheddings or leveraging pre-trained ones. This consistency underscores the robustness of DEPT's transformer body across varying embedding initialization and sampling strategies.
B.7 SCALING EXPERIMENTS
Here, we train smaller multi-domain models with 12 blocks to validate the scaling properties of DEPT across model sizes. In Table T2 we observe that, similar to Table T0. DEPT models outper- form all baselines significantly when starting from random initialization. Importantly, the embed- dings constitute a larger percentage of the model parameters at this model scale. This highlights the robustness of DEPT's modifications to the embedding space in enabling the training of better transformer body.
Interestingly, when using pre-trained embeddings (Table (). the smaller DEPT models perform worse than their larger counterparts. We speculate that the amount of local per-source training performed by DEPT prior to Out OuterOpt should scale with model size. At this scale, the aggrega- tion procedure may be overly harsh on the embedding parameters, particularly for the GLOB and TRIM configurations. This suggests that careful adjustments to the aggregation procedure may be necessary to maintain DEPT's effectiveness at smaller model scales.
B.8 COMPARISON AGAINST SINGLE-CLIENT MODELS
To study the impact of model averaging, we now compare DEPT-based models with models trained on isolated data sources that are never averaged/merged. For fair comparisons, the model of each data source has seen as many tokens as it would have component in DEPT-based training and has undergone continued pre-training for the same number of steps with access to the full dataset. Additionally, DEPT models have undergone continued pre-training from random initialization. If we had compared against such models without the continued pre-training phase, they would have dominated their respective data source while losing all other comparisons, especially in the case of multilingual data.
Tables 14 and 15 show how DEPT models perform when all participants start from random initial- ization. Since the models trained on isolated data sources do not get to keep their highly specialized embeddings, this comparison evaluates how generalizable the abstractions learnt by the transformer
24

---

Published as a conference paper at ICLR 2025
Table 14: Validation perplexity (4) for 24-block models trained on The Pile after continued pre- training with proportional sampling from randomly-initiallle embeddings, compared to models which had been pre-trained on single data source for the same total number of tokens as DEPT has seen from their distributions. DEPT outperforms all baselines. DEPT outperforms all baselines. Baselines whose pre- pre-training dataset matches the evaluation dataset are highlighted in olive.
Name (UNGRAMGH))
EE
CC
UB (7.5)
PC (8)
NH
UL-OOD
AVG
(7.9)
(10)
(7.8)
(7.9)
(7.9)
(7.7)
(9.1)
(8.2)
(8.1)
(10)
15.8
96
25.1
19.1
27.9
72.8
21.9
14.8
130.G
4.9
28.9
11.7
53.5
29.
15.9
95.8
10.5
22.5
15.3
36.5
44.8
97.41
GR
4.9
30.1
43.6
20.8
55.9
30.7
16.5
78.2
33.5
10.9
25.8
16
38.9
54.5
46.7
120.7
37.4
FL
4.9
81.8
48.8
23.5
6.3
68.5
32.4
17.4
84.2
40.9
11.4
21.8
16.8
41
56.8
47.9
122.1
39.1
4.8
25.2
12.2
19.4
5-6
62.9
20
15.5
75
96.6
10.5
22.4
15.3
36.8
50.8
43.4
100
$4.6
RK
48
25.1
12.1
18.9
5.7
51.8
28.4
15.8
74
34..5
10.5
21.9
18.1
35.7
45.88
43.2
90.4
30.9
7.3
140.4
GG9.4
100.6
28
239.5
184
71
G43.9
213.9
31.8
121.7
00.3
213.9
1855
170
96668
226.9
GLOB
4.8
25.7
38.2
17.8
5.4
47.7
25.7
14.7
68.3
32.77
9.9
20
14
32.2
46.38
39.98
94.8
31.6
TRIM
4.8
27.3
39.5
18.5
5.6
51.2
27.8
15.4
71.8
35.1
10.3
21.7
14.8
35.1
40.1
42.2
95.7
33.3
SPSC
4.8
205-
60.1
21.3
14.5
34.3
32.7
48.5
14.1
Imp Imp (%)
1.2
8.3
13.2
7.0
5.6
9.2
Table 15: Validation perplexity (L) for 12-block models trained on MC4 after continued pre- training with unfiorm sampling from randomly-inittalized embeddings, compared to models which had been pre-trained on a single data source for the same total number of tokens as DEPT T has seen from their distributions. Baselines whose pre-training dataset matches the evaluation dataset are highlighted in olive.
In- In-Distribution
out-of-Disttibution
In-Distribution
Out-of--Dstrrbution
Name (UNIGRAM-CE)
ZHH
UR
MS
IT' (7.7)
SR (10.5)
LA
EN
SW
Avg (In-D)
EL
Hl
DE
Avg (OOD)
(9.8)
(10.5)
(9.2)
(7.5)
(10)
(9.3)
(11.1)
(13.9)
(9.7)
(12.6)
2H
187.8
11.6
113.9
988.6
89.1
73.9
128.8
101.6
5744.8
6176.5
1418
1556.1
UR
94.8
27.5
66
06.9
48.7
42.5
79.3
44.
57.5
2596.5
2371.1
000.5
1886
78.6
24.8
42.7
37.3
70
38.9
50.1
2673.3
2329.2
590.4
1967.3
IT
81.3
25.1
59.7
51
43.8
38.2
71.8
39.8
51.3
2617.2
2256.9
615.3
1829.8
85.4
225.7
61
52.4
78.7
2992.1 4
2618.3
657.2
LA
101.8
29.66
T1.7
60.6
G3.5
46
85. 2
47.6
62.4
2838.7 7
2824.8
7-46.6
2136.7
104.4
330.
71.9
61.2
54.3
46.2
85
47.8
62.6
3344.4
3380.6
634.8
2513.3
SW
79.3
24.7
58.1
49.9
43
37.3
69.9
39
50.1
2552.3
2067.5
608.3
1742.7
GLOB
67.7
38
65
45.4
2308.3
1676.5
539.5
1514.7
TRIM
67.7
22.8
55.2
47.5
39.7
35.1
67.2
36.3
46.4
2547.7
1911
567.4
1875.4
SPEC
69.5
23
55.4
47.8
40.3
34.7
68.1
36.3
46.9
2232.1
1578.8
544.7
1451.9
Min Imp (%)
11.8
6.6
4.4
4.2
5.7
2.7
6.8
6.4
0.2
7.6
5.3
3.9
Max Imp (%)
141.1
9.3
7.4
7.7
9.6
9.1
6.5
9.6
9.1
12.5
23.6
9.1
16.7
body are across datasets. In the case of The Pile, shown in Table 14 DEPT models outperform the isolated baselines by 7.8% in terms of average perplexity. Crucially, DEPT models win all com- parisons even though isolated baselines get evaluated on their pre- pre-training dataset, indicating that they have not learned superior abstractions even in this case. For MC4, Table show very similar trend with much higher degree of outperformance for DEPT, 9.8% average on in-distribution data and 6.77% for oul-of-distribution (OOD) data, likely because the transformer body learned for one language has significant difficulty in adapting to multilingual context.
Tables 16 and 17 show the impact keeping the pre-trained embeddings before continued pre- training. The impact of this change is expected: embeddings pre-trained perform well on that dataset; however, they fail to generalize. In the case of The Pile, shown in Table 16. DEPT loses most comparisons the baseline trained on given dataset; however, it outperforms in terms of average perplexity by remarkable 27%. For MC4, shown in Table 17 the outperformance in terms of average perplexity is even more significant, 30.6% for in-distribution data and 14.9% of OOD data.
specific dataset
25

---

Published as a conference paper at ICLR 2025
Table 16: Validation perplexity (4) for 24-block models trained on The Pile after continued pre- training with proportional sampling from pre-trained embeddings, compared to models which had been pre-trained on a single data source for the same total number of tokens as DEPT has seen from their distributions. DEPT significantly outperforms in terms of average perplexity but gets beaten by specialized models on their respective data source. Baselines whose pre-training dataset matches the evaluation dataset are highlighted in olive.
Name (UNEGRASCC))
DM
E.P:
EP
FL.
GH (7:99
CC (7.5)
PA (5:1)
SE (7:1)
PP
WK
AX (1:71
UB
PC
NH [5.1)
GU
HN
UL-DOD [11)
AVG
(5.00
17:9)
[100)
(7:5)
(9.1)
(8:2)
(7.89
(5)
(2-1)
(1:73
(8:1)
10.6
57.6
36.3
32.2
30.8
47.9
133.1
47
52.5
70.2
4.8
25.5
57.5
0.5
14.8
30.77
17.6
23.2
46.4
100.8
4.8
18.7
6.5
25.9
14.7
131.
19.5
13.8
52.9
110.3
47.3
24.4
3.9
97.1
13.3
44.5
25.4
17.9
47.6
60.5
49.8
27.3
31.6
20.4
14.7
17.9
9:3
62.1
70.2
4.9
6.7
15.2
19:4
13..9
30.5
14.4
41.2
116.6
40.8
108.1
424.8
4.5
4.5
17.9
11.2
37.8
22.4
11
35.5
TRTM
4.6
3.6
12
105.9
16.6
11.8
38
56.88
25.7
Min Imop (5)
-2.9
13.9
312
37.44
19.8
21.5
14.1
28.2
13.1
34.8
18.7
5.5
16.9
10..9
3.8
5.3
12.7
18.1
Max Tmp (%6
-1
28.5
02.3
31.1
16.6
-2
1.1
10.5
34.2
20.4
1433
18.1
9.3
11.2
10.1
18.5
10
27
Table 17: Validation perplexity (L) for 12-block models trained on MC4 after continued pre- training with uniform sampling from pre-trained embeddings, compared to models which had been pra-ted on a single data source for the same total number of tokens as DEPT has seen from their distributions. DEPT significantly outperforms in terms of average perplexity but gets beaten by specialized models on their respective data source. Baselines whose pre-training dataset matches the evaluation dataset are highlighted in olive.
In-Distribution
Out-uf-Dissriuttion
Name
ZH
UR
MS (0.2)
II
SR (10.5)
LA
EN
SW
Avg (In-D)
EL
HI
DE
Avg (OOD)
(UNIGRAM.C))
(9.8)
(10.5)
(7:7)
(7:5)
(10)
(9.3)
(14.4)
(13.9)
(0.7)
(12.6)
33.3
87.1
71.3
08.7
58.19
$88.1
5351.5127
$197-97314
936. 293457
3161. 92645
124.7
70.8
57.4
3189.1106
774.645325
802.290588
1588.68217
80.5
27.8
48.1
47.1
38-2
38.3
46.7
2983,89633
2491
620 154907
87781
IT'
91.4
27.4
50.6
26.4
36.0
35.8
59.6
40.7
48.2
2353.07642
3057.45215
344.762726
43083
SR
50.3
28-1
65.1 1
52.5
20.2
41.7
70.1
43.6
52.6
2174.8325
3508. 59985
643.610352
2072. 3474
LA
100.2
30.2
66.5 9
68.1
51
20.1
68.1
45.1
33.7
716.628967
2479.78877
240.702454
1145.73653
1142,11
150.7
276.5
211.1
408.8
147.1
56.2
169.1
323.9
1315877.75
679658 688
11445.9777
6685941137
08.3
47
2782.6792
2673.01815
557.418663
209454
GLOB TRIM
40.1
39.6
1737.5
623
$35.1
41.9
31.3
41.3
40.8
30.8
23.6
1725
345.6
975.3
Min Imp (%)
26.1
288.1
-2.88
68.00
101.0
-53
-43.9
27.8
-142.4
10.4
-43.6
14.9
Max Imp (%)
20.6
22.8
-85
62.1
82.7
47.4
30.6
33.6
30.7
144.4
10.4
43.6
14.9
B.9 COMPARISON AGAINST PYTHIA
The experiments in our work are designed to investigate the outlined research questions instead of producing a state-of-the-art model. However, we believe that providing a comparison against standard baseline may help better contextualize the performance of a given DEPT model. For this purpose, we chose Pythia (Biderman et al. (2023) as it shares a very similar architecture to the one used in work, with the only exception being that Pythia uses untied weights for the embedding matrix and thus all its equivalem sizes have more model parameters. Pythia models are trained on one epoch of the entire The Pile, 300B tokens, regardless of size. Thus, we do not have any OOD dataset for them, and they are expected to perform better on Ubuntu IRC. Since they are trained for many more tokens than DEPT models, we do not perform additional continued pre- training when starting from a treitrained embedding matrix (the extra tokens existed to equalize the amount of work done across baselines); thus, those comparisons show the raw performance of Pythia as published by its authors. When starting from random initialization we use the standard procedure from above.
Table T8 shows comparison between 160M Pythia model and the 125M DEPT models when starting from pre-trained embeddings. At this scale, the additional pre-training Pythia (using 30x the tokens of DEPT) does not provide evident advantage as the model capacity is insuffi- cient to benefit from it. Thus, outside of the expected outperformance on Ubunt IRC (UI), Pythia-1e 60M performs similarly to DEPT models and is slightly outperformed also speculate that using the full 22-dataset version of The Pile during pre- training likely reduced the performance of Pythi 60M as it had to fit broader data distribution. We do not provide
average. We
26

---

Published as a conference paper at ICLR 2025
Table 18: Validation perplexity (4) for our 12-block models trained on The Pile when perform- ing continued pre-training using uniform sampling starting from a pre-trained embedding matrix.
DEPT slightly outperforms Pythia- 6CM at this small scale as its 30x greater number of to- kens is not beneficial with insufficient model capacity. Pythia-1 - 60M was trained on Ubuntu IRC (UI), thus its outperformance is expected as is not an OOD dataset for this model.
UB
CC
PP
(UNGGRASS-H))
(8.1)
(7.9)
(5.2)
(7.8)
(7.8)
(7.9)
(10)
(8.2)
(7.9)
(7.7)
(9:3)
(6.3)
(7.7)
(7.7)
(7.7)
(8.1)
PYIBIA-16OM
31.4
24.1
40.2
64.1
22.
74.5
6.8
GLOB
30.2
29.8
22.0
23.7
20.0
21.9
41.6
61.8
179
19.6
48.0
5.2
13.7
51.1
42.2
90.7
32.4
TRIM
29.5
6.9
29.2
22.4
23.0
19.4
21.1
40.8
60.6
17.4
19.2
46.7
5.1
13.4
52.4
41.0
81.1
31.1
random initialization results for this model size since we found it impossible to make it behave well during continued pre-training and we believe the comparison would be unfair.
Table 19: Validation perplexity (4) for 24-block models trained The Pile after continued pre-training with proportional sampling from handomly-initializee embeddings, compared to Pythia-41 OM. DEPT models come close Pythia- -410 410M despite the latter being trained on 10x more tokens, indicating a comparable if slightly worse transformer body. Pythia-- a-thiia--i 410 OM was trained on Ubuntu IRC (UI), , thus its outperformance is expected as it is not an OOD dataset for this model.
Name
PP
WK
AX
UI-0OD
AVG
(UNIGRASCC)
(0.9)
(7.9)
(11))
(7.8)
(7.9)
(7.9)
(8.2)
(7.7)
(9.1)
(8.2)
(7.7)
(7.9)
(8)
(8.1)
(7.7)
(7.7)
(10)
(8.1)
PYCHEA-440M
4.9
25.9
43.3
17.4
5.1
45.4
24.7
13.9
65.7
31.7
9.8
18.8
13.5
31.3
44.5
58.3
K1.2
50.3
GLOB
4.8
25.7
38.2
17.3
5.4
47.7
25.7
14.7
68.5
52.7
9.9
20
14
32.2
46.5
30.8
91.8
31.6
TRIM
4.8
27.3
39.5
18.5
5.6
51.2
27.8
18.4
71.8
35.1
10.3
21.7
14.8
35.1
49.1
42.2
96.7
33.3
SPEC
4.8
26.7
36.8
18.2
6.5
50.1
27.1
16.1
69.1
34.2
10.1
21.1
14.5
34.3
48.5
41.7
97.6
32.7
SPEC-OPT
25.0
17.5
5.4
48.3
26.1
14.7
06.8
52.8
9.9
20.4
14.1
32.9
47.5
40.5
88.6
31.2
Table 20: Validation perplexity (L) for 24-block models trained on The Pile after continued pre-training with proportional sampling from randomly-inttialize embeddings, compared to Pythia-41 OM. Pythia-A OM significantly outperforms DEPT as its 30x larger number train- ing tokens allow it to train much better embeddings. Pythia-4 OM trained on Ubuntu IRC (UI), thus its outperformance is expected as it is not an OOD dataset for this model.
Namse
PP
AVG
(UNIORMMC))
(6.9)
(7.9)
(10)
(7.8)
(7.9)
(7.9)
(6.2)
(7.7)
(9:1)
(8:2)
(7.7)
(7.8)
(5)
(8.1)
(7.7)
(7.7)
(10)
(8.1)
DYCHIA-411 10PHIA-410N
3.8
9.7
8.9
7.7
11.8
7.2
21.3
12.5
6.9
10.3
7.6
15
17.6
16.9
7.8
10.9
GLOB
4.5
17
16.1
13.2
45
31.5
17.0
11.3
37.8
22.4
8.4
18.4
11
20.6
35.5
28.3
61.2
21.1
TRIM
48
205
1891
4.
38
2022
12
49.9
26.1
87
16.6
1188
26.7
38
32.9
568
237
Tables 19 and 20 show the expected outperformance of the 410M Pythia model over the DEPT models, as this size has sufficient capacity to benefit from the extensive (10x longer compared to DEPT) pre-training. When starting from random initialization, Table T9 the best DEPT vari- ant within 1 average perplexity point of Pythia-44 a-1 0M, indicating that a large portion of the additional token budget is primarily used obtain better embeddings without providing sig- nificantly improved transformer body. When starting from pre-initialized embeddings, Table 20) Pythia-41 OM significantly outperforms DEPT achieving an average perplexity points lower than the best DEPT variant. As discussed above, this is driven by its more extensive pre- pne-training and improved embeddings.
C APPLICATIONS
C.1 FEDERATED TRE-TRAINING OF LLMS ON MULTILINGUAL POPULATION
The challenges of training under data heterogeneity have come back into focus with recent forays into federated pre-trale pre-training (Douillard et al.) 2023 Sani et al.) 2024 Charles al. 2023 Nous Research] 2024). triggered in equal parts by privacy concerns, compute sharing and the search for more data in previously untapped reservoirs.
27

---

Published as a conference paper at ICLR 2025
The way in which datasets are curated, filtered, and combined has significant impact on the per- formance of LLMs (Long et al. 2024). Determining the best methods for data curation, filtering, and mixing from various sources requires extensive experimentation to identify configurations that optimize performance on target evaluation metrics (Meta) 2024). Consequently, the specific details of these processes are often closely guarded by leading LLM developers. Despite careful dataset preparation, data heterogeneity remains inevitable due to the inherent imbalance in data sources. One of the most prominent imbalances is in language representation. For instance, only about 5% of the pre-training data for Llama3 is -on-English, covering over 30 languages, which results in lower expected performance in non-English contexts (Meta (024). similar performance disparity across languages has also been observed with GPT-4 (Achiam et al. 2023
Current datasets used for pre-training are highly geographically concentrated to a few areas of the globe (Faisal et al. 2022), providing the so-called high-resource languages, with high-quality domain-specific data being available predominantly in such languages (Magueresse et al.) 2020).
Such datasets are collected from internet sources and then curated (Brown et al. 2020 Dubey al. 2024). However, bottlenecks in the rate of high-quality data generation Viilalobos et (0.2)) and copyright concerns (Grynbaum & Mac 2023) have led to large organizations making deals with
private data providers such as publishers (OpenAl 2023 Patel & Palazzolo| 2024) in order to meet
the demand of ever-growing models.
Federated pre-training as a methodology allows the model to be taken directly to the training data, potentially enabling training under privacy concerns or legislation that limits data move- ment (Woisetschläger et al. 2024). While this has obvious applications for collaborative training LMs, can also be applied by a single organization as a drop-in replacement for mini-batch SGD during pre-training (Douillard et al.) 2023), which eliminates dataset movement while massively lowering the communication frequency of model training compared to Data mata-parallel algorithms (Ra- jbhandari et al.) (2200) which need to synchronize gradients every batch. The version of the algorithm used a centralized setting, mathematically equivalent to Federated Averaging (McMahan et al.) 2017a), has alternatively been known as: (a) communication--fficiii
SGD (Yu et al. (019)), (b) Lo- cal SGD (Stich) 2019) Ortiz et al., 2021), or (c) as a specific variant of the REPTILE (Nichol et al.) 2018) meta-learning algorithm. Under these various methodologies, it has been shown to (a) confer a linear speedup to convergence similar to increasing batch size, (b) provide better generalization to models compared standard large- batch training, (c) enable meta-learning across various tasks.
While the current centralized pre-training recipe may be stabilized with great effort, such measures largely impractical in federated training scenarios where the participants refuse to offer full control over their data to a third party or where the underlying training distribution may shift as new participants enter federation old ones exit. Furthermore, the complexity of the current pipeline is impractical to all except the best-funded organizations, even in centralized training context.
The inability to directly inspect data sources federated context makes it impossible construct a dedicated vocabulary for a data mixture, ensure a standard curation pipeline on a per-sample basis, or strongly control data sampling rates across all sources. Motivated by this extreme setting, we aim to construct pre- pre-training procedure that is capable of learning from multiple highly hetero- geneous data sources without model divergence while providing a foundation model with greater generalization and more plasticity in adapting to new data.
D TRAINING UNDER DATA HETEROGENEITY
Training LLMs such as Llama 3 (Dubey et al. 2024) requires extensive manual tuning. heuristics, and model-based data selection procedures. This effort aims to achieve the desired mix of categories, such as general knowledge, mathematics, coding, and multilingual data.
This complexity arises due to the wide range of capabilities required by LMs and the risk of neg- ative interference across domains and languages. Current pre-training methodologies are prone to divergence unless data sampling ratios can be meticulously curated based on the characteristics of the data and its fit to the model's 's distribution at any given time (Dubey et al., 2024). Multi-domain ratios manually curated for downstream performance, requiring extensive and expensive tuning, while multilingual pre-training often employs temperature- weighted sampling (Deylin et al., 2019 Conneau et al. 2020 Xue et al. 2021) due to the vast number of languages involved,
28

---

Published as a conference paper at ICLR 2025
As illustrated in Fig. pre-training on heterogeneous data result in model activation diver- gence (Hoffmann et al. 2022). even with a sampling temperature of r = 1.0, which corresponds to proportional sampling based dataset size. Activation divergence is a precursor to significant, often irrecoverable, increases in loss, and necessitating model re-starts from earlier checkpoints with lower learning rates (Zhang et al. [022). Longer training durations could be achieved by dispro- portionately sampling from larger, lower-quality datasets like C4 or high-resource languages like English in multilingual pre-training. Alternatively, methods like active forgetting via embedding resetting (Chen et al. 2023), ACT, may artificially extend the training duration past the natural di- vergence point.
Previous studies show that this Curse of Multilinguality and/or Negative Interference can be at- tributed to vocabulary dilution and capacity contention (Conneau et al. 2020), language-specing
parameter emergence (Wang et al., 2020), and suboptimal tokenization (Rust et al. 2021 Increas- ing model and vocabulary size helps capacity contention (Conneau et al.) 2020; Wang et al., 2020), but this requires immense hardware resources (Dubey et al. 2024) to shard the model across multi- ple GPUS, Addressing vocabulary dilution highly multilingual models even more challenging, providing enough tokens for all languages would result in impractically large models (Rust al., 2021]. These limitations drive us to find scalable methods to incorporate broader data mixtures without significantly increasing the in-memory model size during training.
E FINE-TUNING DEPT MODELS
We evaluate fine-tuning performance three downstream tasks: RACE, MNLI, and STSB. All models are fine-tuned using the recipes provided by |Radford et al. (2018) for each task using the AdamW optimizer with linear learning rate scheduler. For RACE, the model trained for 5 epochs with learning rate of 6e-5 and a batch size 16. For MNLI, fine-tuning is performed over 2 epochs with learning rate of 4e-5 and batch size of 32. Finally, STSB is fine-tuned for 5 epochs using learning rate of 2e-5 and a batch size of 32. The results are reported in Table 21
Table 21:
The performance on downstream tasks (t), following continued pre- pre-training, shows that DEPT models achieve 3% 7.5% relative improvements over the baselines, with TRIM deliver- ing the best results. DEPT consistently outperforms baselines, even with pre-trained embedding initialization, underscoring the importance of an effective transformer body.
Random Init
Pre-trained Init
RACE (ACC)
MNLI
STSB
SST2 (ACC)
RACE (ACC)
MNLI
STSB
SST2
Name
(ACC)
(PC)
(ACC)
(PC)
(ACC)
0)
0.5
0.6
0.66
0.79
0.5
0.71
0.74
0.81
STD (T 1)
0.46
0.68
0.73
0.81
0.53
0.7
0.76
0.83
ACT
0.45
0.66
0.73
0.8
GLOB
0.51
0.72
0.78
0.83
0.51
0.69
0.76
0.82
TRIM
0.53
0.71
0.78
0.83
0.55
0.73
0.81
0.86
SPEC
0.52
0.71
0.79
0.81
SPEC-OPT
0.51
0.69
0.77
0.85
Min Imp (%) 2.9%
4.6%
5.9%
0.7%
- 3377 3.2%
- 33.2%
0.5%
-1.8%
Max Imp (%)
5.8%
6.1%
7.5%
4.1%
3%
6.6%
3.2%
F USING SPEC MODELS FOR INFERENCE
SPEC models do not inherently support inference on a broad corpus after initial pre-training. Suppose local vocabularies and embedding matrices are available without privacy concerns. In that case, inference can be performed using the embedding matrix of the broadest data source or the one closest to the target application. For instance, target- ing English text would utilize EN embeddings for MC4 CC embeddings for The Pile. While effective, this limits generalization beyond the broadest dataset in the pre-training distribution.
As discussed in Sections 2.4 3.5 and 6.1
To handle corpus resembling mixture of all pre-training data sources or unseen ones, SPEC models require a broader embedding matrix for good performance. This can be achieved through multi-phased adaptive continued pre-training, starting with random embedding matrix or the
29

---

Published as a conference paper at ICLR 2025
broadest pre-training one, demonstrated in this work. Alternatives include vocabulary/embeeddnn transfer (Remy et al., 2024) or vocabulary matching (Xu et al., 2024). If these methods fail to reach the desired performance, additional optimization may be necessary to align the embeddings with the transformer body.
30


Reviews:
[{"review_id": "lztSK9wXgB", "reviewer": "Reviewer_GATh", "summary": "The authors introduce the method DEPT, “Decoupled Embeddings for Pre-trained Language models,” an algorithm that trains a transformer model on heterogeneous datasets such as domains or languages individually and subsequently aggregates the parameters by averaging them. They introduce three variants where increasingly less embedding information is shared. DEPT achieves lower perplexity than baselines.", "questions": "1. How would the model DEPT-SPEC be used for inference? Would there first be a need to determine the used domain / language as tokenizers are separate (L129)?\n2. The GLOB model seems to work best. Is there any hypothesis why this is the case? What is the implication of this fact?\n3. The mentioned “Performance metrics” in all captions is always perplexity?\n4. Based on which estimation is communication cost 675x lower?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. The authors tackle an important problem. The usage of data mixtures during pre-training is not well understood but is an essential part of modern foundation models.\n2. While the idea of using model averaging after an inner loop of training on dedicated subsets of data is not particularly novel, it might have a big impact on pre-training, given the encouraging results.", "weaknesses": "1. Writing can be improved or misses important information. For example, for the experimental setup, I struggle to understand L205-311, and information on software/hardware, such as how many FLOPS  or hours training took, is missing.\n2. Some claims are overstated: M-T outperforms DEPT in 5/11 datasets in Table 2. I am not convinced that Trim and Glob perform identically (L377).\n3. An important additional baseline would be models trained on individual data sets. This would give insights into the advantages/disadvantages of model averaging.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "6ivDmuwHPk", "reviewer": "Reviewer_vuvz", "summary": "In this work, the authors propose to train Language Models with multiple independent tokenizer and language modeling heads (one per language or data source) but with a single Transformer backbone. One touted advantage of this approach is that each source can train on its own GPU node, with a reduced need for cross-node communications (since a significant share of the weights does not need to be replicated across them). Another touted advantage is that forcing one model to work with multiple tokenization schemes makes the model more adaptable to new languages and domains post-training, similar to how \"embedding reinitialization during the training\" was found to have a similar effect (at a cost that however did not warrant its use when training large models, unlike the proposed methodology that does not requiring wasting training cycles).", "questions": "- [addressed] Is there any comparable experiment whose results you can compare your models against?\n- [addressed] How did you perform hyperparameter search?\n- [addressed] Did you perform scaling experiments?\n\n[addressed] More than answering questions, I expect the authors to do a significant effort during the rebuttal to improve the paper's form. I would downgrade my rating if this does not happen. (I prefer to give authors a chance to amend, because I think the paper is otherwise interesting if you can get past the shortcomings of the writing.)", "limitations": "", "rating": 8, "confidence": 5, "soundness": 4, "presentation": 2, "contribution": 4, "strengths": "- The setup proposed in this paper looks very satisfying, and it seems to solve several problems both in the industry and in research labs.\n- The value proposition seems clear to me.\n- The deployed methodology appears novel.\n- The literature research looks satisfactory to me, given the scope of the paper.", "weaknesses": "- [addressed] The paper's form is well below the required writing standards. To address this, I'd suggest specific improvements, such as:\n  - Standardizing method names throughout the paper and tables (SPED vs SPEC, GlOB vs GLOB vs Glob, ...)\n  - Clearly defining the performance metrics used and specifying explicitly whether lower or higher values are better\n  - Adding a reference to Table 1 in the main text\n  - Improving table readability by adding summary statistics (averages...), using bold or color highlighting, or splitting into multiple tables (moving some languages to an Appendix).\n- [addressed] Not enough arguments are brought forward to justify the issue with diverging models during training. I have myself never experienced this phenomenon in similar setups. As such, it's difficult to rule out that it might be the result of bugs in the training code or poor hyperparameter choices, rather a general phenomenon. \n  - A better description of the exact training methodology and the hyperparameter search would help alleviate concerns, here.\n  - An ablation study or an explanatory paragraph isolating factors that contribute to divergence would also help.\n- [addressed] The lack of comparisons with baselines not trained by the authors is worrying. \n  - I would prefer for external baselines to be added, even if some added context is necessary to explain away unfair comparisons (could be an appendix).\n- [addressed] Without devising a clear methodology to perform inference on SPEC-type models, the paper feels a bit incomplete. \n  - I'd suggest to the authors to briefly outline a proposed inference methodology for SPEC models, and to discuss the challenges and potential approaches for inference with these models in more detail.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ZaB6XIGIMO", "reviewer": "Reviewer_xZ5j", "summary": "This paper proposes a pre-training framework DEPT which allows the model to train without being bound to a shared global vocabulary. Three variants are introduced: Glob, Trim, and Spec. Although each pipeline has the same setting for the transformer block, the difference lies in how they deal with the embeddings. Among the three, Glob is close to the standard pre-training: a single global embedding matrix is used. Trim keeps a local embedding matrix for each data source, and each token in the matrix is also contained in the global vocabulary. With the federated learning framework, the updates of a specific token are then aggregated to the same token in the global embedding matrix. Spec is a fully decoupled version where there is a non-sharing local embedding matrix for each data source. The authors evaluate the DEPT by investigating the efficiency, generalization as well as plasticity.", "questions": "- abstract: I believe \"negative interference\" and \"curse of multilingual\" are not parallel concepts. Curse of multilingual is a special type of interference because of the training data containing many languages. \n\n- I don't find the authors clearly stating the \"performance metrics\" in Tables 1, 2, 3, 4.\n\n- The authors are suggested to have aggregated statistics for Tables 1, 2, 3, and 4 (e.g., average) so that the authors can have a better understanding of which model generally performs the best. Additionally, the authors can bold the best number in each column.\n\n- line 278: 50,257 instead of 50257.0\n\n- Line 341 (b) tn the ?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "- The paper is well-written and easy to follow.\n\n- The idea of decoupling embedding matrix and transformer block in pre-training within the federated learning framework is novel.\n\n- The authors answer the raised research questions with meaningful and extensive experiments.\n\n- The results generally confirm that DEPT can improve the generalization and plasticity of the models.", "weaknesses": "- The data sources are not always clear given a dataset. The proposed pipeline only works if the domains are known. Otherwise, some manual or automatic clustering has to be used to create different sets of data.\n\n- The multi-domain data is almost only in English. But for the multilingual data, the data of each language should also contain various domains. Therefore there are confounding variables. A natural question would be whether the model can generalize to the same domains across different languages.\n\n- No downstream tasks in natural language understanding or generation are evaluated on the resulting models. But such further evaluation is important.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Oral)