Segatron: Segment-Aware Transformer for Language Modeling and Understanding

Authors

  • He Bai University of Waterloo
  • Peng Shi University of Waterloo
  • Jimmy Lin University of Waterloo RSVP.ai
  • Yuqing Xie University of Waterloo
  • Luchen Tan RSVP.ai
  • Kun Xiong RSVP.ai
  • Wen Gao Peking University
  • Ming Li University of Waterloo RSVP.ai

DOI:

https://doi.org/10.1609/aaai.v35i14.17485

Keywords:

Applications

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. Our code is available on GitHub.

Downloads

Published

2021-05-18

How to Cite

Bai, H., Shi, P., Lin, J., Xie, Y., Tan, L., Xiong, K., Gao, W., & Li, M. (2021). Segatron: Segment-Aware Transformer for Language Modeling and Understanding. Proceedings of the AAAI Conference on Artificial Intelligence, 35(14), 12526-12534. https://doi.org/10.1609/aaai.v35i14.17485

Issue

Section

AAAI Technical Track on Speech and Natural Language Processing I