XLNet: Generalized Autoregressive Pretraining for Language Understanding

In plain words

Before you fine-tune a language model on a task like question answering, you pretrain it on a lot of raw text. There were two main ways to do that. A left-to-right language model (like GPT) predicts each word from the words before it, so it never learns to use what comes after a word. BERT hides some words behind a [MASK] placeholder and guesses them from both sides. But that placeholder never shows up in real inputs later, and BERT guesses each hidden word separately, as if the hidden words had nothing to do with each other.

XLNet keeps left-to-right-style prediction but shuffles the order in which words are predicted. The text stays in place. Only the order changes in which the model reveals and predicts words. Over many random orders, every word ends up being predicted from words on both sides, and no placeholder token is needed. Picture filling a crossword in a different random order each time. Each answer is guessed from whatever letters are already filled in, wherever they sit.

The catch is that a normal Transformer cannot tell which position it is meant to predict next. XLNet fixes this with two parallel streams of hidden states. One stream knows where the target is but not what word sits there. The other stream knows the words, as usual. It also borrows two pieces from Transformer-XL, its long-context memory and its relative position encoding. Trained on the same data with the same recipe as BERT, it beat BERT on every task the authors compared.

TL;DR

XLNet replaces BERT's masked-token objective with permutation language modeling. For each sequence a factorization order z is sampled, and the model is trained autoregressively along z. Attention masks implement the order, while tokens keep their original positions. Only the last ≈1/K tokens of z are predicted (K = 6 for Large). To make predictions depend on which position is the target, the network runs two-stream self-attention: a content stream (ordinary self-attention, which includes the token itself) and a query stream (starts from a learned vector, attends only to strictly earlier tokens in z, and produces the logits). The backbone is Transformer-XL, with relative positional encodings, segment-level memory and relative segment encodings. At fine-tuning the query stream is dropped. Author-reported: with BERT's data and hyperparameters, XLNet-Large beats the best of three BERT-Large variants on all 11 benchmarks in tab. 1. Scaled to 32.89B tokens, it beats RoBERTa on RACE, SQuAD and most of GLUE. The paper describes several ingredients only loosely (span-based target sampling, the "bidirectional data pipeline", the memory length), and this card did not check them against the code.

Problem

Autoregressive (AR) LM pretraining maximizes a forward (or backward) product of conditionals, so each representation sees only one side of the context, while downstream understanding tasks need both sides §1 eq. 1. BERT-style denoising autoencoding (AE) sees both sides, but it has two flaws. (1) Its objective is only an approximation of the joint likelihood: the masked tokens are reconstructed independently of each other given the unmasked ones §2.1 eq. 2. (2) The artificial [MASK] symbol appears in pretraining but never at fine-tuning time, a pretrain-finetune discrepancy. Replacing masks with original tokens only partly helps, because that can be done only with small probability §2.1. The goal is an objective that captures bidirectional context, uses no input corruption, and keeps the exact product-rule factorization §1 §2.2.

Method

Section numbering follows the compiled paper: §1 Introduction, §2 Proposed Method (§2.1 Background, §2.2 Permutation LM objective, §2.3 Two-stream self-attention, §2.4 Transformer-XL ideas, §2.5 Multiple segments, §2.6 Discussion), §3 Experiments (§3.1 Pretraining and implementation, §3.2 Fair comparison with BERT, §3.3 Comparison with RoBERTa, §3.4 Ablation), §4 Conclusions. Appendix A has subsections A.1 (failure example), A.2 (two-stream details), A.3 (datasets), A.4 (hyperparameters), A.5 (discussion), A.6 (attention patterns), A.7 (visualizations). Numbered equations: eq. 1 AR objective, eq. 2 BERT objective, eq. 3 permutation LM, eq. 4 target-aware softmax, eq. 5 partial prediction. The two-stream and memory updates are displayed unnumbered.

Baselines being generalized §2.1:

\max_\theta \log p_\theta(\mathbf{x}) = \sum_{t=1}^{T} \log \frac{\exp(h_\theta(\mathbf{x}_{1:t-1})^\top e(x_t))}{\sum_{x'} \exp(h_\theta(\mathbf{x}_{1:t-1})^\top e(x'))} eq. 1

