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
Primary area: foundation or frontier models, including LLMs
Keywords: decentralized training, federated learning, multi-domain training, multilingual training

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.

Decision: accept (oral)
Reviewer scores: 8, 8, 8

Reviews and discussion:
<thread>
creation date:
1729081368635

thread_type:
official_review

<thread_object>
creator: reviewer_xz5j

<content>
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.

strengths:
- The paper is well-written and easy to follow.

- The idea of decoupling embedding matrix and transformer block in pre-training within the federated learning framework is novel.

- The authors answer the raised research questions with meaningful and extensive experiments.

- 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.

- 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.

- No downstream tasks in natural language understanding or generation are evaluated on the resulting models. But such further evaluation is important.

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. 

- I don't find the authors clearly stating the "performance metrics" in Tables 1, 2, 3, 4.

- 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.

- line 278: 50,257 instead of 50257.0

- Line 341 (b) tn the ?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730138742397

thread_type:
official_review

<thread_object>
creator: reviewer_vuvz

<content>
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).

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.
- The value proposition seems clear to me.
- The deployed methodology appears novel.
- 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:
  - Standardizing method names throughout the paper and tables (SPED vs SPEC, GlOB vs GLOB vs Glob, ...)
  - Clearly defining the performance metrics used and specifying explicitly whether lower or higher values are better
  - Adding a reference to Table 1 in the main text
  - Improving table readability by adding summary statistics (averages...), using bold or color highlighting, or splitting into multiple tables (moving some languages to an Appendix).
- [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. 
  - A better description of the exact training methodology and the hyperparameter search would help alleviate concerns, here.
  - An ablation study or an explanatory paragraph isolating factors that contribute to divergence would also help.
- [addressed] The lack of comparisons with baselines not trained by the authors is worrying. 
  - 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).
- [addressed] Without devising a clear methodology to perform inference on SPEC-type models, the paper feels a bit incomplete. 
  - 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.

questions:
- [addressed] Is there any comparable experiment whose results you can compare your models against?
- [addressed] How did you perform hyperparameter search?
- [addressed] Did you perform scaling experiments?

[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.)
</content>
</thread_object>

</thread>
<thread>
creation date:
1731188691922

thread_type:
official_review

<thread_object>
creator: reviewer_gath

<content>
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.

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.
2. 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.
2. 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).
3. An important additional baseline would be models trained on individual data sets. This would give insights into the advantages/disadvantages of model averaging.

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)?
2. The GLOB model seems to work best. Is there any hypothesis why this is the case? What is the implication of this fact?
3. The mentioned “Performance metrics” in all captions is always perplexity?
4. Based on which estimation is communication cost 675x lower?
</content>
</thread_object>

</thread>
<thread>
creation date:
1731784524145

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Rebuttal

comment:
# Dear Reviewers, ACs, and PCs,

Thank you very much for your dedication, support, and insightful feedback. We are delighted you have found our work and experimental design **insightful** and **sound** [vuvz], addressing an **important** [GATh] and **novel** [xZ5j, vuvz] problem by decoupling the embedding matrix and transformer body. Your constructive comments have immensely helped us improve our paper's quality. We have reviewed all the comments, addressed all questions, and provided new experimental results. Below, we summarize the revisions we made:

---

## Revisions and Updates

### Experimental Results

- **Additional experimental results. A total of 38 new experiments added:**
    - [GATh] DEPT variants win **82.2% of experiments** in terms of perplexity and performance on downstream tasks.
    - [xZ5j] Fine-tuned DEPT models on tasks like **Natural Language Inference**, **Question Answering**, and **Sentence Similarity**, outperforming baselines in all configurations (**Table 7**).
    - [GATh] Studied sampling ratio choices during pre-training on **The Pile** dataset:
        - Switching from uniform to proportional sampling improved results (**Tables 3** and **5**).
    - [GATh, vuvz] Initiated **26 additional experiments** (results forthcoming):
        - Baselines where local models are trained on individual data sources.
        - Comparisons with externally trained models.

