Selective Prompt Anchoring for Code Generation

Yuan Tian, Tianyi Zhang

TL;DR: A novel approach to adjust the influence of selected prompt tokens and improve code generation performance by addressing the attention dilution issue.

Abstract: Recent advances in large language models (LLMs) have transformed software development by automatically generating code from natural language. Yet challenges remain in generating fully correct code that aligns with user intent. Our study reveals that LLMs tend to pay less attention to user prompts as more code tokens are generated. We hypothesize that this attention dilution issue is an important reason for code generation errors. To mitigate this issue, we propose ***S**elective **P**rompt **A**nchoring* (SPA) to guide code LLMs to pay more attention to user intent when generating code. We evaluate SPA using six base LLMs across six benchmarks. Our results demonstrate that SPA enhances Pass@1 by up to 12.9%, consistently outperforming SOTA code generation methods in all settings. Our code is available at https://github.com/magic-YuanTian/Selective-Prompt-Anchoring.

Lay Summary: Large language models (LLMs) are changing the way people write code, making it possible to generate programs simply via natural language. However, LLMs sometimes fail to capture the user's intent and make mistakes in the generated code. We identified a surprising reason for code generation mistakes: as the model writes more code, it gradually “forgets” to focus on the user’s original instructions (i.e., user prompt). To fix this, we introduce a model-agnostic and training-free mechanism to control LLMs’ attention. We use this attention-controlling mechanism to ensure that LLMs attend to users’ instructions throughout the code generation process. We show that this idea significantly and consistently improves code generation performance. To help other researchers explore more ideas, we have released a general and easy-to-use API for controlling the attention of any LLMs.

Link To Code: https://github.com/magic-YuanTian/Selective-Prompt-Anchoring

Primary Area: Deep Learning->Attention Mechanisms

Keywords: Large Language Models (LLMs), Prompt, Code Generation, Self-attention, Taylor expansion, Logits, Anchoring

Flagged For Ethics Review: true

Submission Number: 7967