Dear Reviewer xZ5j,

We express our gratitude for the valuable feedback of their review and for acknowledging our proposal's novelty, quality, and soundness. Please see our detailed responses below:

• RE: “Clarification on the data sources”
  We thank the reviewer for raising this interesting point. One of our assumptions is, indeed, that data sources are partitioned into different domains or languages, as they are for The Pile [1] and MC4 [4] datasets that we use. We believe that this is not a strong assumption.

  • First, in many scenarios, this assumption holds: in Federated Learning settings, in Multilingual datasets, and Multi-domain datasets, the partitions are already identifiable without additional manual or automatic clustering.
  • Second, our experimental results show that, in the case of IID data shown in Section B.3, DEPT performs no worse than standard pre-training.

• RE: “Investigating multiple domains for multiple languages simultaneously”
  This particular context pointed out by the Reviewer represents a very interesting setting for DEPT.

  • Our experimental space is indeed limited to investigating multiple domains and multiple languages separately. This limitation arises from, to the best of our knowledge, the lack of open-source datasets allowing for such experimentation.
  • The construction of such datasets, which would require gathering in-the-wild multi-domain data in multiple languages, needs careful consideration and extensive efforts to make such text corpora usable.
  • We expect our method to generalize to the same domain across different languages. However, we would be very excited to evaluate DEPT on open-source datasets of this kind if the Reviewer can kindly point them to us.

• RE: “Evaluating the downstream performance of DEPT models”
  Thanks to the Reviewer's comment, we acknowledge the importance of such an evaluation and provide an additional evaluation in the revision we posted.

  • We invite the Reviewer to kindly inspect the new revision, especially Table 7, where they can find such an assessment.
  • We evaluate the natural language understanding tasks: Natural Language Inference (MNLI), Question Answering (RACE), and Sentence Similarity (STSB) and show that DEPT wins all of the comparisons, bringing improvements in accuracy (RACE, MNLI) and Pearson correlation (STSB), with relative improvements of 3%-7.5% over the baselines.
  • For such comparisons, TRIM performs best with GLOB and SPEC on equal footing, proving further advantages to our novel embedding decoupling methods beyond memory and communication efficiency and language modeling performance.
  • If the Reviewer has any other recommendations for downstream tasks, we are happy to incorporate them.

• RE: “Negative interference and curse of multilinguality not being parallel concepts”
  We acknowledge that negative interference is a superset of the curse of multilinguality and encompasses a much broader set of issues.

  • The parallels we draw between the two are based on the work of [3] and are meant simply to connect why our methods are aimed at both multi-domain and multilingual data, despite highly heterogeneous subsets such as mathematics not being generally perceived as a fully different language.
  • Our new abstract now states that heterogeneous data sources suffer from “challenges such as” negative interference or the curse of multilinguality to avoid making them seem equivalent.

• RE: “Identifying performance metrics, adding aggregated statistics, and bolding the best number in each column”
  We have fixed table formatting issues, clearly identified every performance metric, and made the paper presentation consistent.

• Thank you so much for catching the typos.
  We have fixed them (Line 246, Line 279 in the original submission).

We are firmly committed to improving our submission, having uploaded a new revision that addresses most of the Reviewer's comments, and we are in the process of completing the additional evaluations we have been kindly asked to produce.

We intend to progressively improve our work, and we kindly ask the Reviewer to raise their scores if they believe our rebuttal addresses their concerns in light of our additional downstream task results, clarifications, revisions, and improved presentation.

References

[1] Gao et al., The pile: An 800GB dataset of diverse text for language modeling.
[2] Sani et al., The future of large language model pre-training is federated.
[3] Wang et al., On Negative Interference in Multilingual Models: Findings and A Meta-Learning Treatment.
[4] Xue et al., mT5: A massively multilingual pre-trained text-to-text transformer.

Kind regards,
The Authors

Dear Reviewer vuvz,

We genuinely appreciate the Reviewer's thorough assessment of our submission, which recognizes the soundness, novelty, and careful literature review shaping the foundation of our proposal. Their review is very insightful and has improved our submission. Below we elaborate on the following comments:

• RE: “Writing standard”
  We greatly appreciate the reviewer’s valuable suggestions. In the revised version of the paper:
  • We have fixed and applied all the suggestions.
  • Method names are made consistent, and performance metrics are clearly named in the tables.
  • Arrows indicating whether greater or lower values are better have been added.
  • All tables are now referenced in the main text.
  • Tables report the average value as a separate column, highlight the best performance in bold, and include the improvement (%), as described in the captions.

• RE: “Model divergence”
  Thank you for raising this important point. The phenomenon of model divergence during training has been documented in several works (e.g., [1, 2, 4]), with recent theoretical insights provided in [6]. Divergence is often preceded by activation instability and can occur even in smaller transformers under specific conditions, such as:

  • Unsuitably high learning rates [4, 6]
  • Data heterogeneity [1]

  Widely accepted solutions to divergence include:

  • Extensive hyperparameter tuning
  • Gradient clipping
  • Expensive techniques like using “a token-distribution Kullback-Leibler divergence to filter out documents containing excessive numbers of outlier tokens compared to the training corpus distribution” [1].

  DEPT mitigates this issue significantly, offering a more robust and cost-effective framework for model development.

  To address the Reviewer’s concern:

  • We have expanded Section 3.4 to discuss contributing factors to divergence and their mitigations.
  • All these explanations are added as per the reviewer’s suggestion.

  Code-related bug concerns:

  • Our training pipeline is built on the rigorously tested MosaicML Composer library for LLM pre-training and the Flower framework for federated learning.
  • For our InnerOPT implementation, we rely on MosaicML's hyperparameters and infrastructure, modifying only the embedding matrix as per Algorithm 1.
  • Standard baselines are trained using the unmodified MosaicML codebase, which is validated by wide adoption and peer-reviewed publications [7].
  • Our codebase is available as an anonymous GitHub repo listed in A.1.1.

  Additional clarifications about the training setup and hyperparameter choices have been incorporated into Sections A.1, A.1.1, and A.1.2 of the revision.

• RE: “Hyperparameter search”
  We emphasize the following:

  • For robustness plots, models are compared with the exact same learning rate.
  • For performance comparisons, we always use well-tuned checkpoints, as reported in Table 8 and the new A.1.2 section.

  Details of hyperparameter tuning:

  • MosaicML provides hyperparameter-tuned models on the C4 dataset, and we use their learning rate schedule and number of training steps as a starting point.
  • For DEPT, we find that MosaicML parameters work well since the OuterOpt acts as a regularizer.
  • For standard training baselines, the learning rate is gradually lowered, starting from the one recommended by MosaicML (as shown in Table 8).
    • Learning rates are reduced progressively by $5 \times 10^{-5}$.
    • Starting learning rates range from $6 \times 10^{-4}$ to $2 \times 10^{-4}$, with larger models using smaller learning rates.

  Cosine cycle and validation perplexity:

  • We stop when we find a learning rate allowing us to complete a full cosine cycle for the pre-determined compute-optimal number of training tokens [3] (approximately 20 tokens per parameter).
  • The best-performing checkpoint is chosen based on validation perplexity, across all experiments (reported in Table 8).

  Gradient clipping is always applied with a 1.0 norm threshold.

  While this hyperparameter search does not cover all possible parameters (e.g., gradient clip norm, batch size), we are happy to incorporate any suggestions.

• RE: “External baselines”
  Thanks for this excellent suggestion.

  • Results will be made available in the appendix with sufficient context for interpretation as soon as experiments finish.

  Key points:

  • Our goal is to propose a novel training paradigm, not to surpass the SOTA on downstream evaluation.
  • Unlike SOTA small models, our models have not been trained on hundreds of billions of tokens from carefully curated data sources.

Dear Reviewer GATh,

We thank the Reviewer for their valuable feedback and for acknowledging that our submission tackles an important problem and shows promising results. We now address the rest of the Reviewer’s comments and requests:

• The Reviewer raises an excellent point regarding the origins of DEPT in parameter averaging methods. Despite its origins, DEPT significantly extends past traditional FL and parameter averaging methods, where identical models with the same objective function are trained in parallel and then aggregated:
  • TRIM modifies the optimization objective by reducing the vocabulary to each data source’s specific data:
    • Loss Domain: The cross-entropy loss is computed over a restricted token embedding space derived from the TRIM-specific vocabulary.
      • Probabilities for a token at time-step $x_t$ are calculated as:
        • $p(x_t \mid x_{1:t-1}; \phi_{\text{TRIM}}) = \text{softmax}(\mathrm{score}(h_{t-1}, \phi_{\text{TRIM}}))$
        • Where $x_{1:t-1}$ is the sequence of previous tokens and $\text{score}$ is a function which outputs a scalar summarizing the relation between the last hidden state $h_{t-1}$ and every token embedding in the $\phi_{\text{TRIM}}$ embedding matrix
      • The cross-entropy loss on data source $k$ is:
        • $\mathcal{L}^k_{\text{TRIM}} = - \sum_{t=1}^n \log p(x_t \mid x_{1:t-1}; \phi_{\text{TRIM}})$
  • SPEC extends heterogeneity further, allowing completely distinct tokenization schemes and vocabularies across clients:
    • Autoregressive Task: Identical data may be tokenized into entirely different sequences (with varying lengths) based on each client’s unique vocabulary.
    • Loss Domain: The cross-entropy loss is calculated over a data-source-specific set of tokens:
      • $p(x_t^k \mid x_{1:t-1}^k; \phi_k) = \text{softmax}(\mathrm{score}(h_{t-1}, \phi_k))$
      • The cross-entropy loss for data source $k$ is:
        • $\mathcal{L}^k_{\text{SPEC}} = - \sum_{t=1}^{n} \log p(x_t^k \mid x_{1:t-1}^k; \phi_k)$
  • Thus, both variants significantly depart from standard methods for the scoring function, objective function, and, in the case of SPEC, the data representation itself.

• RE: Writing clarity and missed information

  We apologize for the confusion. An improved description of the use of baselines and metrics is available in Sections 3.3 and 3.4 in the revised manuscript. The software, which we make available, is based on the Flower framework [1] and the MosaicML Composer library [2]. As for the hardware, we use a mixture of cloud-based instances equipped with a range of Nvidia GPUs, e.g., H100 and A100, in several regions (e.g., Europe, US, and Canada). The details are in the revised version (Sec. 3.3, 3.4, and A.1.1), we also make our FL pre-training framework available in A.1.1 in an anonymized form.

• RE: Wall-clock training time

  While DEPT takes the same number of training steps as baselines, assuming the same batch size and number of GPUs, two factors influence time:

  • Worker throughput: For GLOB, it matches standard pre-training as the model is identical. For TRIM and SPEC, reduced memory requirements allow increasing micro-batch sizes in some hardware-dependant settings and hence faster training—e.g.,, doubling it on an H100 for a 24-block model for the small-vocabulary DeepMind Mathematics tasks.
  • Communication topology: On cloud machines with 10 Gbps bandwidth and Ring AllReduce, DEPT reduces training time by 33% for a 1-billion parameter model. On faster connections like InfiniBand, differences are primarily due to throughput.

• RE: M-T performance

  Insightful observation. We note that DEPT wins 82.2% of comparisons overall and hope this figure sufficiently justifies our outperformance claims. For example:

  • We win all comparisons on downstream tasks and a vast majority of the language modeling comparisons.
  • Table 2, in the original submission, shows results starting from a pre-trained embedding matrix, which are sensitive to the sampling used by the baselines during pre-training. Looking at the updated version of Table 2 (Table 6 in the latest revision), we see that DEPT embeddings are slightly worse for high- and medium-resource languages like EN, IT, or SR, or for LA (closely related to EN and IT). However, DEPT wins the comparisons for very low-resource UR and SW languages by large margins.
  • Furthermore, in the case of The Pile (Table 5 in the new version), DEPT still wins 12/17 comparisons under this exact setting, showing that it is not a permanent disadvantage.
  • Since embeddings are much easier to transfer or retrain, we argue that obtaining a better transformer body is generally preferable.