### Writing Standards and Clarity

- [GATh, vuvz, xZ5j] Improved clarity of methods and experiments by providing detailed explanations and further context.

### Presentation of Results

- Enhanced **tables**, **figures**, **captions**, and **descriptions** to improve readability and presentation consistency. Perplexity values too large to fit in tables are now denoted by a “-“.

### Updated Workflow Illustration

- Modified **Figure 1** to better showcase DEPT’s workflow and compare its variants to standard pre-training pipelines.

---

## Section-by-Section Walkthrough of Revisions

**Revisions are highlighted in blue in the manuscript.** Key changes include:

- **Abstract:**
    - Added figures illustrating DEPT’s improved perplexity and downstream fine-tuning performance.

- **Section 1:**
    - Provided numerical results for DEPT’s **generalization performance**.
    - Detailed the **memory footprint** and **communication efficiency** in asymptotic terms.

- **New Section 2.2:**
    - Formally described the **TRIM variant trimming procedure**.

- **Section 2.4 (previously Section 2.3):**
    - Clarified inference steps for **SPEC-trained models**, with further details in **Section 3.5** and **Appendix F**.

- **New Table 1 (Section 2):**
    - Summarized **asymptotic complexities** for memory and communication costs.

- [GATh] **New Table 2 (Section 4):**
    - Addressed scalability and communication efficiency (**714x improvement**):
        - **500x reduction** from infrequent averaging.
        - **30% reduction** from excluding embedding parameter communication.

- **Appendix A.1:**
    - Detailed **hyperparameter choices** and **sampling strategies** for continued pre-training.

- [GATh] **Appendix A.1.1:**
    - Outlined **software/hardware configurations** and provided an anonymous GitHub repository with DEPT execution code.
    - Reported hardware throughput and time differences between DEPT and baseline methods.

- [vuvz] **Appendix A.1.2:**
    - Detailed the **hyperparameter tuning methodology** for both DEPT and baselines.

- **Appendix A.1.3:**
    - Adapted active forgetting methodology for decoder-only models.

- **Appendix A.2:**
    - Additional statistics on data source heterogeneity, including unique token counts.

- **Appendix A.2.1:**
    - Described the unigram normalized cross-entropy, presented in **Tables 3–6**, as a reference metric for data source difficulty.

- **[GATh, vuvz, xZ5j] Appendix B:**
    - Presented results on:
        - **Larger models** (**B.1**).
        - **Model plasticity** (**B.2**).
        - **IID data performance** (**B.3**).
        - **Generalization** (**B.4**).
        - **Scaling from smaller models** (**B.7**).

- [xZ5j] **Appendix E:**
    - Added fine-tuning parameter details.

- [GATh, vuvz] **Appendix F:**
    - Presented multiple inference methods for SPEC.

---

## Request for Feedback

We kindly ask the reviewers to inspect our revisions and individual responses. We are happy to address any remaining questions. If our responses and additional results adequately address your concerns, we kindly ask you to consider increasing your scores. We deeply appreciate your engagement and suggestions, which have been invaluable for improving this work.

**Best regards,**

*The Authors*
</content>
</thread_object>

</thread>
<thread>
creation date:
1732618750928

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Rebuttal Reflections

comment:
# Dear Reviewers, ACs, and PCs,

We sincerely thank the reviewers for their time, constructive feedback, and engagement throughout the discussion phase. Your insights have been invaluable in improving the clarity, presentation, and depth of our submission. We are thrilled to hear that our revisions and additional experiments addressed your concerns, and we deeply appreciate the positive updates to your evaluations.

### **Discussion Highlights**