\max_\theta \log p_\theta(\bar{\mathbf{x}} \mid \hat{\mathbf{x}}) \approx \sum_{t=1}^{T} m_t \log \frac{\exp(H_\theta(\hat{\mathbf{x}})_t^\top e(x_t))}{\sum_{x'} \exp(H_\theta(\hat{\mathbf{x}})_t^\top e(x'))} eq. 2

Here m_t = 1 marks a masked token, \hat{\mathbf{x}} is the corrupted input and \bar{\mathbf{x}} the masked tokens §2.1.

1. Permutation LM objective §2.2. Let \mathcal{Z}_T be all permutations of [1, \dots, T], with z_t the t-th element and \mathbf{z}_{<t} the first t−1 elements:

\max_\theta \; \mathbb{E}_{\mathbf{z} \sim \mathcal{Z}_T} \Big[ \sum_{t=1}^{T} \log p_\theta(x_{z_t} \mid \mathbf{x}_{\mathbf{z}_{<t}}) \Big] eq. 3

Parameters are shared across all orders, so in expectation each token is trained with context from both sides §2.2. Only the factorization order is permuted. The input stays in natural order with its original positional encodings, and the permutation is realized purely through attention masks. The paper says this is necessary because fine-tuning sees only natural-order text §2.2 fig. 4.

2. Target-aware softmax §2.3. With a standard parameterization h_\theta(\mathbf{x}_{\mathbf{z}_{<t}}), two permutations that share the prefix \mathbf{z}_{<t} but have different next positions i ≠ j get the identical predicted distribution, so the model cannot know which position it is predicting §2.3 app. A.1. XLNet conditions on the target position:

p_\theta(X_{z_t} = x \mid \mathbf{x}_{\mathbf{z}_{<t}}) = \frac{\exp(e(x)^\top g_\theta(\mathbf{x}_{\mathbf{z}_{<t}}, z_t))}{\sum_{x'} \exp(e(x')^\top g_\theta(\mathbf{x}_{\mathbf{z}_{<t}}, z_t))} eq. 4

3. Two-stream self-attention §2.3. g_\theta(\cdot, z_t) must use the position z_t but not the content x_{z_t} when predicting x_{z_t}, yet it must encode x_{z_t} to serve as context for later targets. These requirements conflict, so there are two representations per position:

• content stream h_{z_t}: encodes \mathbf{x}_{\mathbf{z}_{\le t}}, including the token itself;
• query stream g_{z_t}: encodes \mathbf{x}_{\mathbf{z}_{<t}} and the position z_t only.

Initialization: g_i^{(0)} = w (one trainable vector), h_i^{(0)} = e(x_i) §2.3. Schematic per-layer update, with parameters shared between streams §2.3:

g_{z_t}^{(m)} \leftarrow \text{Attention}(Q = g_{z_t}^{(m-1)},\; KV = \mathbf{h}_{\mathbf{z}_{<t}}^{(m-1)};\; \theta)

h_{z_t}^{(m)} \leftarrow \text{Attention}(Q = h_{z_t}^{(m-1)},\; KV = \mathbf{h}_{\mathbf{z}_{\le t}}^{(m-1)};\; \theta)

Note that the query stream's keys and values are content states. The last-layer g_{z_t}^{(M)} feeds eq. 4 §2.3. The full layer with Transformer-XL memory \tilde{\mathbf{h}}^{(m-1)}, as given in app. A.2:

\hat h_{z_t}^{(m)} = \text{LayerNorm}\big(h_{z_t}^{(m-1)} + \text{RelAttn}(h_{z_t}^{(m-1)}, [\tilde{\mathbf{h}}^{(m-1)}, \mathbf{h}_{\mathbf{z}_{\le t}}^{(m-1)}])\big), \quad h_{z_t}^{(m)} = \text{LayerNorm}\big(\hat h_{z_t}^{(m)} + \text{PosFF}(\hat h_{z_t}^{(m)})\big)

\hat g_{z_t}^{(m)} = \text{LayerNorm}\big(g_{z_t}^{(m-1)} + \text{RelAttn}(g_{z_t}^{(m-1)}, [\tilde{\mathbf{h}}^{(m-1)}, \mathbf{h}_{\mathbf{z}_{< t}}^{(m-1)}])\big), \quad g_{z_t}^{(m)} = \text{LayerNorm}\big(\hat g_{z_t}^{(m)} + \text{PosFF}(\hat g_{z_t}^{(m)})\big)

