Semformer: Transformer Language Models with Semantic Planning

TL;DR

Semformer introduces semantic planning into Transformer language models, improving their ability to predict responses accurately and mitigating shortcut learning, with strong results in minimal planning tasks and downstream NLP applications.

Contribution

The paper proposes a novel training method for Transformers that explicitly models semantic planning, addressing shortcut learning and enhancing performance on various tasks.

Findings

Near-perfect performance in graph path-finding task

Effective mitigation of shortcut learning

Improved perplexity and in-context learning results

Abstract

Next-token prediction serves as the dominant component in current neural language models. During the training phase, the model employs teacher forcing, which predicts tokens based on all preceding ground truth tokens. However, this approach has been found to create shortcuts, utilizing the revealed prefix to spuriously fit future tokens, potentially compromising the accuracy of the next-token predictor. In this paper, we introduce Semformer, a novel method of training a Transformer language model that explicitly models the semantic planning of response. Specifically, we incorporate a sequence of planning tokens into the prefix, guiding the planning token representations to predict the latent semantic representations of the response, which are induced by an autoencoder. In a minimal planning task (i.e., graph path-finding), our model exhibits near-perfect performance and effectively mitigates shortcut learning, a feat that standard training methods and baseline models have been unable to accomplish. Furthermore, we pretrain Semformer from scratch with 125M parameters, demonstrating its efficacy through measures of perplexity, in-context learning, and fine-tuning on summarization tasks.