Semformer: Transformer Language Models with Semantic Planning (2024)

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.

## introduction neural language models (lms) @xcite , a fundamental component of natural language processing (nlp), have witnessed significant advancements in recent years. by scaling up model sizes and pretraining on extensive text, large language models (llms) have successfully learned language and world knowledge, which has resulted in promising performance across various tasks and even demonstrated reasoning capabilities @xcite @xcite @xcite . the success of these models can be attributed to a straightforward training paradigm: next-token prediction with * corresponding author 1 https://github.com/aries-lm/semformer.git teacher forcing @xcite , in which the models are trained to predict tokens using all preceding ground truth tokens as input. recent studies, however, have raised concerns about the efficacy of the aforementioned training scheme in facilitating the learning of an accurate problem solver or planner @xcite @xcite @xcite . for instance, the graph path-finding task-which necessitates lookahead and planning-demonstrates that teacher forcing can lead to a clever hans cheat phenomenon characterized by shortcut learning @xcite . consequently, the later nodes such as 3 and 1 in figure 1 become easier to predict, while the first node of the answer (i.e., 22) becomes more challenging to learn. this could result in a highly inaccurate next-token predictor, which would struggle to generalize to unseen problems, even without considering out-ofdistribution and length generalization. humans, intuitively, do not rely solely on historical context to solve a problem @xcite . instead, they formulate an abstract plan based on the problem at hand, which subsequently guides them towards the final answers. for the problem in figure 1 , the quickest solution is to look ahead at the later nodes to identify a unique path correspond-ing to the problem, and then reverse this found path to generate the correct answer. in general, a language model should internalize the process of looking ahead or "thinking about the future". the semantics of finding the response path is predicted by internal computation, with token output guided by the intended semantics. to this end, we incorporate semantic planning into next-token prediction in a decoder-only transformer @xcite , which we refer to as semformer. our semformer is composed of a language model and an autoencoder that is used only during training. for the language model, we introduce a semantic planning token sequence that follows the prefix of the input. this sequence disregards the general next-token prediction loss and is utilized to predict the latent representations of the subsequent tokens. the autoencoder learns to generate the sequence of latent representations, compressing the subsequent tokens into a low-dimensional space. on the graph path-finding problem @xcite , our semformer achieves almost 100% accuracy scores on the settings of different levels of difficulty, showing superiority to the related baselines. only introducing dummy tokens in the sequence (i.e., pause transformer) @xcite fails to learn the planning task. moreover, our semformer learns to solve the problem significantly faster than the baselines, merely one epoch based on the gpt2-large @xcite . further, to validate the effectiveness of this architecture on general lm pretraining, we train transformer models with 125m parameters from scratch on openwebtext. semformer results in improvements on perplexity evaluation, in-context learning, and fine-tuning on abstractive summarization. beyond the next-token prediction, various training paradigms have been proposed including nonautoregressive models @xcite , diffusion lm @xcite , and multiple token prediction @xcite @xcite . predicting multiple future tokens is originally for accelerating inference, and @xcite recently show that it can also avoid the localness issue of next-token prediction with teacher forcing. @xcite introduce a latent diffusion model to generate paragraph representations induced by a variational autoencoder (kingma and welling, 2014), and feed them into the language model to help paragraph generation. rather than significantly changing the model architecture, we internalize the planning ability into the language model, achieved through the semantic representation prediction of the subsequent sequence. custom tokens in language modeling custom tokens can be used to increase model capacity used as additional memory @xcite @xcite . for example, @xcite propose to apply custom tokens recurrently, leading to improvement on long sequence modeling and algorithmic tasks. compressing long prompts into fixed length sequence can alleviate the heavy burden of a large key-value cache during inference @xcite @xcite . custom tokens are also used to optimize pretrained models language model figure 2 : illustration of our semformer. we introduce trainable tokens in language modeling. the representations of the tokens encoded by the language model are regressed to the latent representations of the response with l 2 loss. we can share the parameters between the language model and the encoder, and utilize a small decoder to enhance training efficiency. to accomplish specific downstream tasks, i.e., parameter efficient fine-tuning @xcite . for vision transformer, @xcite find that appending trainable tokens to image patches leads to smoother representation learning. incorporating trainable tokens has been demonstrated as an effective way to enhance the transformer's reasoning and planning capabilities. @xcite find that such a method leads to small perplexity gains on reasoning tasks, and @xcite investigate its effectiveness on the setting of pretraining on c4 with the evaluation on math and question answering. @xcite propose the addition of new tokens preceding each cot step. @xcite generate rationales post every token to elucidate future thinking using reinforce learning. our work is in line with the above studies in introducing additional tokens. the differences lie in that our purpose is to alleviate the shortcut learning induced by teacher forcing, and the tokens are used to generate the semantic plan representations of the subsequent tokens. more importantly, we use a simple and general representation prediction method to guide the function learning of the custom tokens. ## experiments on graph path-finding the graph path-finding task, as introduced by bachmann and nagarajan (2024), involves a unique structure known as a path-star graph g(d, l, n ). each graph features a central node from which d distinct paths emerge, each comprising l nodes, including the central node. the parameter n represents the range of node values, randomly selected from the set {0, 1, ..., n -1}, and may exceed the total number of nodes in the graph. the input of language models includes all of the edges of the star-graph, the start node, and the end node. the objective is to accurately predict the sole correct path between the designated start and end nodes. in particular, both the training and test graphs are derived from the same distribution, maintaining consistent topology characterized by fixed values of d, l, and n . this setup ensures that the observed failures are attributable to in-distribution errors rather than lack of compositional or length generalization capabilities. given that each graph is uniquely labeled and features a randomized adjacency list, the model is required to deduce a general algorithmic solution. following bachmann and nagarajan (2024), the dataset comprises 200,000 training samples and 20,000 test samples. the number of node values n is set as the product of l and d, facilitating a diverse range of graph instantiations. ## experiments of pretraining in this section, we extend the proposed model to pre-training, and validate its effectiveness in terms of perplexity, in-context learning, and supervised fine-tuning. ## conclusion in this paper, we presented semformer which explicitly models semantic planning in addition to next-token prediction. semformer introduces a sequence of trainable planning tokens to induce the planning within the internal computation, and the planning tokens in the language model are supervised by predicting the latent planning representations generated by an autoencoder. the results on the graph path-finding problem show that semformer can achieve nearly perfect accuracy in such a minimal lookahead task, alleviating the shortcut learning caused by teacher forcing. extending semformer to a general pertaining on openwebtext demonstrates the advantages of the paradigm. future research will focus on validating our model with larger sizes and training corpus and exploring its application on reasoning-related tasks such as math and coding. additionally, investigating hierarchical or block-wise prediction of semantic vectors presents a promising avenue for further exploration.