The layer is post-LayerNorm. The content-stream update is exactly standard self-attention, so at fine-tuning the query stream is dropped and the network is used as a normal Transformer-XL §2.3. In mask terms: the content stream at z_t may attend to positions z_i with i \le t, while the query stream may attend only to i < t and cannot see its own token app. A.7 fig. 5 fig. 6.

4. Partial prediction §2.3. Optimizing the full eq. 3 converged slowly in preliminary experiments, so z is split at a cut point c and only the tail is predicted:

\max_\theta \; \mathbb{E}_{\mathbf{z} \sim \mathcal{Z}_T} \big[ \log p_\theta(\mathbf{x}_{\mathbf{z}_{>c}} \mid \mathbf{x}_{\mathbf{z}_{\le c}}) \big] = \mathbb{E}_{\mathbf{z} \sim \mathcal{Z}_T} \Big[ \sum_{t=c+1}^{|\mathbf{z}|} \log p_\theta(x_{z_t} \mid \mathbf{x}_{\mathbf{z}_{<t}}) \Big] eq. 5

The tail is chosen because it has the longest context. K sets |\mathbf{z}| / (|\mathbf{z}| - c) \approx K, so about 1/K of tokens are targets. Query states for non-targets need not be computed, which saves time and memory §2.3.

5. Span-based prediction §3.1. Sample a length L \in [1, \dots, 5], then select a consecutive span of L tokens as targets "within a context of (K L) tokens" §3.1. The paper gives no further detail (see Pitfalls).

6. Transformer-XL integration §2.4. Relative positional encodings are computed from the original positions. Segment recurrence works like this: process segment \tilde{\mathbf{x}} = \mathbf{s}_{1:T} under its own permutation, cache its per-layer content states \tilde{\mathbf{h}}^{(m)}, then for the next segment \mathbf{x} = \mathbf{s}_{T+1:2T}:

h_{z_t}^{(m)} \leftarrow \text{Attention}(Q = h_{z_t}^{(m-1)},\; KV = [\tilde{\mathbf{h}}^{(m-1)}, \mathbf{h}_{\mathbf{z}_{\le t}}^{(m-1)}];\; \theta) §2.4

Positional encodings depend only on actual positions, so the memory can be reused without knowing the previous segment's factorization order. The query stream uses memory the same way §2.4.

7. Multiple segments §2.5. Pretraining input follows BERT's format [CLS, A, SEP, B, SEP], with two segments sampled from the same context or from different ones. Memory is reused only when it belongs to the same context. XLNet-Large uses no next-sentence prediction §2.5 tab. 6.

8. Relative segment encodings §2.5. There are no absolute segment embeddings. For positions i, j, set \mathbf{s}_{ij} = \mathbf{s}_+ if they are in the same segment, else \mathbf{s}_-. Both are learnable per head. Compute a_{ij} = (\mathbf{q}_i + \mathbf{b})^\top \mathbf{s}_{ij}, where \mathbf{b} is a learnable per-head bias, and add a_{ij} to the normal attention score. This allows fine-tuning with more than two segments §2.5.

9. Bidirectional data pipeline §3.1. Because of recurrence, the forward and backward directions each take half of the batch §3.1.

Algorithm: XLNet pretraining step (paper-level description; code not checked)
Input: batch of length-512 sequences [CLS, A, SEP, B, SEP]; memory from previous segment of the same context
1. For each sequence, sample a factorization order z and a cut c with |z|/(|z|-c) ≈ K,
   using span-based target selection (L ∈ [1..5], span of L tokens in a context of K·L)   [§2.3] [§3.1]
2. Build two attention masks from z over [memory, current segment]:
     content mask: position z_t may see memory and z_1..z_t
     query mask:   position z_t may see memory and z_1..z_{t-1}                            [app. A.2] [app. A.7]
3. h^(0) = token embeddings; g^(0) = shared learned vector w, only at target positions      [§2.3]
4. For m = 1..M: update h with the content mask and g with the query mask; both use
   relative positional encodings, relative segment encodings and shared weights,
   with keys/values from content states only                                              [app. A.2] [§2.5]
5. Loss = −Σ over targets of log softmax(e(x)^T g^(M)_{z_t}) over the vocabulary           [eq. 4] [eq. 5]
6. Cache this segment's content states h^(m) as memory for the next segment                [§2.4]
Output: updated θ; at fine-tuning, drop the query stream                                   [§2.3]

