XLNet, RoBERTa, ALBERT (2019) › The Solutions

XLNet addressed the masked language modeling limitations by introducing a generalized autoregressive pretraining approach that could capture bidirectional context without masking. The key insight was that autoregressive models could be made bidirectional by considering all possible factorization orders of the input sequence.

Permutation Language ModelingLink Copied

Instead of masking tokens, XLNet used permutation language modeling. For a sequence of length $n$, there are $n!$ possible permutation orders. XLNet sampled a permutation order during training and predicted tokens in that order. This approach maintained the autoregressive property needed for valid probability distributions while allowing the model to see both left and right context for any given position.

For example, given the sequence "The capital of France is Paris", XLNet might sample permutation order $[3, 1, 5, 2, 4]$ where positions are indexed from 1. The model would predict position 3 first (seeing positions 1, 2, 4, 5), then position 1 (seeing positions 2, 4, 5), then position 5 (seeing positions 1, 2, 4), and so on. By averaging over many different permutations during training, the model learned representations that captured bidirectional context.

The permutation approach solved several problems. First, it eliminated the pretrain-finetune discrepancy because the model always saw actual tokens, never masked placeholders. Second, it naturally captured dependencies between tokens because the autoregressive structure modeled conditional distributions. Third, it maintained the ability to be generative while capturing bidirectional context, unlike BERT's purely bidirectional approach.

Two-Stream Self-AttentionLink Copied

To implement permutation language modeling efficiently, XLNet introduced two-stream self-attention. The query stream handled the position being predicted, while the content stream handled positions already observed. This dual-stream architecture allowed the model to properly implement the autoregressive structure while maintaining computational efficiency comparable to BERT.

The content stream used both content and position embeddings, allowing tokens to access their own content and position information. The query stream used only position embeddings. This ensured that when predicting a token, the model could not directly access that token's content. This structure enforced the autoregressive constraint while allowing bidirectional context through the permutation mechanism.

Relative Positional EncodingLink Copied

XLNet also used relative positional encodings, similar to those introduced in Transformer-XL. Rather than encoding absolute positions, the model encoded relative distances between query and key positions. This approach improved generalization to sequences longer than those seen during training and better captured positional relationships.

RoBERTa maintained BERT's architecture but systematically improved its training procedure. This showed that better training methodology could yield substantial improvements without architectural changes.

Training Data and DurationLink Copied

RoBERTa used significantly more training data than BERT. Where BERT had trained on BooksCorpus and English Wikipedia totaling about 16GB, RoBERTa expanded to include CommonCrawl news data, web text, and stories, totaling about 160GB. This tenfold increase in training data provided more diverse examples and improved generalization.

The model also trained for longer. BERT had trained for 1M steps, but RoBERTa extended this significantly, allowing the model to fully converge and extract maximum benefit from the larger dataset. The combination of more data and longer training enabled the model to learn richer representations.

Dynamic MaskingLink Copied

RoBERTa replaced BERT's static masking with dynamic masking. Instead of applying the same masking pattern to each sequence across epochs, RoBERTa generated a new masking pattern each time a sequence was processed. This prevented the model from overfitting to specific masks and encouraged representations that generalized across masking patterns.

Dynamic masking was straightforward to implement but had significant impact. By varying the masking patterns, the model learned to handle unpredictable contexts, which better matched the variability in real-world text where word relationships could take many forms.

Removing Next Sentence PredictionLink Copied

RoBERTa removed the next sentence prediction task entirely. Experiments showed that NSP was not necessary and might even hurt performance. Instead, the model focused solely on masked language modeling, using longer sequences and more training data to learn sentence-level relationships implicitly.

The removal of NSP simplified the training procedure and allowed the model to allocate more capacity to learning word-level and sequence-level patterns through masked language modeling alone. The results validated that masked language modeling was sufficient to learn high-quality representations.

Optimized HyperparametersLink Copied

RoBERTa used carefully tuned hyperparameters optimized for the larger dataset and longer training. The learning rate schedule, batch size, and other training details were adjusted based on empirical findings. While these changes were incremental, their cumulative effect contributed to improved performance.

ALBERT redesigned BERT's architecture to use parameters more efficiently, enabling training of larger models within fixed computational budgets.

Factorized Embedding ParameterizationLink Copied

ALBERT separated the vocabulary embedding size from the hidden dimension size. In BERT, these were the same, meaning embedding parameters scaled as $V \times H$ where $V$ is vocabulary size and $H$ is hidden dimension. ALBERT introduced an intermediate embedding dimension $E$ smaller than $H$, creating embeddings of size $V \times E$ that were then projected to size $H$ through a linear transformation.

This factorization reduced embedding parameters from $V \times H$ to $V \times E + E \times H$. For typical values where $H = 768$ and $E = 128$, this reduced embedding parameters substantially. The projection layer added only $E \times H$ parameters, which was much smaller than the savings from reducing the embedding matrix.

Cross-Layer Parameter SharingLink Copied

ALBERT shared parameters across all transformer layers, rather than having independent parameters for each layer. This meant that all layers learned the same transformations, dramatically reducing the parameter count. For a 12-layer model, this reduced parameters by roughly a factor of 12 for the layer-specific components.

Parameter sharing forced the model to learn more general transformations that worked across all positions in the network. While this constraint might seem limiting, experiments showed that shared-parameter models could achieve comparable performance to unshared models while using far fewer parameters. The shared parameters learned to represent information at different abstraction levels implicitly through the layer depth.

Sentence Order PredictionLink Copied

ALBERT replaced next sentence prediction with sentence order prediction (SOP). Instead of predicting whether two sentences were consecutive, SOP predicted whether two sentences appeared in the correct order. This task proved more effective because it required understanding sentence-level coherence, not just adjacency.

SOP worked by swapping the order of consecutive sentences half the time during training. The model learned to distinguish whether sentences appeared in their natural order or had been swapped. This objective better captured sentence-level relationships and improved downstream performance on tasks requiring sentence understanding.

Effects of Parameter EfficiencyLink Copied

These architectural innovations enabled ALBERT to train much larger models. While BERT-base had 110M parameters and BERT-large had 340M parameters, ALBERT-base achieved similar performance with 12M parameters, and ALBERT-large achieved better performance than BERT-large with 18M parameters. The extreme parameter efficiency meant that even larger ALBERT models could be trained within reasonable computational budgets.