PaperlayerView exact revision on arXiv

Source-linked AI summary

ALBERT: A Lite BERT for Self-supervised Learning of Language Representations

Zhenzhong Lan, Mingda Chen, Sebastian Goodman, Kevin Gimpel, Piyush Sharma, Radu Soricut

arXiv:1909.11942v6cs.CLcs.AI

TL;DR

Scaling pretrained language models is constrained by hardware memory and communication overhead. ALBERT reduces parameters with factorized embeddings and cross-layer sharing, adds sentence-order prediction, and achieves state-of-the-art results across GLUE, SQuAD, and RACE while using fewer parameters than BERT-large.

  • Problem

    Scaling pretrained language models is constrained by hardware memory limits and communication overhead that grows with model parameters.

  • Method

    ALBERT combines factorized embedding parameterization, cross-layer parameter sharing, and a sentence-order prediction loss focused on inter-sentence coherence.

  • Results

    89.4% RACE accuracy was achieved alongside state-of-the-art results on GLUE, SQuAD, and RACE with fewer parameters than BERT-large.

  • Takeaways & Limitations

    ALBERT supports larger pretrained configurations with fewer parameters than BERT-large while achieving significantly better benchmark performance.

  • Takeaways & Limitations

    Although ALBERT-xxlarge has fewer parameters and better results than BERT-large, its larger structure makes it computationally more expensive.

Abstract

from arXiv · show

Increasing model size when pretraining natural language representations often results in improved performance on downstream tasks. However, at some point further model increases become harder due to GPU/TPU memory limitations and longer training times. To address these problems, we present two parameter-reduction techniques to lower memory consumption and increase the training speed of BERT. Comprehensive empirical evidence shows that our proposed methods lead to models that scale much better compared to the original BERT. We also use a self-supervised loss that focuses on modeling inter-sentence coherence, and show it consistently helps downstream tasks with multi-sentence inputs. As a result, our best model establishes new state-of-the-art results on the GLUE, RACE, and \squad benchmarks while having fewer parameters compared to BERT-large. The code and the pretrained models are available at https://github.com/google-research/ALBERT.

1 INTRODUCTION

ALBERT addresses the memory and communication barriers to scaling pretrained language models through parameter reduction and a sentence-order prediction objective. Its larger configurations use fewer parameters than BERT-large while achieving state-of-the-art results on GLUE, SQuAD, and RACE.

  • Motivation: Large pretrained networks improve state-of-the-art NLP performance, but scaling them is constrained by hardware memory limits and parameter-dependent communication overhead.Current state-of-the-art models can contain hundreds of millions or billions of parameters, making these constraints difficult to avoid.
  • ALBERT architecture: ALBERT uses significantly fewer parameters than traditional BERT to address both memory limitations and distributed-training communication overhead.The paper distinguishes this approach from existing solutions that address memory but not communication overhead.
  • ALBERT architecture: Factorized embedding parameterization decomposes the vocabulary embedding matrix into two smaller matrices, separating hidden-layer size from vocabulary-embedding size.This makes it easier to increase hidden size without significantly increasing vocabulary-embedding parameters.
  • Self-supervised objective: ALBERT introduces sentence-order prediction, a self-supervised loss focused on inter-sentence coherence and designed to address weaknesses of BERT’s next-sentence prediction objective.The objective is intended to improve ALBERT’s performance on inputs involving multiple sentences.
  • Results: 89.4% is the reported RACE accuracy, 89.4 is the GLUE benchmark score, and 92.2 is the SQuAD 2.0 F1 score for ALBERT’s state-of-the-art results.These results were achieved by larger ALBERT configurations that still used fewer parameters than BERT-large.

2 RELATED WORK