Why it beats BERT on dependencies §2.6 app. A.5. For targets T and non-targets N, BERT optimizes \sum_{x \in \mathcal{T}} \log p(x \mid \mathcal{N}) and XLNet optimizes \sum_{x \in \mathcal{T}} \log p(x \mid \mathcal{N} \cup \mathcal{T}_{<x}). Every dependency covered by BERT is therefore covered by XLNet, and those that involve earlier targets are covered only by XLNet. Example from the paper, targets [New, York] given "is a city": XLNet adds the term \log p(\text{York} \mid \text{New, is a city}) §2.6.

Hyperparameters

Pretraining, XLNet-Large app. A.4 tab. 7 §3.1:

Pretraining data (full model) §3.1: BooksCorpus + English Wikipedia (13 GB together), Giga5 (16 GB), ClueWeb 2012-B (19 GB after filtering), Common Crawl (110 GB after filtering). In SentencePiece pieces: Wikipedia 2.78B, BooksCorpus 1.09B, Giga5 4.75B, ClueWeb 4.30B, Common Crawl 19.97B, for 32.89B in total. The filtering heuristics are not described beyond removing short or low-quality articles §3.1.

Other variants: XLNet-Large-wikibooks uses only BooksCorpus + Wikipedia and "all pretraining hyper-parameters as in the original BERT" §3.1. These values are not restated in the paper. XLNet-Base (ablation) uses BERT-Base model hyperparameters, also not restated §3.4.

Fine-tuning app. A.4 tab. 8:

Layer-wise decay: if the top (24th) layer uses learning rate l and the decay rate is α, layer m uses l \alpha^{24-m} app. A.4. Fine-tuning hyperparameters for the other tasks (other GLUE tasks, IMDB, Yelp-2, DBpedia, AG, Amazon) are not stated. Otherwise, fine-tuning "follows BERT" §3.1.

Task-specific heads app. A.3: SQuAD2.0 uses a logistic-regression answerability loss plus a standard span-extraction loss. GLUE multi-task models are trained jointly on MNLI, SST-2, QNLI and QQP and then fine-tuned on the other tasks. QNLI test submissions use pairwise relevance ranking, while QNLI dev uses standard classification. WNLI uses the loss of Kocijan et al. ClueWeb09-B uses XLNet word embeddings without fine-tuning, fed to a kernel pooling ranker app. A.3.

Results

All numbers are author-reported. Unless noted, models are 24-layer, about BERT-Large size.

Fair comparison: same data (BooksCorpus + Wikipedia) and BERT hyperparameters. The BERT row takes the best of three variants per task: original, whole-word masking, and no next-sentence prediction §3.2 tab. 1. SQuAD is EM/F1.

Scaled-up model (full data). The authors describe it as reusing RoBERTa's hyperparameters §3.3. RACE test accuracy and ClueWeb09-B test ranking tab. 2. * marks an ensemble, † marks the authors' own implementation:

SQuAD tab. 3 (EM / F1). † marks the authors' runs of the official BERT code, * an ensemble, and ‡ an older XLNet model, because test results for the latest one could not be obtained:

Text classification, test error rate in % (lower is better) tab. 4:

GLUE tab. 5. Dev results are the median of 10 runs. * marks an ensemble, and † a single-task result inside a multi-task row:

Ablation: 12-layer BERT-Base-size models trained on BooksCorpus + Wikipedia, median of 5 runs §3.4 tab. 6. BERT's RACE number comes from another paper. The other BERT numbers are the authors' runs with the same hyperparameter search space as XLNet. The column order F1, EM follows the paper:

The authors' reading of tab. 6: Transformer-XL (row 2 vs 1) and the permutation objective (rows 3–4 vs 2) both help. Removing memory hurts most on RACE, the longest-context task. Span-based prediction and the bidirectional pipeline both matter. Next-sentence prediction does not reliably help, so it is excluded §3.4. Rows 5–8 are changes relative to row 4 (K=6), as the indentation in tab. 6 shows.

Paper vs code

