Segatron: Segment-Aware Transformer for Language Modeling and Understanding

TL;DR

Segatron introduces a segment-aware encoding that injects paragraph, sentence, and token positions into Transformer models. By applying this to Transformer-XL (Segatron-XL) and to MLM pre-training (SegaBERT), the approach yields notable perplexity reductions on WikiText-103 and improved performance on GLUE, SQuAD, and RACE, with SegaBERT-large surpassing RoBERTa-large on zero-shot STS. The work provides both quantitative gains and qualitative attention analyses, supporting the value of richer hierarchical positional information for language modeling and understanding. Overall, segment-aware position encoding enhances long-range dependency modeling and transfer learning capabilities in both autoregressive and masked settings, with practical implications for robust NLP systems.

Abstract

Transformers are powerful for sequence modeling. Nearly all state-of-the-art language models and pre-trained language models are based on the Transformer architecture. However, it distinguishes sequential tokens only with the token position index. We hypothesize that better contextual representations can be generated from the Transformer with richer positional information. To verify this, we propose a segment-aware Transformer (Segatron), by replacing the original token position encoding with a combined position encoding of paragraph, sentence, and token. We first introduce the segment-aware mechanism to Transformer-XL, which is a popular Transformer-based language model with memory extension and relative position encoding. We find that our method can further improve the Transformer-XL base model and large model, achieving 17.1 perplexity on the WikiText-103 dataset. We further investigate the pre-training masked language modeling task with Segatron. Experimental results show that BERT pre-trained with Segatron (SegaBERT) can outperform BERT with vanilla Transformer on various NLP tasks, and outperforms RoBERTa on zero-shot sentence representation learning.

Figures (8)

• Figure 1: Input representation of Segatron-XL and SegaBERT.
• Figure 2: Valid perplexities during the training processes of language modeling.
• Figure 3: Test perplexities of Segatron-XL and Transformer-XL with different input lengths.
• Figure 4: Valid losses during the pre-training.
• Figure 5: Self-attention heat maps of the first, the sixth, and the last layer of SegaBERT and BERT when encoding the first 512 tokens of a Wikipedia article.