Related work tracks the shift from word-embedding pretraining toward pretrained language representations, alongside methods for reducing the computational burden of large models. It also explores cross-layer parameter sharing and pretraining objectives modeling discourse coherence, which motivate ALBERT’s design.

  • Representation learning: Pretraining has shifted from standard or contextualized word embeddings toward broader pretrained natural-language representations.These representations have been widely adopted across NLP tasks.
  • Memory-efficient training: Large-model experimentation is constrained by GPU/TPU memory, motivating gradient checkpointing and related memory-reduction methods.Gradient checkpointing reduces memory requirements sublinearly but requires an extra forward pass.
  • Parameter sharing: Cross-layer parameter sharing has been explored in Transformers, but prior work primarily targeted standard encoder-decoder training rather than pretraining and finetuning.Other studies reported differing performance outcomes for shared-parameter networks, including improvements on language modeling and subject-verb agreement.
  • Parameter sharing: Other Transformer studies examined equilibrium behavior or combined shared and standard layers, with reported observations of convergence, oscillation, or increased parameter counts.The cited work includes Deep Equilibrium Models and hybrids that further increase the standard Transformer’s parameters.
  • Coherence objectives: Pretraining objectives have also modeled discourse coherence, and ALBERT predicts the ordering of two consecutive text segments.Prior research has studied coherence and cohesion phenomena connecting neighboring text segments, often using simple objectives.

3 THE ELEMENTS OF ALBERT

ALBERT introduces three design changes over BERT: factorized embedding parameterization, cross-layer parameter sharing, and a sentence-order prediction loss focused on inter-sentence coherence. These choices substantially improve parameter efficiency, with ALBERT-large using 18M parameters versus BERT-large’s 334M.

  • Factorized embedding parameterization: ALBERT factorizes embedding parameters, reducing their complexity from O(V × H) to O(V × E + E × H), especially when H ≫ E.The model projects one-hot vectors into an E-dimensional embedding space before projecting them into the H-dimensional hidden space.
  • Cross-layer parameter sharing: ALBERT shares all parameters across encoder layers by default to improve parameter efficiency.The paper also considers sharing only feed-forward or only attention parameters, but all experiments use all-parameter sharing unless otherwise specified.
  • Cross-layer parameter sharing: Weight sharing produces smoother layer-to-layer transitions in ALBERT than BERT, while input-output embedding distances and similarities remain nonzero after 24 layers.The measurements use BERT-large and ALBERT-large configurations and indicate that weight sharing stabilizes network parameters.
  • Inter-sentence coherence loss: ALBERT replaces BERT’s next-sentence prediction with sentence-order prediction, using swapped consecutive segments as negative examples to focus on inter-sentence coherence.The authors argue that NSP is comparatively easy because it conflates topic prediction with coherence prediction.
  • Parameter efficiency: 18M versus 334M parameters gives ALBERT-large about 18x fewer parameters than BERT-large.ALBERT-xlarge with H = 2048 has 60M parameters, while ALBERT-xxlarge with H = 4096 has 233M parameters, around 70% of BERT-large’s parameters.

4 EXPERIMENTAL RESULTS … 4.5 CROSS-LAYER PARAMETER SHARING

The experiments use a standardized pretraining setup and evaluate ALBERT on intrinsic and downstream tasks. Results show improved parameter efficiency and throughput, while embedding size and cross-layer sharing create important performance trade-offs.

  • 4.1 EXPERIMENTAL SETUP: Models use BOOKCORPUS and English Wikipedia, sequences capped at 512 tokens, batch size 4096, LAMB optimization, and 125,000 training steps unless specified otherwise.Training uses Cloud TPU V3 systems with 64–512 TPUs depending on model size.
  • 4.2 EVALUATION BENCHMARKS: A held-out development set from SQuAD and RACE monitors MLM and sentence-classification convergence without affecting downstream model selection.The set is used only to check convergence.
  • 4.2 EVALUATION BENCHMARKS: Downstream evaluation covers GLUE, two versions of SQuAD, and the RACE reading-comprehension dataset.These benchmarks are used following prior evaluation protocols.
  • 4.3 OVERALL COMPARISON BETWEEN BERT AND ALBERT: With around 70% of BERT-large’s parameters, ALBERT-xxlarge achieves significant improvements over BERT-large on representative downstream development-set tasks.The passage identifies parameter efficiency as ALBERT’s most important design advantage and reports task-specific score differences in Table 2.
  • 4.3 OVERALL COMPARISON BETWEEN BERT AND ALBERT: 1.7 times faster data iteration is observed for ALBERT-large than BERT-large, whereas ALBERT-xxlarge is about 3 times slower under the same training configuration.The throughput differences are attributed to less communication and fewer computations, offset by ALBERT-xxlarge’s larger structure.
  • 4.4 FACTORIZED EMBEDDING PARAMETERIZATION: Embedding size E = 128 is selected for future settings because it performs best under all-shared parameterization, whereas larger embeddings help only modestly without sharing.The comparison uses an ALBERT-base configuration.
  • 4.5 CROSS-LAYER PARAMETER SHARING: Sharing all cross-layer parameters reduces performance by -1.5 on Avg for E = 128 and -2.5 on Avg for E = 768, with most loss caused by sharing FFN parameters.Attention sharing causes +0.1 on Avg for E = 128 and -0.7 on Avg for E = 768; smaller sharing groups improve performance but increase total parameters.

