PaperlayerView exact revision on arXiv

Source-linked AI summary

Rethinking embedding coupling in pre-trained language models

Hyung Won Chung, Thibault Févry, Henry Tsai, Melvin Johnson, Sebastian Ruder

arXiv:2010.12821v1cs.CLcs.LG

TL;DR

The paper asks whether tying input and output embeddings unnecessarily constrains parameter allocation in pre-trained language models, particularly multilingual ones. It studies decoupled embedding dimensions and reallocates capacity to embeddings and Transformer layers, finding stronger transfer performance and more general representations without increasing fine-tuning parameters.

  • Problem

    Coupled embeddings force input and output matrices to share dimensionality, limiting parameter-allocation flexibility in multilingual models with large vocabularies.

  • Method

    The paper systematically evaluates decoupled embedding dimensions, reallocates saved input-embedding parameters to Transformer layers, and analyzes transferability across tasks and languages.

  • Results

    RemBERT outperforms XLM-R with the same number of fine-tuning parameters while using 3.5× fewer training tokens; larger output embeddings also improve transferability and representation generality.

  • Takeaways & Limitations

    Decoupled embeddings provide a more flexible allocation of capacity, with reduced input embeddings and enlarged output embeddings supporting efficient multilingual transfer.

  • Takeaways & Limitations

    The study focuses on encoder-only models and excludes encoder-decoder architectures such as T5.

Abstract

from arXiv · show

We re-evaluate the standard practice of sharing weights between input and output embeddings in state-of-the-art pre-trained language models. We show that decoupled embeddings provide increased modeling flexibility, allowing us to significantly improve the efficiency of parameter allocation in the input embedding of multilingual models. By reallocating the input embedding parameters in the Transformer layers, we achieve dramatically better performance on standard natural language understanding tasks with the same number of parameters during fine-tuning. We also show that allocating additional capacity to the output embedding provides benefits to the model that persist through the fine-tuning stage even though the output embedding is discarded after pre-training. Our analysis shows that larger output embeddings prevent the model's last layers from overspecializing to the pre-training task and encourage Transformer representations to be more general and more transferable to other tasks and languages. Harnessing these findings, we are able to train models that achieve strong performance on the XTREME benchmark without increasing the number of parameters at the fine-tuning stage.

1 INTRODUCTION

Pre-trained language models inherit coupled input and output embeddings, but this assumption constrains parameter allocation, especially in multilingual models. The paper studies decoupling as a way to improve flexibility, transfer performance, and representation generality.

  • Pre-trained language models typically use Transformer-based self-supervised pre-training followed by downstream fine-tuning.
  • Coupling input and output embeddings forces both matrices to share a dimensionality, limiting modeling flexibility.
  • 47–71% of multilingual models’ parameters can reside in high-dimensional embedding matrices, suggesting inefficient allocation.
  • The paper systematically evaluates embedding decoupling and reallocates capacity between embeddings and Transformer layers.
  • Larger output embeddings improve pre-training, transferability, and representation generality by preventing upper layers from overspecializing.

2 RELATED WORK

Related work situates embedding coupling within language modeling and translation, while prior studies show that upper Transformer layers can become specialized and less transferable. Multilingual models also face capacity pressures from shared vocabularies across many languages.

  • Embedding coupling: Embedding coupling was introduced to improve perplexity and reduce complexity in neural machine translation.
  • Transferability of representations: Prior work finds that representations transition from general to task-specific across model depth.
  • Transferability of representations: The last few Transformer layers can specialize to masked language modeling and become less transferable.
  • Multilingual models: Multilingual models use shared subword vocabularies across around 100 languages.
  • Multilingual models: The curse of multilinguality requires sufficient capacity for each language, producing large vocabularies and embedding sizes.
  • Efficient models: Efficiency research commonly uses pruning or distillation to reduce model size or transfer capacity from larger teachers.

3 EXPERIMENTAL METHODOLOGY