- Reviewer **GATh** initially recognized our work as **tackling an important and novel problem** while providing a **method that could significantly impact pre-training methodologies**. After reviewing our detailed responses, improved presentation, and additional experiments, the Reviewer updated their score to 8, highlighting that **our updates adequately addressed their concerns** and endorsed the acceptance of our paper with **high** confidence.
- Reviewer **vuvz** praised the **clear value proposition and novelty** in their initial feedback. After seeing our improvements in presentation, hyperparameter search explanations, and inference methods, they affirmed their support, stating, **“Your detailed answer was a pleasure to read”,** and maintained their score of 8 with firm support for acceptance with **very-high** confidence.
- Reviewer **xZ5j** commended the paper for being **well-written, easy to follow, and methodologically novel**. Our response, particularly the downstream task evaluations, and clarifications, led them to declare our work as **“a good paper worth accepting”** and increase their score to 8  with **high** confidence.

### A complete set of new experimental results

Throughout this rebuttal, we have provided a new set of experimental results based on the reviewers' feedback. The first set accompanied revision 1 of the paper and included 26 additional experiments covering downstream tasks[**xZ5j**] and continued pre-training sampling policies. We provide the second set of results now in revision 2, covering 51 new experiments, including comparisons against the externally-trained Pythia 160M and 410M baselines[**vuvz**], an additional downstream task[**xZ5j**], and comparisons against new baselines trained on isolated data sources without model averaging[**GATh**].  For details on both sets of experiments, please look at our previous [general rebuttal comment](https://openreview.net/forum?id=vf5aUZT0Fz&noteId=d7Krp7Zo3V) and its [reply](https://openreview.net/forum?id=vf5aUZT0Fz&noteId=GLNSHb0pot).

### Impact of DEPT

We envision DEPT as the foundation of a **new paradigm for pre-training language models**. This framework **uses decoupled pre-training on heterogeneous data** to obtain a more generalizable and adaptive/plastic transformer body than the one produced by standard pre-training methods.   

The implications of this paradigm are:

- **Flexibility:** Practitioners can train on heterogeneous and even private data sources without the complexities of creating a shared vocabulary.
- **Efficiency:** The reduced communication and memory costs allow pre-training across geographically distributed GPUs, enabling a wider-scale deployment than standard pre-training.
- **Generality:** DEPT-trained models can serve as versatile foundation models that perform well across various tasks and domains. Compared to standard pre-training, this improved foundation model can enhance the efficacy of continued training and fine-tuning.

This approach not only simplifies training pipelines but also extends the applicability of pre-trained models to settings where diverse, distributed, and sensitive datasets must be leveraged.

### Conclusion

We are delighted that the reviewers now unanimously consider our paper to be **worth accepting** and appreciate the value and novelty of DEPT. We hope DEPT’s contributions and vision resonate with the community as a step forward in multilingual and multi-domain model pre-training.

Thank you again for your support, organization, and invaluable feedback. 

Best regards,

The Authors
</content>
</thread_object>

</thread>


Metareview:
metareview:
This work proposes the DEPT algorithm (Decoupled Embeddings for Pre-Trained language models) to train LMs on diverse data sources more effectively. The method iteratively trains a transformer model on different subsets of data with separate embeddings for each data type (such as domain or language); the work proposes and compares three variants of this approach. The experiments show that this approach performs well for continued pretraining in both different English domains and in the multilingual setting, and the analyses also demonstrate the method's efficiency.

Strengths:
- This paper addresses an important problem (tokenization and embedding representations over diverse data sources) and provides a compelling approach to solving this problem.
- Results are (for the most part) consistent and convincing for both the perplexity and downstream task experiments.
- The approach is more efficient than standard pretraining, as cross-node communication isn't needed during the steps where the model is trained independently on each data type (InnerOPT).
- The paper provides clear research goals beyond task performance that the experiments and analyses address, such as model efficiency and generalization.


Weaknesses: 
- Multiple reviewers raised concerns with the writing, and some details are still confusing after revisions (e.g., defining STANDARD as the continued pretraining baseline in Table 2 and then referring to it with the names based on sampling methods in Tables 3 and 4). The authors should ensure the notation and discussion are consistent and clear in the next version of the paper. This also includes other points raised by the reviewers and clarified in the rebuttal.
- Table 7 is too small to read easily and should be larger in the next version.