Dear Reviewers, ACs, and PCs,

We have completed another paper revision, addressing all reviewers' requests and providing the promised baselines as detailed in the general comments above.

1. Impact of Model Averaging (GATh's Recommendation)

To evaluate the impact of model averaging, we compared models undergoing continued pre-training from a DEPT initialization with models undergoing continued pre-training starting from models trained on isolated data sources. For a fair comparison, each isolated model sees as many tokens as its DEPT-based counterpart and undergoes continued pre-training for the same number of steps, with access to the full dataset. The results of this comparison are detailed in Section B.8, they show that the model averaging allows DEPT models to generalize better across distributions.

Continued Pre-Training from Random Initialization

• This comparison tests the generalizability of the abstractions learned by the transformer body.
  • [Table 14] The Pile: DEPT-based models outperform the baselines by 7.8% in average perplexity.
  • [Table 15] MC4: DEPT-based models outperform baselines by 9.8% on in-distribution data and 16.7% on out-of-distribution (OOD) data.
• Despite isolated models receiving embeddings suitable for the multi-dataset context through the continued pre-training process, their transformer body abstractions do not generalize well across distributions.

Continued Pre-Training from Pre-Trained Embeddings

• Models initialized from isolated data-source models are biased toward their pre-training data, performing well on specific datasets but failing to generalize:
  • [Table 16] The Pile: DEPT-based models outperform baselines by a remarkable 27% in average perplexity. However, DEPT models tend to be outperformed by isolated baselines for their respective data sources.
  • [Table 17] MC4: DEPT-based models achieve significant outperformance, with a 30.6% improvement in average perplexity for in-distribution data and 14.9% for OOD data. This is likely because embeddings generalize worse across languages than across domains.

2. Comparison Against Externally-Trained Baseline (vuvz's Recommendation)

We compared our models with the externally trained Pythia[1] model suite, which uses a similar architecture but untied weights for the embedding matrix. We chose Pythia as the closest baseline in terms of model structure and pre-training dataset. As a result of the untied embedding weights, Pythia models with the same number of blocks have more parameters. Pythia models are trained on one epoch of The Pile (300B tokens compared to 10B to 30B for DEPT models), thus, they are expected to perform better on Ubuntu IRC than DEPT models since it is not OOD for Pythia. These comparisons have been added to section B.9, they show that DEPT is competitive with Pythia at smaller scales and when starting continued pre-training from random initialization (indicating a comparable transformer body); however, the superior embeddings of the larger Pythia model, pre-trained for 10x more tokens, allow it to outperform DEPT models when using pre-trained embeddings.

160M Model Comparison

• [Table 18] For Pythia-160M vs. DEPT-125M:
  • DEPT models (with pre-trained embeddings) outperform Pythia by 1 point in average perplexity.
  • Pythia's extensive pre-training (30x DEPT's token budget) does not provide an evident advantage due to insufficient model capacity to benefit fully from the additional data.
• Random initialization results for this model size are omitted as Pythia models were unstable during the continued pre-training, making the comparison unfair.

410M Model Comparison

• For Pythia-410M:
  • [Table 19] Random Initialization: DEPT's best variant is within 1 point of Pythia-410M in average perplexity, suggesting that much of Pythia's additional token budget enhances embeddings without significantly improving the transformer body.
  • [Table 20] Pre-Trained Embeddings: Pythia-410M significantly outperforms DEPT, achieving 10 points lower in average perplexity, likely due to training on 10x more tokens than the equivalent DEPT models and thus achieving superior embeddings.

3. Addition of a New Downstream Task (xZ5j's Recommendation)

We added a new downstream task, Sentence Classification evaluated on the SST-2 dataset[2], and reported the results in Table 7. DEPT models outperform all baselines on this task, achieving a 3.2% to 4.1% relative improvement in accuracy.

These experiments have helped significantly improve our work, and we thank the reviewers again for recommending them.

Kind Regards,

The Authors

[1] Biderman, et al., “Pythia: A Suite for Analyzing Large Language Models Across Training and Scaling”

[2] Socher, et al., Recursive deep models for semantic compositionality over a sentiment treebank