4.6 SENTENCE ORDER PREDICTION (SOP) … 4.9 CURRENT STATE-OF-THE-ART ON NLU TASKS

ALBERT’s SOP loss improves multi-sentence downstream performance, while its efficiency and training choices yield stronger results than BERT under equal training time. Additional data, removing dropout, and the best configuration support state-of-the-art scores on GLUE, SQuAD, and RACE.

  • 4.6 SENTENCE ORDER PREDICTION (SOP): SOP loss achieves 78.9% accuracy on NSP and 86.5% on SOP, whereas NSP reaches only 52.0% on SOP.The 52.0% NSP result is similar to random guessing, indicating that NSP models topic shift rather than sentence order.
  • 4.6 SENTENCE ORDER PREDICTION (SOP): SOP consistently improves downstream task performance for multi-sentence inputs compared with the none and NSP conditions.The comparison uses ALBERT-base configurations and evaluates intrinsic and downstream tasks.
  • 4.7 WHAT IF WE TRAIN FOR THE SAME AMOUNT OF TIME?: +1.5% Avg improvement and +5.2% on RACE make ALBERT-xxlarge significantly better than BERT-large after roughly equal training time.BERT-large processes data about 3.17x faster than ALBERT-xxlarge, so the comparison controls for actual training hours rather than training steps.
  • 4.8 ADDITIONAL TRAINING DATA AND DROPOUT EFFECTS: Additional training data significantly boosts dev-set MLM accuracy and improves downstream tasks except the Wikipedia-based SQuAD benchmarks.The SQuAD decline is attributed to negative effects from out-of-domain training material.
  • 4.8 ADDITIONAL TRAINING DATA AND DROPOUT EFFECTS: Removing dropout significantly improves MLM accuracy and helps downstream tasks on intermediate ALBERT-xxlarge evaluations around 1M training steps.The models still do not overfit after 1M steps, motivating removal of dropout to increase model capacity.
  • 4.9 CURRENT STATE-OF-THE-ART ON NLU TASKS: The single-model ALBERT setup uses ALBERT-xxlarge with combined MLM and SOP losses and no dropout.State-of-the-art evaluation reports single-model and ensemble results using single-task fine-tuning and median development-set performance over five runs.
  • 4.9 CURRENT STATE-OF-THE-ART ON NLU TASKS: ALBERT achieves a GLUE score of 89.4, a SQuAD 2.0 test F1 score of 92.2, and a RACE test accuracy of 89.4.The RACE result is +17.4% absolute points over BERT, +7.6% over XLNet, +6.2% over RoBERTa, and 5.3% over DCMI+.

5 DISCUSSION

ALBERT-xxlarge uses fewer parameters than BERT-large and achieves significantly better results, but its larger structure makes it computationally more expensive. The discussion therefore highlights faster attention methods and orthogonal approaches for increasing representation power as future directions.

  • ALBERT-xxlarge has fewer parameters than BERT-large while achieving significantly better results.
  • Its larger structure makes ALBERT-xxlarge computationally more expensive than BERT-large.
  • Sparse attention and block attention are proposed as methods to improve ALBERT’s training and inference speed.
  • Hard example mining is identified as an orthogonal research direction that could provide additional representation power.

A APPENDIX · A.1 EFFECT OF NETWORK DEPTH AND WIDTH