The experiments compare models using fine-tuning parameter counts, establish a coupled mBERT baseline, and evaluate variants on multilingual XTREME tasks. Results are averaged across three fine-tuning runs and generally use development sets.

  • Models are compared primarily by their number of parameters during fine-tuning, with pre-training and fine-tuning counts also reported.
  • The baseline matches mBERT with 12 Transformer layers, hidden size 768, coupled embeddings, and 177M total parameters.
  • Experiments use XTREME tasks including XNLI, NER, PAWS-X, XQuAD, MLQA, and TyDiQA-GoldP.
  • Results are averaged across three fine-tuning runs and evaluated on development sets unless otherwise stated.

4 EMBEDDING DECOUPLING REVISITED

Decoupling enables independent control of input and output embedding sizes, revealing that capacity is more productively allocated to output embeddings and Transformer layers. Smaller input embeddings preserve performance better than smaller output embeddings while freeing parameters.

  • Naïvely decoupling embeddings slightly improves average performance by 0.4, but gains are not uniform across tasks.
  • Decoupling allows input and output embedding dimensions to be changed independently.
  • With Ein = 128 and Eout = 768, the model has 43% fewer fine-tuning parameters yet slightly outperforms the Ein = 768, Eout = 128 comparison.
  • Reducing input embedding dimension saves parameters at a noticeably smaller accuracy cost than reducing output embedding size.
  • The output embedding contributes importantly to transferability, while a smaller input embedding frees parameters at modest performance cost.

5 EMBEDDING AND LAYER RESIZING FOR MORE EFFICIENT FINE-TUNING

The experiments resize embeddings and Transformer capacity while preserving efficient fine-tuning parameter budgets. Reinvesting input-embedding parameters into Transformer layers improves XTREME performance, with deeper models performing best.

  • Adding capacity via layers: Adding 11 Transformer layers during pre-training and discarding them afterward performs poorly on question answering tasks.The comparison matches the Eout = 768 model in pre-training and fine-tuning parameter counts.
  • Reinvesting input embedding parameters: Reducing Ein from 768 to 128 frees 77M parameters, which can be reinvested by increasing Transformer width or depth.The tested strategies increase H from 768 to 1024 or L from 12 to 23 while maintaining similar fine-tuning parameter counts.
  • Reinvesting input embedding parameters: An average improvement of 3.9 over the coupled-embedding baseline is achieved by increasing Transformer depth with the same overall fine-tuning parameter count.Reinvesting parameters in either H or L improves every task, while increasing L performs best.
  • A rebalanced mBERT: RemBERT combines reduced Ein, enlarged Eout, resized layers, and 32 Transformer layers for a multilingual model covering 110 languages.It matches XLM-R’s 559M fine-tuning parameters while using Ein = 256 and Eout = 1536.
  • A rebalanced mBERT: RemBERT outperforms XLM-R on every considered XTREME task despite using 3.5× fewer pre-training tokens and covering 10 more languages.Both models are fine-tuned using only the English training data for each corresponding task.

6 ON THE IMPORTANCE OF THE OUTPUT EMBEDDING SIZE

The analysis shows that larger output embeddings make Transformer representations more transferable across tasks and languages, while shifting output embeddings toward MLM-specific representations. These effects support more general upper-layer representations without increasing fine-tuning parameters.

  • 6.1 NATURE OF INPUT AND OUTPUT EMBEDDING REPRESENTATIONS: Larger output embeddings make output representations more specialized to MLM while allowing them to diverge from input embeddings.As Eout increases, output embeddings capture less traditional lexical similarity, while MLM accuracy increases; this flexibility avoids encoding relationships unnecessary for pre-training.
  • 6.2 CROSS-TASK TRANSFERABILITY OF TRANSFORMER LAYER REPRESENTATIONS: Removing the last layer improves XNLI performance for Eout = 128 and Eout = 768, including with 7.1M fewer parameters.With Eout = 128, performance remains similar after removing several final layers, suggesting those layers are not critical for transferability.
  • 6.2 CROSS-TASK TRANSFERABILITY OF TRANSFORMER LAYER REPRESENTATIONS: As Eout increases, the last Transformer layers become more transferable; with Eout = 3072, removing any layer worsens performance.For Eout = 768, removing more than one layer sharply reduces performance, whereas every layer remains useful at Eout = 3072.
  • 6.2 CROSS-TASK TRANSFERABILITY OF TRANSFORMER LAYER REPRESENTATIONS: Increasing Eout improves probing performance across all evaluated NLP tasks without increasing the number of parameters.The results indicate that larger output embeddings enable Transformer layers to learn more general representations.
  • 6.3 CROSS-LINGUAL TRANSFERABILITY OF TRANSFORMER LAYER REPRESENTATIONS: As Eout increases, Transformer layers become more language-agnostic and preserve language-agnostic representations to higher layers.Nearest-neighbor English-to-German accuracy rises, and the layer-7-to-11 slope becomes flatter; the last layer remains less language-agnostic than the previous one.

