Breaking the Attention Trap in Code LLM s: A Rejection Sampling Approach to Enhance Code Execution Prediction (2025)

Abstract: Code-specific Large Language Models (Code LLMs) have greatly improved performance across code-related tasks, offering substantial benefits in practical applications. However, existing research reveals significant performance bottlenecks in Code Execution tasks, which requires models to predict the execution results of given code snippets. This study identifies that, the Attention Trap phenomenon in training data constitutes a key constraint on model performance. To address this phenomenon, we propose the Attention Cracking with Rejection Sampling (AC-RS) method. The method first applies structural optimization to training data to eliminate attention traps. Then, it conducts secondary training on the outputs generated by the fine-tuned model to mitigate potential negative impacts from manual data intervention. Experimental results show that AC-RS significantly enhances the accuracy of Code Execution while preserving models’ original capabilities. Notably, the optimized 7B model achieves Code Execution accuracy comparable to 32B model and GPT-4o.

## introduction with the rapid advancement of large language models (llms) @xcite @xcite @xcite @xcite , code llms have attracted substantial academic and industrial attention due to their applicability and broad potential. from early models like starcoder @xcite and codellama @xcite to recent advancements including deepseek coder @xcite and qwen coder @xcite , code llms have shown remarkable performance across code-related tasks. however, studies (austin et al., 2021; nye et al., 2021; gu et al., 2024) indicate that current code llms underperform in code execution tasks. austin et al. (2021) reveals that even 137b model struggles to predict execution results of basic python code, and fine-tuning only provides minimal performance gains. nye et al. (2021) attributes this to the lack of explicit step-by-step reasoning before giving the predicted results. while previous work focuses on reasoning deficiencies, our work reveals that attention traps in widely-used leetcode training data fundamentally constrain execution prediction capabilities. when models process training data with lexical similarities, their attention mechanisms become overly focused on surface-level token correlations while neglecting deeper abstract relationships between data components. lexical similarity-induced cognitive bias exhibits universality in deep learning systems @xcite @xcite . for instance, models tend to misclassify samples containing categorical lexical cues (e.g., texts with "cinema" being erroneously categorized as "film"). we term this phenomenon "attention trap" in code execution training scenarios utilizing leetcode data, and investigate how the attention trap affects the training process. figure 1 demonstrates how trained model distributes attention during predicting code execution result token. the target outputs in example sections (e.g., the values 5 and 8 in "output: 5" and "output: 8") exactly match the current token, which attracts high attention weights. during learning, models excessively attend to these target outputs in the input queries, preventing proper modeling of the multi-step reasoning chain connecting problem descriptions, program code, and execution results. full example are provided in appendix a. to eliminate attention traps and analyze its effects on training process, we propose attention cracking with rejection sampling. our method contains two stages: (1) attention cracking (ac) modifies training data to eliminate attention traps; ## related works the field of code llms originated from datacentric methodologies and has gradually developed into a thriving research area @xcite . early studies in code-related domains adopted data construction methods from general-purpose domains. for instance, chaudhary (2023) employed the self-instruct @xcite approach to automatically generate code instruction dataset codealpaca. luo et al. (2023) further enhanced this dataset through evol-instruct (xu et al., 2023), training the wizardcoder model. additionally, magicoder (wei et al., 2024) attempted to generate high-quality instruction tuning data using opensource code. as data-related challenges were progressively addressed, multiple high-performance open-source code models emerged. representative examples include the qwen coder series and deepseek coder series. concurrently, researchers achieved notable progress in other dimensions of code-related tasks. frameworks like mftcoder (liu et al., 2024a) and models like phi (abdin et al., 2024) advanced the field through multi-task training strategies and parameter efficiency improvements, respectively. ## method this section details the implementation of ac-rs method. to ensure fair comparison and better prepare for subsequent training data generation, we first train lc-base model using leetcode data through rejection sampling (① in figure 2 ). ac-rs method then introduces two formal stages: attention cracking and rejection sampling. the full pseudocode of ac-rs is provided in appendix c. ## conclusion our study proposes ac-rs method. the ac stage eliminates attention traps in training data through data restructuring. the rs stage addresses performance degradation by training models with selfgenerated outputs. experimental results demonstrate that our 7b model trained with ac-rs achieves superior performance on livecodebench. notably, it matches the code execution accuracy of 32b parameter model and performs comparably to gpt-4o.