Is Next Token Prediction Sufficient for GPT? Exploration on Code Logic Comprehension

TL;DR

This work questions whether next-token prediction is sufficient for genuine code logic understanding in GPT-like LLMs. It introduces a diagnostic task, Logically Equivalent Code Selection, to probe underlying code logic, and a Next Token Prediction+ pretraining objective to align embeddings with semantic structure without compromising generation. Empirical results show that current models struggle on logic-based distinctions but substantially improve after Next Token Prediction+ pretraining, attaining stronger performance on both the logic task and code completion. The findings suggest that augmenting pretraining with logic-aware objectives can meaningfully enhance code reasoning capabilities in autoregressive models, with practical implications for software engineering tasks.

Abstract

Large language models (LLMs) has experienced exponential growth, they demonstrate remarkable performance across various tasks. Notwithstanding, contemporary research primarily centers on enhancing the size and quality of pretraining data, still utilizing the next token prediction task on autoregressive transformer model structure. The efficacy of this task in truly facilitating the model's comprehension of code logic remains questionable, we speculate that it still interprets code as mere text, while human emphasizes the underlying logical knowledge. In order to prove it, we introduce a new task, "Logically Equivalent Code Selection," which necessitates the selection of logically equivalent code from a candidate set, given a query code. Our experimental findings indicate that current LLMs underperform in this task, since they understand code by unordered bag of keywords. To ameliorate their performance, we propose an advanced pretraining task, "Next Token Prediction+". This task aims to modify the sentence embedding distribution of the LLM without sacrificing its generative capabilities. Our experimental results reveal that following this pretraining, both Code Llama and StarCoder, the prevalent code domain pretraining models, display significant improvements on our logically equivalent code selection task and the code completion task.

Figures (5)

• Figure 1: The hypothesis of different perspective between human and LLM about code. We hypothesis that humans tend to understand the logical structure behind the code (the part in the blue box) while the model tends to treat it as unordered keywords(the part in the gray box).
• Figure 2: The accuracy distribution of different models on different embedding similarity analysis subtasks. The blue, orange, and gray columns respectively represent the accuracy of large language models of different sizes (7B, 13B, 15B) on the six distance analysis tasks we have designed. Additionally, we have included OpenAI's specialized embedding model, text-embedding-ada-002 (depicted by a yellow column), as a control for comparison.
• Figure 3: The next token prediction task on line shuffled code. In the illustrative diagram, the white box symbolizes the normal token, while the orange box represents the shuffled token. The intermediate output is pointed towards a uniquely designed "void" token <v>. This methodology enables the model to bypass the problematic portions of the code, whilst retaining its comprehension of the normal sections.
• Figure 4: The next token prediction task on obfuscated code. This task continues to adhere to the training format of the causal language model. In the depicted diagram, the white boxes denote normal tokens, while the orange boxes represent obfuscated tokens. Throughout the training process, the obfuscated code acquires the output embedding via the model, and attempts to optimize the loss with the target token through the linear layer. This technique aims to diminish the output distribution gap between the obfuscated code and the corresponding normal code.
• Figure 5: Cosine Similarity Visualization. The left image represents the sentence embedding from Code Llama 7B model, while the right image depicts the result that after our continued pretraining. We differentiate the similarity distance values of various types of codes to the original code with distinct colors. The auxiliary line denotes their average similarity value.