Research map / Large language models / Language, safety and interpretability
DEPT: Decoupled Embeddings for Pre-training Language Models
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes DEPT, a communication-efficient pre-training framework that decouples token/positional embeddings from the transformer body (GLOB, TRIM, SPEC variants), enabling vocabulary-agnostic federated pre-training of up to 1.3B decoder-only models.
Problem
Pre-training language models on heterogeneous multi-domain and multilingual data suffers from negative interference, the curse of multilinguality, vocabulary dilution, and high communication/memory costs, especially in distributed or federated settings.
Approach
Decouple embeddings from the transformer body and train the body on data sources with independent (possibly custom) vocabularies, aggregating only the body (and optionally trimmed embeddings) via an outer optimizer like FedAvg, with three variants of increasing specialization.
Key results
- The best DEPT variant improves average validation perplexity by 15.3% on The Pile and 17.3% on MC4 over the best baseline.
- SPEC-OPT reduces the 1.3B model's per-step communication to 0.001x of standard distributed training and cuts embedding memory by 4-5x.
- DEPT variants win 82.2% (51/62) of main comparisons across The Pile, MC4 and downstream tasks; TRIM gives the best downstream results with 3-7.5% relative gains.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| use DEPT variants rather than standard pre-training | transformer body on The Pile · The Pile validation | average perplexity: 36.9 -> 31.2 | Helped |
| use DEPT variants rather than standard pre-training | transformer body on MC4 · MC4 in-distribution validation | average perplexity: 54.9 -> 45.4 | Helped |
| use DEPT variants rather than standard pre-training | transformer body on MC4 · MC4 out-of-distribution validation | average perplexity: 2175.8 -> 1451.9 | Helped |
| trim global token embeddings to each source vocabulary (ablation) | GLOB on The Pile · The Pile validation | average perplexity: 31.6 -> 33.3 | Hurt |
| specialize token and positional embeddings per source (ablation) | GLOB on The Pile · The Pile validation | average perplexity: 31.6 -> 32.7 | Mixed |
| use optimized per-source vocabularies (ablation) | SPEC on The Pile · The Pile validation | average perplexity: 32.7 -> 31.2 | Helped |
| trim global token embeddings to each source vocabulary (ablation) | GLOB on MC4 · MC4 in-distribution validation | average perplexity: 45.4 -> 46.4 | Hurt |
| specialize token and positional embeddings per source (ablation) | GLOB on MC4 · MC4 out-of-distribution validation | average perplexity: 1514.7 -> 1451.9 | Helped |
| trim token embeddings to each source vocabulary (ablation) | GLOB on RACE · RACE | accuracy: 0.51 -> 0.53 | Helped |
| use specialized token and positional embeddings (ablation) | GLOB on STSB · STSB | Pearson correlation: 0.78 -> 0.79 | Helped |
| use optimized per-source vocabularies (ablation) | SPEC on SST-2 · SST-2 | accuracy: 0.81 -> 0.85 | Helped |
| start continued pre-training from pre-trained rather than randomly initialized embeddings (ablation) | DEPT on MC4 · MC4 in-distribution validation | average perplexity: 6.4% improvement in average perplexity | Mixed |
Surprising findings
- Trimmed embeddings (TRIM) perform comparably to GLOB on perplexity despite smaller embedding matrices, suggesting out-of-vocabulary mistakes do not drastically hurt performance.
- With pre-trained embeddings, DEPT does not uniformly beat standard pre-training on perplexity; it sometimes underperforms on OOD MC4 data, which the authors attribute to embedding sampling ratios.
Limitations
- SPEC variants do not produce a global embedding matrix and need an extra continued pre-training phase for inference.
- Pre-trained-embedding comparisons show mixed results, and continued pre-training design must be carefully chosen.
- Experiments are limited to decoder-only models up to 1.3B parameters.
Experimental setup
Datasets: The Pile MC4 C4 MNLI RACE STSB SST-2 Ubuntu IRC
Benchmarks: The Pile validation perplexity MC4 validation perplexity MNLI RACE STSB SST-2
Metrics: perplexity accuracy Pearson correlation communication cost memory cost
Baselines: STD (uniform, temperature-weighted, proportional sampling) ACT (active forgetting)
Builds on: FedAvg Local SGD active forgetting MosaicML composer AdamW
Closest papers by meaning
- Repurposing Language Models into Embedding Models: Finding the Compute-Optimal Recipe (NeurIPS 2024 · similarity 0.90)
- Leviathan: Decoupling Input and Output Representations in Language Models (NeurIPS 2026 · similarity 0.90)
- Test-Time Training on Nearest Neighbors for Large Language Models (ICLR 2024 · similarity 0.89)
- Scaling Embedding Layers in Language Models (NeurIPS 2025 · similarity 0.89)
- Analyzing Similarity Metrics for Data Selection for Language Model Pretraining (NeurIPS 2025 · similarity 0.89)
- Data-Constrained Language Model Pretraining: Improved Regularization and Scaling Laws (NeurIPS 2026 · similarity 0.89)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.