This appendix examines how network depth and hidden-layer width affect ALBERT performance. Increasing width improves performance with diminishing returns, while very large width can substantially hurt performance; warm-starting helps convergence but not downstream accuracy.

  • A APPENDIX · A.1 EFFECT OF NETWORK DEPTH AND WIDTH: The study evaluates how the number of layers and hidden size affect ALBERT performance.It uses an ALBERT-large configuration for depth experiments and a 3-layer ALBERT-large configuration for width experiments.
  • A APPENDIX · A.1 EFFECT OF NETWORK DEPTH AND WIDTH: Different layer counts are compared for an ALBERT-large configuration in Table 11.Networks with at least 3 layers are fine-tuned from checkpoints trained at the preceding depth.
  • A APPENDIX · A.1 EFFECT OF NETWORK DEPTH AND WIDTH: The depth experiments use staged fine-tuning, such as initializing the 12-layer network from the 6-layer checkpoint.The passage notes that a similar technique was used by Gong et al. (2019).
  • A APPENDIX · A.1 EFFECT OF NETWORK DEPTH AND WIDTH: Increasing hidden size improves performance with diminishing returns, but performance declines significantly at a hidden size of 6144.These results come from a 3-layer ALBERT-large configuration summarized in Table 12.
  • A APPENDIX · A.1 EFFECT OF NETWORK DEPTH AND WIDTH: Warm-starting does not improve ALBERT-large downstream performance, but it helps the 48-layer network converge.The same approach was also applied to ALBERT-xxlarge by warm-starting from a 6-layer network.
  • A APPENDIX · A.1 EFFECT OF NETWORK DEPTH AND WIDTH: The configurations are evaluated using SQuAD1.1, SQuAD2.0, MNLI, SST-2, RACE, and their average.The reported table fields also include hidden size and parameter count.

A.2 DO VERY WIDE ALBERT MODELS NEED TO BE DEEP(ER) TOO?

For very wide ALBERT-xxlarge models, increasing depth from 12 to 24 layers yields negligible downstream-accuracy differences, so deeper-than-12-layer configurations are unnecessary when all cross-layer parameters are shared.

  • A.2 DO VERY WIDE ALBERT MODELS NEED TO BE DEEP(ER) TOO?: For ALBERT-xxlarge (H=4096), the downstream-accuracy difference between 12-layer and 24-layer configurations is negligible, with identical Avg scores.This tests whether the earlier ALBERT-large finding also holds for much wider models.
  • A.2 DO VERY WIDE ALBERT MODELS NEED TO BE DEEP(ER) TOO?: When all cross-layer parameters are shared in the ALBERT style, models deeper than 12 layers are unnecessary.The conclusion follows from the negligible accuracy difference between the 12-layer and 24-layer ALBERT-xxlarge configurations.

A.3 DOWNSTREAM EVALUATION TASKS

The downstream evaluation covers GLUE’s language-understanding tasks, extractive question answering on SQuAD, and multi-choice reading comprehension on RACE.

  • GLUE: GLUE comprises 9 tasks spanning linguistic acceptability, sentiment, paraphrase, textual similarity, question pairs, and natural-language inference.The listed tasks include CoLA, SST, MRPC, STS, QQP, MNLI, QNLI, and RTE.
  • SQuAD: SQuAD evaluates extractive question answering by requiring models to predict answer spans from Wikipedia context paragraphs.The evaluation uses SQuAD v1.1 and v2.0; v2.0 adds 50,000 unanswerable questions to v1.1’s 100,000 annotated pairs.
  • RACE: RACE evaluates multi-choice reading comprehension on nearly 100,000 examination questions, each offering 4 candidate answers.Models concatenate the passage, question, and each candidate answer, then use the [CLS] representation to predict answer probabilities.

A.4 HYPERPARAMETERS

Downstream-task hyperparameters for ALBERT are presented in Table 14, using standard configuration labels for optimization, batching, regularization, training, warmup, and sequence length. The settings are adapted from prior work by Liu et al. (2019), Devlin et al. (2019), and Yang et al. (2019).

  • A.4 HYPERPARAMETERS: The downstream-task settings are adapted from Liu et al. (2019), Devlin et al. (2019), and Yang et al. (2019).
  • A.4 HYPERPARAMETERS: Table 14 reports ALBERT’s downstream-task hyperparameters, including learning rate, batch size, dropout rate, training steps, warmup steps, and maximum sequence length.The table abbreviates these settings as LR, BSZ, DR, TS, WS, and MSL.
Loading 1909.11942v6…