DEPT: Decoupled Embeddings for Pre-training Language Models

TL;DR

DEPT introduces a decoupled embedding framework for pre-training language models that reduces communication costs, enhances robustness to data heterogeneity, and allows custom vocabularies, demonstrated at billion-scale federated settings.

Contribution

The paper presents a novel decoupled embedding approach enabling vocabulary-agnostic federated pre-training with significant efficiency gains and improved model performance.

Findings

Reduces communication costs by orders of magnitude.

Cuts embedding memory by 4-5 times.

Improves perplexity and downstream task performance.

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.