7 CONCLUSION

The paper concludes that decoupling embedding shapes provides the main benefit by enabling more flexible parameter allocation. Reducing input embeddings and enlarging output embeddings improves efficiency and representation transfer, including in RemBERT.

  • 7 CONCLUSION: Decoupling embedding shapes enables flexible parameter allocation, allowing input embeddings to be reduced and output embeddings to be enlarged.Reinvesting saved parameters improves performance while preserving the fine-tuning parameter count.
  • 7 CONCLUSION: RemBERT outperforms XLM-R with the same number of fine-tuning parameters after training on 3.5× fewer tokens.The conclusion also reports that larger output embeddings produce more transferable and general representations, particularly in upper Transformer layers.

A.1 EFFICIENCY COMPARISON BASED ON PARAMETER COUNT DURING FINE-TUNING

The paper compares model efficiency using parameter counts during fine-tuning because models are generally pre-trained once but may be fine-tuned or used for inference repeatedly.

  • A.1 EFFICIENCY COMPARISON BASED ON PARAMETER COUNT DURING FINE-TUNING: Efficiency is compared by the number of parameters during fine-tuning, excluding the task-specific head.This choice reflects repeated downstream use after one-time pre-training and avoids weighting discarded pre-training-only parameters equally.

A.2 BASELINE MODEL DETAILS

The baseline follows multilingual BERT with 12 Transformer layers, hidden size 768, 12 attention heads, and coupled embeddings matching the hidden size.

  • A.2 BASELINE MODEL DETAILS: The baseline has 12 Transformer layers, hidden size H = 768, and 12 attention heads with 64 dimensions each.Its coupled input and output embeddings use Eout = Ein = H, and the model has 177M parameters during pre-training and fine-tuning.
  • A.2 BASELINE MODEL DETAILS: For encoder-only models, parameters after the final Transformer layer, including output embeddings and the pooling layer, are discarded after pre-training.This distinguishes pre-training parameters from the parameters retained during fine-tuning.

A.3 TRAINING DETAILS

The experiments evaluate embedding factorization and reallocation strategies across multilingual and English models, using XTREME and related task evaluations. They also specify RemBERT training and a nearest-neighbor translation analysis.

  • Evaluation: XTREME evaluation covers cross-lingual inference, paraphrase identification, part-of-speech tagging, and named-entity recognition among its multilingual tasks.The benchmark includes XNLI, PAWS-X, Universal Dependencies POS data, and Wikiann named-entity data.
  • Embedding factorization: Reducing both input and output embeddings drastically reduces multilingual-model performance, unlike the smaller 0.8% accuracy drop reported for the monolingual model.The multilingual model drops by 2.6%, making direct ALBERT-style factorization unsuitable for multilingual models.
  • English-model experiments: English BERTBase experiments decouple embeddings, reduce Ein to 128, and vary output embedding size or pre-training depth on MNLI and SQuAD.These comparisons isolate input-dimension reduction from added output or layer capacity.
  • RemBERT training: RemBERT uses 995M parameters during pre-training and 575M during fine-tuning, with Wikipedia and Common Crawl data spanning 110 languages.Its hyperparameters and architecture are summarized in Table 14, while fine-tuning settings are listed in Table 15.
  • Leaderboard reporting: The XTREME leaderboard comparison reports both averages across tasks and averages across task categories.Table 16 provides per-task comparisons, while the category-average convention corresponds to Table 7.
  • Cross-lingual analysis: Translation analysis samples 5,000 English-German sentence pairs, averages token activations by layer, and evaluates nearest-neighbor retrieval using ℓ2 distance.The correct-pair retrieval frequency measures how often the nearest German sentence is the paired translation.
Loading 2010.12821v1…