The official repository is https://github.com/zihangdai/xlnet, linked from the paper's abstract footnote. It was not inspected for this card, so no commit or code license is recorded, and no paper-vs-code comparison has been done. Items worth checking in the code: the exact span/permutation sampling procedure, memory length, the form of the bidirectional pipeline, Adam betas, SentencePiece vocabulary size, and whether the query stream is computed only for target positions.

Pitfalls

• Permute the masks, not the input. Tokens stay in natural order and keep their original (relative) positions. Only the attention masks encode z. Physically shuffling the sequence breaks the method, because fine-tuning sees natural-order text §2.2 §2.4.
• The query stream must not see its own token, at any layer. Its keys and values come from content states at strictly earlier positions in z. Its own query input carries only w plus position information through relative attention §2.3 app. A.2. Any leakage, for example letting the query mask include the diagonal, makes the objective trivial.
• The query stream is initialized from one shared vector w, not from position embeddings. Target-position information enters only through relative positional encodings in attention §2.3 §2.4. The paper does not say how relative encodings are formed for the query stream's own position. Check the code.
• Without the query stream the model is broken: a plain Transformer with permuted masks predicts the same distribution for every target position after a given prefix app. A.1.
• Partial prediction is essential for optimization. The full permutation objective converged slowly §2.3. K=6 means about 1/6 of tokens are targets, roughly 85 per 512-token sequence. K=7 was slightly better than K=6 on SQuAD2.0 and MNLI but worse on RACE and SST-2 tab. 6, so the choice is task-dependent.
• Span-based prediction is underspecified. "Select a consecutive span of L tokens as prediction targets within a context of (K L) tokens" §3.1 does not say how spans relate to the sampled permutation, how many spans are drawn per sequence, or how L is distributed. The ablation shows it matters tab. 6, row 6.
• "Bidirectional data input pipeline" is described in one sentence: each direction takes half the batch §3.1. What "backward" means for a permutation model is not explained. Removing it costs about 0.3 to 0.7 points on most tasks tab. 6, row 7.
• Hidden cost: two streams roughly double attention work for target positions in pretraining (query states are skipped for non-targets §2.3). Pretraining XLNet-Large took 512 TPU v3 chips for about 5.5 days, and the model was still underfitting §3.1.
• Relative segment encodings replace absolute segment embeddings. They encode only same-segment versus different-segment, with per-head \mathbf{s}_+, \mathbf{s}_- and bias \mathbf{b}, added to the attention score §2.5. Adding BERT-style absolute segment embeddings is not what the paper does.
• Memory is reused only across segments of the same context §2.5. With randomly paired A/B segments, the pipeline must track context boundaries.
• Which hyperparameters apply to which model is ambiguous. §3.1 lists the 500K-step, batch-8192 recipe for the scaled-up model, while §3.3 says the RoBERTa comparison "reuses the hyper-parameters of RoBERTa, as described in section 3.1". It is unclear whether tab. 7 is a RoBERTa-matched recipe or the original one. The wikibooks and Base variants use BERT's recipe, which the paper does not restate §3.1 §3.4.
• Fine-tuning recipes are given for only four tasks tab. 8. SQuAD uses layer-wise LR decay 0.75, the others 1.0.
• Benchmark caveats. The SQuAD1.1 test numbers come from an older XLNet model tab. 3. QNLI test uses a ranking formulation while dev uses classification app. A.3. GLUE test rows mix single-task and multi-task results tab. 5. ALBERT was deliberately left out of the comparisons because of its larger hidden sizes §3.3. ClueWeb09-B uses frozen embeddings, not fine-tuning app. A.3.
• The section on attention patterns refers to "Fig. 1-(d)", but Figure 1 has only panels (a) to (c) app. A.6 fig. 1. The intended reference is unclear.

Lineage

• Builds on: Transformer-XL (arXiv:1901.02860), which supplies the backbone, relative positional encodings and segment recurrence §2.4. BERT (arXiv:1810.04805) supplies the AE baseline, input format and fine-tuning procedure §2.5 §3.1. Orderless NADE (Uria et al., 2016) supplies the permutation-based AR idea. XLNet differs in its motivation (bidirectional context for pretraining rather than density estimation) and in using two-stream attention to make predictions target-aware §1.
• Compared against: BERT, RoBERTa (arXiv:1907.11692), GPT, MT-DNN, ULMFiT, DPCNN, kernel-pooling rankers tab. 1 tab. 2 tab. 3 tab. 4 tab. 5.