Selective Prompt Anchoring for Code Generation

## Abstract

Abstract Recent advances in large language models (LLMs) such as Copilot and ChatGPT have transformed software development by automating coding tasks. Despite these advancements, challenges remain in reducing error rates and fully meeting user expectations. Our empirical study reveals LLMs tend to dilute their self-attention on the initial prompt as more code tokens are generated.
We hypothesize this self-attention dilution issue is one of the root causes of inaccuracies in LLM-generated code.
To mitigate this issue, we propose S elective P rompt A nchoring ( Spa ). Spa amplifies the influence of the selected parts in the initial prompt, which we refer to as “anchored text”, during code generation. Specifically, Spa calculates the logit distribution difference with and without the anchored text. We prove this difference approximates the anchored text’s contextual contribution to the output logits. Spa creates an augmented logit distribution by linearly combining the original logit distribution and the logit difference.
We evaluate Spa with five LLMs on four benchmarks. Our results demonstrate that using Spa can consistently improve Pass@1 rates by up to 9.7% in all settings. Notably, with selective text anchoring, a small version of DeepSeek-Coder (6.7B) can achieve better performance than an original much larger version (33B). Our code is available at https://github.com/magic-YuanTian/Selective-Prompt-Anchoring .

## 1 Introduction

Figure: Figure 1: Pipeline of Spa
Refer to caption: extracted/5804348/figures/pipeline.png

Large language models (LLMs) have emerged as powerful tools that assist programmers by automating parts of the coding process. These models leverage the vast capabilities of LLMs to interpret task descriptions in natural language and generate corresponding code. While LLMs have achieved unprecedented performance on coding tasks, they still produce incorrect solutions to many tasks or generate code that does not fully meet user expectations. The prevalence of such generation errors undermines their reliability and limits their utility in real-world software development.

To improve the performance of LLMs on coding tasks, many efforts have been made to develop high-quality training data [, , ] and design new domain-specific training objectives [, ]. However, these approaches require tremendous computational resources.
Training-free approaches have been explored to address this challenge by enhancing the prompting method or incorporating external knowledge, such as retrieval-augmented generation [], chain-of-thoughts [, , ], self debugging and planning [, ], etc.
While they have been proven to be effective in improving model performance, there exist limitations such as being sensitive to the quality of prompt design and retrieved data [].
On top of these great attempts, this work aims to study and improve LLMs in an orthogonal direction—adjusting LLMs’ attention.

The impressive performance of transformer-based LLMs is largely attributed to its self-attention mechanism [], which enables models to focus dynamically on crucial parts of the given prompt.
Despite the success of the self-attention mechanism, prior works found Language models exhibit simple attention patterns [, ].
Furthermore, an empirical study [] found that given a coding task, there often exists a misalignment between LLM attention and human attention.
LLMs often focus on different parts of a natural language description as human programmers do when generating code.
Another work [] shows that when the model’s attention aligns more closely with human programmers’ attention, the model generates a more accurate SQL query.
Inspired by these findings, we hypothesize that a root cause of inaccuracy in LLM-generated code stems from the suboptimal model attention.
To verify our hypothesis, we conduct an empirical study that analyzes the shift in LLMs’ attention distribution during code generation.
We observe that LLMs’ attention to the initial prompt gradually dilutes as generating more code.
This is also supported by another study [], which observes that self-attention can be diluted as the input to the transformer becomes longer.

In standard decoding algorithms, LLMs calculate a conditional probability for the next token based on the preceding context.
However, the autoregressive nature of LLMs considers both the initial prompt and possibly wrong self-generated tokens together as the correct context and pays comparable attention to them.
We argue LLMs should pay more attention to the absolutely correct prompt and less attention to the following self-generated content that could potentially be wrong.

To mitigate this limitation, we propose Selective Prompt Anchoring (Spa), a model-agnostic approach that optimizes LLMs’ attention by amplifying the contextual contribution of selective prompt, towards each generated token.
Spa is inspired by the anchoring effect [] in psychology, which refers to people being influenced by specific information given before decision-making.
In Spa, we refer to this information as anchored text, a group of selected tokens within the prompt.

Figure 1 illustrates the pipeline of Spa.
Given the anchored text, Spa creates an original embedding matrix as well as a masked embedding matrix by replacing the embeddings corresponding to anchored text with mask embeddings (\footnotesize1⃝ & \footnotesize2⃝).
We mathematically show that the anchored text’s contextual contribution can be approximately measured by the difference between the logit distribution generated from the original prompt and the prompt with the anchored text masked (\footnotesize3⃝).
To amplify the contribution of anchored text in the model output, Spa multiplies this logit distribution difference by a weighting hyper-parameter (\footnotesize4⃝), and then adds it to the original logit distribution (\footnotesize5⃝).
We find the optimal weighting hyper-parameter varies across different models and tasks, but it is easy to tune.
To select the anchored text, our ablation study shows anchoring the natural language description alone (exemplified in Figure 1) performs the best for code generation tasks.

We evaluate Spa on four benchmarks with five different-sized LLMs.
The result shows Spa can significantly and consistently boost the Pass@$1$ Rate across all models and benchmarks, highlighting a new direction for controlling LLMs’ high-level attention and effectively improving performance.

## 2 An Empirical Analysis of Attention Dilution

We first conduct an empirical study to analyze the attention dilution issue in LLMs during code generation. To improve the generalizability of our findings, we experimented with two different methods to compute the attention scores over input tokens. First, we used a self-attention-based method [, , , , ] to obtain scores from the last self-attention layer of LLMs. Second, we used a gradient-based method [, , ] that treats the entire LLM as a function and measures to what extent each input token contributes to the output.
Based on these two methods, we calculate the percentage of attention on the initial prompt. Calculation details are provided in Appendix A.1.

Figure: Figure 2: Shift of LLMs’ self-attention to the initial prompt. The attention is calculated from the last self-attention layer of the LLM.
Refer to caption: extracted/5804348/figures/attention_change_self.png

We experiment on five LLMs [, , ] with various numbers of parameters on HumanEval [], a widely-used benchmark for code generation.
Figure 3 and Figure 3 show the evolution of the density of LLMs’ attention on the initial prompt when generating the first 400 tokens.
The results demonstrate that as the model generates more tokens, model attention on the initial prompt gradually becomes smaller, which we refer to as attention dilution. As a result, the code generation process becomes more influenced by the code tokens generated in previous time steps, rather than the prompt from users as the generated code sequence becomes longer. This can be problematic in two ways. First, generation errors in the previous time steps are very likely to propagate to the following steps as the model pays more attention to the preceding code tokens. Second, for complex tasks that require the generation of a long code sequence (e.g., multiple if statements), the model is likely to miss critical descriptions as it pays little attention to the user prompt deep in the code generation process.

## 3 Approach

### 3.1 Limitations of Standard Inference

Given a transformer-based large language model (LLM) denoted as ${f_{\theta}}$ and an initial prompt $x$, the model generates tokens autoregressively by calculating conditional probabilities for the next token $t_{i}$ at step $i$. Let $t_{1},t_{2},\ldots,t_{i-1}$ be the sequence of tokens already generated by the model.

The input to the model ${f_{\theta}}$ is an $n\times m$ embedding matrix, where $n$ stands for the dimension of each embedding (column) and $m$ represents the fixed context window size.
Each embedding is mapped from the corresponding token in the prompt.
Formally,

$$ $\mathbf{E}_{i}=\text{embedding}(x,t_{1},t_{2},\ldots,t_{i-1})=[\mathbf{E}^{x}, e_{1},e_{2},\ldots,e_{i-1},\texttt{PAD}].$ (1) $$

Here, $\mathbf{E}_{i}$ is the input embedding matrix to the model ${f_{\theta}}$ at step $i$. $\text{E}^{x}$ is the collection of embeddings for the tokens in the prompt $x$, serving as a submatrix of $\mathbf{E}_{i}$. $e_{i},\ldots,e_{i-1}$ are the embeddings corresponding to generated tokens $t_{i},\ldots,t_{i-1}$ respectively.
PAD is the padding submatrix of $\mathbf{E}_{i}$, composed of repetitive padding embeddings.
Then $\mathbf{E}_{i}$ is then injected with positional encoding.
Since the positional encoding does not affect our approach, we ignore it in the following computation for simplicity.

At each step, the model produces logits, denoted as ${f_{\theta}}(\mathbf{E}_{i})$, which are the raw outputs of LLM’s final layer prior to normalization.
Each logit element corresponds to a token in the vocabulary.
The logits ${f_{\theta}}(\mathbf{E}_{i})$ are then transformed into probabilities using the softmax function, ensuring it is a valid probability distribution over the vocabulary:

$$ $P_{\theta}(t_{i}|x,t_{0},t_{1},\ldots,t_{i-1})=\text{softmax}({f_{\theta}}( \mathbf{E}_{i})).$ (2) $$

Based on the conditional probability distribution for tokens at step $i$, a certain sampling method is applied to select the particular token. For example, greedy sampling simply selects the token with the highest probability at each step:

$$ $t_{i}=\arg\max_{t}P_{\theta}(t|x,t_{0},t_{1},\ldots,t_{i-1}).$ (3) $$

However, the model treats the initial prompt $x$ and self-generated content $t_{0},t_{1},\ldots,t_{i-1}$ equally as the input.
The self-attention dilution issue may arise as the number of self-generated tokens increases.
It can potentially cause the model to overlook the initial prompt $x$, pay excessive attention to self-generated content $t_{0},t_{1},\ldots,t_{i-1}$, and gradually deviate from the original purpose.

### 3.2 Semantic Adjustment

To mitigate the attention dilution issue, we propose Selective Prompt Anchoring (Spa) to augment the output logits by amplifying the contextual contribution of the selective tokens within the prompt, which we refer to as “anchored text”.

Spa introduces the mechanism of adjusting the semantic impact of arbitrary groups of token embeddings in the input matrix $\mathbf{E}_{i}$ towards the output logits ${f_{\theta}}(\mathbf{E}_{i})$.
For simplicity, here we make the entire initial prompt $x$ as the anchored text.

$\mathbf{E}_{i}$ is an $n\times m$ input embedding matrix at step $i$, and $\mathbf{E}^{x}$ represents a $n\times k$ submatrix within $\mathbf{E}_{i}$ covering the first $k$ columns (corresponding to the prompt $x$). They are visualized below:

$$ $\mathbf{E}_{i}=\left[\begin{array}[]{cc}\underbrace{\begin{array}[]{ccc}e_{11} &\cdots&e_{1k}\\ e_{21}&\cdots&e_{2k}\\ \vdots&\ddots&\vdots\\ e_{n1}&\cdots&e_{nk}\end{array}}_{\mathbf{E}^{x}}&\begin{array}[]{ccc}e_{1,k+1 }&\cdots&e_{1m}\\ e_{2,k+1}&\cdots&e_{2m}\\ \vdots&\ddots&\vdots\\ e_{n,k+1}&\cdots&e_{nm}\end{array}\end{array}\right].$ (4) $$

We construct two $n\times m$ matrices, $\mathbf{X}$ and $\mathbf{G}_{i}$, which add up to $\mathbf{E}_{i}$.
Matrix $\mathbf{X}$ is created by preserving the first $k$ columns of $\mathbf{E}_{i}$ corresponding to $\mathbf{E}^{x}$ and setting all other columns to zero (note that $\mathbf{E}^{x}$ and $\mathbf{X}$ remain unchanged during generating new tokens).
Matrix $\mathbf{G}_{i}$ is constructed by setting the first $k$ columns of $\mathbf{E}_{i}$ that correspond to $\mathbf{E}^{x}$ to zero, and retaining all other elements from the remaining columns.
They are visualized as follows:

$$ $\mathbf{X}=\begin{bmatrix}e_{11}&e_{12}&\cdots&e_{1k}&0&\cdots&0\\ e_{21}&e_{22}&\cdots&e_{2k}&0&\cdots&0\\ \vdots&\vdots&\ddots&\vdots&\vdots&\ddots&\vdots\\ e_{n1}&e_{n2}&\cdots&e_{nk}&0&\cdots&0\\ \end{bmatrix},\mathbf{G}_{i}=\begin{bmatrix}0&0&\cdots&0&e_{1,k+1}&\cdots&e_{1 m}\\ 0&0&\cdots&0&e_{2,k+1}&\cdots&e_{2m}\\ \vdots&\vdots&\ddots&\vdots&\vdots&\ddots&\vdots\\ 0&0&\cdots&0&e_{n,k+1}&\cdots&e_{nm}\\ \end{bmatrix}.$ (5) $$

The sum of $\mathbf{X}$ and $\mathbf{G}_{i}$ reconstructs the original matrix $\mathbf{E}_{i}$:

$$ $\mathbf{E}_{i}=\mathbf{X}+\mathbf{G}_{i}.$ (6) $$

Suppose we want to amplify the semantic impact of the submatrix $\mathbf{X}$ by a weighting value ${\omega}$.
${\omega}>1$ indicates amplification of semantics, while ${\omega}<1$ indicates diminishing the semantics.
We define a semantic adjustment function ${\Phi}(\mathbf{X},{\omega})$ that scales the influence of $\mathbf{X}$ by ${\omega}$.
Here we assume the semantic consistency []: the scaled embeddings maintain their semantic interpretations in a manner comprehensible by the model ${f_{\theta}}$.
Note that the original embedding matrix $\mathbf{E}_{i}$ corresponds to when ${\omega}$ equals $1$:

$$ $\mathbf{E}_{i}={\Phi}(\mathbf{X},1)+\mathbf{G}_{i}.$ (7) $$

Given ${f_{\theta}}$ is differentiable for backpropagation,
to amplify the semantic impact of the anchored prompt $x$ in the final logits,
it is essentially calculating the integral along with the directional derivative of $\mathbf{X}$. Let ${F_{\theta,i,x}}({\omega})$ represent the augmented logits calculated by model ${f_{\theta}}$ at step $i$, where the impact of anchored text $x$ is scaled by ${\omega}$.
Formally,

$$ $\displaystyle{F_{\theta,i,x}}({\omega})$ $\displaystyle={f_{\theta}}({\Phi}(\mathbf{X},{\omega})+\mathbf{G}_{i})$ (8) $\displaystyle\approx{f_{\theta}}(\mathbf{E}_{i})+\int_{1}^{\omega}\nabla_{ \mathbf{X}}{f_{\theta}}({\Phi}(\mathbf{X},t)+\mathbf{G}_{i})\,dt$ (9) $\displaystyle={F_{\theta,i,x}}(1)+\int_{1}^{\omega}\frac{d{F_{\theta,i,x}}(t)} {dt}\,dt$ (10) $\displaystyle={F_{\theta,i,x}}(0)+\int_{0}^{\omega}\frac{d{F_{\theta,i,x}}(t)} {dt}\,dt,$ (11) $$

where $t$ is the variable of integration.

### 3.3 Augmented Logits by Approximation

Given the computational complexities of LLMs, directly solving $\int_{0}^{\omega}\frac{d{F_{\theta,i,x}}(t)}{dt}\,dt$ is impractical.
Therefore, we approximate it by employing the Taylor expansion:

$$ ${F_{\theta,i,x}}({\omega})={F_{\theta,i,x}}(0)+{\omega}\cdot{F_{\theta,i,x}}^{ \prime}(0)+\frac{{\omega}^{2}}{2!}{F_{\theta,i,x}}^{\prime\prime}(0)+\ldots$ (12) $$

Since LLMs are inherently non-linear, higher-order derivatives of the logits function are non-zero.
However, for computational efficiency, we truncate the series after the first derivative, yielding:

$$ ${F_{\theta,i,x}}({\omega})\approx{F_{\theta,i,x}}(0)+{\omega}\cdot{F_{\theta,i ,x}}^{\prime}(0),$ (13) $$

where the the integral part $\int_{0}^{\omega}\frac{d{F_{\theta,i,x}}(t)}{dt}\,dt$ in equation 11 is approximated by ${\omega}\cdot{F_{\theta,i,x}}^{\prime}(0)$.
Here we assume the logit function ${F_{\theta,i,x}}$ is smooth and differentiable so it is possible to compute the first-order derivative.

To calculate ${F_{\theta,i,x}}(0)$, we mask tokens in the anchored text $x$ using masked embeddings.
Each LLM provides at least one special token reserved for text masking, which almost has no semantic influence, e.g., <unk> for Code Llama [] and <pad> for DeepSeek-Coder [].
Each special token corresponds to a masked embedding.
By replacing embeddings of $x$ with masked embeddings, we get a masked input matrix $\mathbf{E}^{mask}_{i}$. It ablates the semantic influence of the anchored text $x$ while the positional encoding is not affected.
Thus, we can get

$$ ${F_{\theta,i,x}}(0)={f_{\theta}}(\textbf{E}^{mask}_{i}).$ (14) $$

To calculate ${F_{\theta,i,x}}^{\prime}(0)$, we use finite-difference methods to get an approximation. Assuming the interval of $1-0$ is sufficiently small for ${F_{\theta,i,x}}$, we get:

$$ ${F_{\theta,i,x}}^{\prime}(0)\approx\frac{{F_{\theta,i,x}}(1)-{F_{\theta,i,x}}( 0)}{1-0}.$ (15) $$

Combining equation 13, 14, and 15, we get the augmented logits by first-order approximation:

$$ $\displaystyle{F_{\theta,i,x}}({\omega})$ $\displaystyle\approx{F_{\theta,i,x}}(0)+{\omega}\cdot({F_{\theta,i,x}}(1)-{F_{ \theta,i,x}}(0))$ (16) $\displaystyle={\omega}\cdot f_{\theta}({\Phi}(\mathbf{X},1)+\mathbf{G}_{i})+(1 -{\omega})\cdot f_{\theta}({\Phi}(\mathbf{X},0)+\mathbf{G}_{i})$ (17) $\displaystyle={\omega}\cdot f_{\theta}(\textbf{E}_{i})+(1-{\omega})\cdot{f_{ \theta}}(\textbf{E}^{mask}_{i}).$ (18) $$

Based on the augmented logits ${F_{\theta,i,x}}({\omega})$ where the impact of the anchored text is adjusted by a given value ${\omega}$, a certain sampling algorithm is applied to select the particular token.
Spa can be used to augment different existing sampling methods, including greedy sampling, beam search, nucleus sampling [], and more.

### 3.4 Higher Order Approximation

It is feasible to calculate higher-order approximation.
For example, if we want to keep the term with the second-order derivative $\frac{{\omega}^{2}}{2!}{F_{\theta,i,x}}^{\prime\prime}(0)$ in Equation 12, it can still be computed using finite-difference methods:

$$ ${F_{\theta,i,x}}^{\prime\prime}(0)\approx\frac{{F_{\theta,i,x}}(1)-2{F_{\theta ,i,x}}(0)+{F_{\theta,i,x}}(-1)}{(1-0)^{2}}.$ (19) $$

${F_{\theta,i,x}}(-1)$ can be solved by Equation 16 where ${F_{\theta,i,x}}(0)$ and ${F_{\theta,i,x}}(1)$ are the logits generated from the original input and the logits generated from the masked input.
Similarly, whatever how many terms we keep in the Taylor expansion, it is essentially a linear combination of ${F_{\theta,i,x}}(0)$ and ${F_{\theta,i,x}}(1)$.
For simplicity, Spa calculates augmented logits using the first-order approximation in this work.

### 3.5 Tuning Weighting Values

The weighting value ${\omega}$ serves as a hyperparameter.
Our experiments demonstrate an unimodal relationship between ${\omega}$ and the performance.
As the weight ${\omega}$ increases, the performance first improves, reaching an optimum, and then declines with further increases of ${\omega}$.
It is simple to tune this single parameter.
More details are discussed in Section 5.2.

### 3.6 Selective Anchored Text

While Spa can anchor the entire prompt, the initial prompt can significantly vary due to task differences. In some scenarios, the prompts can be lengthy and not all information is important.
According to our observation, we find narrowing down anchored text can lead to better performance.
This is intuitive since anchoring too many tokens may decrease their impact on differentiating the logit distribution, e.g., anchoring all the tokens will make no difference.
Our goal is to identify and anchor the most informative tokens, which LLMs should consistently focus on, while excluding trivial details in the prompt.
For code generation tasks, the prompt commonly comprises four possible components: (1) Natural language instruction or docstring; (2) Starting code snippet; (3) A list of test cases; (4) Few-shot examples.
Intuitively, natural language instruction provides high-level guidance that LLM should continually consider.
Our ablation study in Section 5.3 confirms it and shows anchoring the natural language instruction alone yields the best performance.

## 4 Experiments

### 4.1 Comparison Baselines

Spa requires access to the full logits generated by the large language models (LLMs), so we are unable to evaluate closed-source models, such as ChatGPT.
We select five representative open-source code LLMs: CodeGen-Mono [] (350M), CodeLlama [] (7B), and three different-sized DeepSeek-Coder-Instruct [] (1.3B, 6.7B, 33B) models.
These models have been fine-tuned for code generation tasks.
Notably, the DeepSeek-Coder-Instruct models have been fine-tuned by instruction-tuning [], while CodeGen-Mono and CodeLlama are standard text completion models.
This selection aims to cover diverse SOTA code LLMs of different types and sizes.

### 4.2 Benchmarks

HumanEval []. It includes 164 Python tasks designed by OpenAI developers. It was initially designed to evaluate Codex [] and has since become a common benchmark for code generation.

MBPP []. It includes 974 crowd-sourced Python tasks.
However, due to the crowd workers’ ambiguous or insufficient task descriptions, the MBPP authors created a sanitized version containing 427 tasks with clearer descriptions. We evaluate Spa on the sanitized version.

HumanEval+ and MBPP+. Although HumanEval and MBPP are considered de facto standards for assessing code LLMs, a recent study [] found they lack sufficient test cases and precise problem descriptions.
This has been demonstrated as an issue that can lead to an unreliable assessment of LLM-generated code [].
Liu et al. [] subsequently released HumanEval+ and MBPP+, which supplement HumanEval and MBPP with additional test cases and better instruction.
We also evaluate Spa performance on HumanEval+ and MBPP+.

### 4.3 Experiment Setup and Metrics

Model Deployment.
We download and deploy LLMs from Huggingface.
To speed up evaluations, we apply 8-bit quantization [, ] to all the models.
Prior studies [, ] show 8-bit quantization has very little impact on LLM performances.
All the experiments are conducted on eight A5500 GPUs (each with 24GB VRAM) with 192GB memory for five weeks.

Prompt Design.
We use the original task descriptions from the datasets as prompts for the text-completion models, CodeLlama and CodeGen-Mono.
For the three DeepSeek-Coder-Instruct models, we format the prompts using the official chat template from HuggingFace. All experiments are conducted in a zero-shot setting.

Evaluation Metric. Following prior works [, , ], we measure model performance using the Pass@$k$ metric, which measures whether any of $k$ candidates can pass all the test cases.
As a common setting, we set $k$ to 1 in our experiment.
For each task, LLMs only generate one code snippet using greedy sampling. The generated code is considered correct when it passes all the test cases.

## 5 Results

### 5.1 Main Results

**Table 1: Pass@$1$ rates (%) with and without using Spa**
| Model | #Params | HumanEval | HumanEval+ | MBPP | MBPP+ |
| --- | --- | --- | --- | --- | --- |
| CodeGen-Mono [] | 350M | 12.8 | 10.4 | 19.7 | 15.9 |
| + Spa |  | 14.6 (+1.8) | 14.0 (+3.6) | 26.3 (+6.6) | 20.6 (+4.7) |
| DeepSeek-Coder [] | 1.3B | 62.8 | 59.8 | 58.3 | 52.4 |
| + Spa |  | 66.5 (+3.7) | 63.4 (+3.6) | 61.9 (+3.6) | 53.4 (+1.0) |
| DeepSeek-Coder [] | 6.7B | 73.8 | 70.7 | 67.2 | 58.6 |
| + Spa |  | 83.5 (+9.7) | 76.8 (+6.1) | 72.2 (+5.0) | 61.2 (+2.6) |
| CodeLlama [] | 7B | 29.9 | 25.0 | 50.8 | 40.8 |
| + Spa |  | 37.8 (+7.9) | 33.5(+8.5) | 53.1 (+2.3) | 43.1 (+2.3) |
| DeepSeek-Coder [] | 33B | 80.5 | 73.2 | 73.5 | 63.2 |
| + Spa |  | 84.8 (+4.3) | 77.1 (+3.9) | 77.3 (+3.8) | 68.5 (+5.5) |

We compare each LLM’s performance with and without using Spa in Table 1.
The results indicate that using Spa can consistently improve the Pass@$1$ rate across all benchmarks and all LLMs.
The improvement is up to 9.7% on Humaneval for DeepSeek-Coder (6.7B).
It’s worth noting that through selective text anchoring, the small version of DeepSeek-Coder (6.7B) surpasses a much larger version (33B).
This supports our hypothesis that the inaccuracies in LLMs mainly stem from poor attention, not their generative capabilities.
To demonstrate Spa can effectively anchor LLM’s attention on the initial prompt, we include two examples in Appendix A.3.

While Spa achieved a consistent improvement on LLMs with different sizes, we do not observe a monotonic relationship between model size and the improvement. Furthermore, there is no obvious correlation between the original model performance and the improvement. It is an interesting future direction to investigate how different model attributes affect the improvement achieved by Spa.

### 5.2 Analysis of Weighting Values

Figure: Figure 4: Analysis of weighting values
Refer to caption: extracted/5804348/figures/weight_analysis.png

Spa introduces a single weighting hyperparameter ${\omega}$, which is used to adjust the magnitude of the anchoring effect of Spa.
Figure 4 reports Pass@$1$ rates given various values of ${\omega}$ for each model and benchmark (${\omega}=1$ means the original model).
We observe that the optimal ${\omega}$ slightly varies across different models and benchmarks.
However, there roughly exists an unimodal relationship between the weight ${\omega}$ and the performance.
As the weight ${\omega}$ increases, the performance first improves, reaching an optimum, and then declines with further increases of ${\omega}$.
While this value can be easily tuned on a sampled validation set from each benchmark due to this feature, this work does not focus on evaluating the quality of validation sets.
We report the optimal weighting values in Appendix A.2.
We observe any values less than $1.25$ help improve performance.
Figure 4 also shows that while optimal weights may vary across different models, each model has close optimal weights on different datasets. This suggests that the optimal weight strongly depends on the model.

### 5.3 Analysis of Anchored Test Selection

**Table 2: Improvements of Pass@$1$ rates (values in %) for different anchored text**
| Anchored Text | HumanEval | HumanEval+ | MBPP | MBPP+ |
| --- | --- | --- | --- | --- |
| NL | + 5.48 | + 5.08 | + 4.26 | + 3.22 |
| NL + Test | + 5.11 | + 4.89 | + 4.05 | + 3.11 |
| NL + Code | + 4.87 | + 4.65 | - | - |
| NL + Code + Test | + 4.76 | + 4.57 | - | - |

As mentioned in Section 3.6,
we analyzed how anchored text selection impacts the performance of SPA.
Prompts in HumanEval and HumanEval+ include the function signature (referred to as Code), task descriptions in natural language (NL), and test cases (Test).
Prompts in MBPP and MBPP+ consist of task descriptions (NL) followed by test cases (Test). For HumanEval and HumanEval+, we create four conditions by removing test cases and source code. For MBPP and MBPP+, we create two conditions by removing test cases. Since task descriptions serve as the core of a prompt, we choose to anchor it in all conditions.

For each condition and benchmark, we calculate the average Pass@$1$ rate improvement across all five LLMs.
Table 2 shows that anchoring the task description alone yields the best performance.
This implies that anchoring more tokens in the prompt may not necessarily be helpful. Narrowing down the anchored text to fewer but critical, informative tokens can lead to better performance.

## 6 Related work

Code Generation.
In recent years, there has been rapid progress in the development of code generation approaches [, , , ] and benchmarks [, , , , , , , , ]. With the advent of large language models (LLMs), such as GPT-4 [], Gemini [], and the LLaMA family [, ], code generation has become a standard capability. Subsequent research [, , , , , , , , ] has focused on fine-tuning these pre-trained LLMs on coding datasets, achieving state-of-the-art performance across a variety of coding tasks.

Despite their remarkable ability to follow natural language instructions, LLMs still face challenges when generating long and complex code. To enhance the code generation capabilities of LLMs, recent studies have explored train-free approaches such as prompt engineering [, ], in-context learning [, , ], and retrieval-augmented generation [, ]. Additionally, self-debugging techniques [] enable LLMs to debug code based on error messages and execution results, while self-planning [] allows LLMs to decompose tasks into subtasks and implement solutions step-by-step. The chain-of-thought approach [, , ] facilitates a step-by-step reasoning process in LLMs.

Building on these advancements, Spa introduces an orthogonal approach particularly suitable for code generation. It can be integrated with existing methods to further improve performance.

Controllable Generation.
Compared to fine-tuning a language model (LM) at the decoding time, controllable generation aims to steer the pre-trained LMs to
match a sentence-level attribute (e.g., a topic on sports).
Existing approaches usually require additional models or training, such as fine-tuning a smaller LM [, , , ], a reward model [, ], or a fine-tuned model with controlling codes [, , ].
The mechanism used in Spa can also be used to control the generation by adjusting weights over the input text.
Compared to the aforementioned works, Spa does not require any additional models or training.

Logit Arithmetic.
There has been a growing body of methods that perform arithmetic on multiple logit distributions to enhance text generation. These methods include contrasting logits from multiple LMs [, , , ], logits of LMs of different sizes [], logits from different layers of a model [, ], and logits from the same model given different inputs [, , , , ]. Similar ideas have also been explored in diffusion models [, ].

Spa can be considered analogous to contrasting logits from the same model when given different inputs. However, we delve deeper into modeling and characterize Spa as a first-order approximation of the integral over arbitrary groups of embeddings.
Additionally, Spa is designed to mitigate the attention dilution issue in LLMs during code generation. By contrast, none of the aforementioned works explored code-generation tasks.
They primarily focus on reducing hallucinations [, , ], enhancing coherence [], factuality [], and controllable text generation [, , ].
Besides, Spa focuses on perturbation of the original prompt through masking rather than providing additional context [, , ] or changing to a completely new prompt [].
The most similar work to ours is Visual Contrastive Decoding [], which contrasts outputs derived from both original and distorted visual inputs. However, instead of distorting visual inputs, Spa focuses on masking textual inputs.
Furthermore, Spa focuses on perturbing selective parts in the prompt that are especially useful for code generation, rather than the entire input.

## 7 Limitations

We employed 8-bit quantized LLMs to expedite all experiments. Although this method has been shown to have minimal impact on performance, we did notice some degradation. Furthermore, we did not evaluate very large LLMs (e.g., more than 100B) due to computational constraints.
Despite the unimodal feature, it is infeasible to enumerate all the weighting hyperparameter ${\omega}$ on the continuous distribution. The real optimal ${\omega}$ should perform slightly better than the values reported in Section 4.
While our experimented benchmarks have become standards for evaluating code generation models, we did not evaluate Spa in real-world scenarios. More evaluations can be performed in future studies.

## 8 Conclusion & Future Work

In this paper, we propose Spa, a training-free approach designed to enhance the quality of code generated by large language models (LLMs) by mitigating the attention dilution issue. Spa employs a novel technique to adjust the influence of arbitrary groups of input tokens, based on a mathematical approximation. Our empirical studies indicate that LLMs may overlook the initial prompt as generating more new tokens. To mitigate this issue, Spa amplifies the impact of the initial prompt throughout code generation.
Our evaluations demonstrate that Spa consistently and significantly enhances the performance of code LLMs of various sizes on multiple benchmarks.

The effectiveness of Spa highlights its potential in other domains, especially for generation tasks.
Besides, the underlying principle of Spa is not confined to transformer-based LLMs and could be adapted for use in other model architectures (e.g., RNN).
In this work, we pre-define the method to select anchored tokens and make the weighting hyperparameter fixed when generating code. We believe this only serves as a baseline. In future work, both of anchored text and the weighting value can be dynamically determined according to different contexts and sampling stages.

## Appendix A Appendix / supplemental material

### A.1 Attention Calculation

Self-attention.
Most LLMs are based on the decoder of transformer [] which has multiple self-attention layers.
Roughly speaking, given an LLM ${f_{\theta}}$ and an input sequence of tokens $t_{0},t_{1},\ldots,t_{n}$ where $t_{i}$ represents the $i$th token.
The transformer calculates relevance scores between every pair of tokens. The self-attention score for a token $t_{i}$ in the sequence can be roughly formulated as:

$$ $\text{attention}(t_{i})\approx\frac{\sum_{j=1}^{n}\text{relevance}(t_{i},t_{j} )}{\sum_{i=1}^{n}\sum_{j=1}^{n}\text{relevance}(t_{i},t_{j})},$ (20) $$

where the relevance function approximates the computation among $Q,K,V$ in transformers [].
However, different layers have different attention distributions.
According to a study [], deeper self-attention layers can better capture long-distance dependencies and program structure, so we calculate the attention by aggregating attention from multiple heads at the last layer.
Nevertheless, this still excludes the influence from the last forward layer.

Gradient-based Attention.
Compared to using self-attention layers in transformers, the gradient-based method can be generalized to different model architectures and consider the entire model as a whole.
It computes the model’s attention by calculating the gradients relative to the input.
Intuitively, a token that induces a larger gradient is considered more influential, suggesting that the model pays greater attention to it.
Formally, the attention over the token $t_{i}$ is calculated by

$$ $\text{attention}(t_{i})=\frac{\partial{f_{\theta}}(t_{0},t_{1},\ldots,t_{n})}{ \partial t_{i}}.$ (21) $$

Attention Percentage to the Prompt.
Based on these two methods, we analyze how the attention of LLMs to the initial prompt shifts.
Formally, given the prompt $x$ and the following generated tokens $t_{0},t_{1},\ldots,t_{i-1}$, we calculate the percentage of attention $\alpha(x)$ over the initial prompt

$$ $\alpha(x)=\frac{\text{attention}(x)}{\text{attention}(x)+\sum^{n}_{i=1}\text{ attention}(t_{i})}$ (22) $$

Given attention analysis requires open sourcing, we select five SOTA code LLMs with various sizes.
We run the experiments on HumanEval [], one of the most popular benchmarks for evaluating code generation models.
We run five LLMs [, , ] on all 164 Humaneval tasks.
Figure 3 shows the self-attention shift and Figure 3 shows the gradient-based attention shift when generating the first 400 tokens.
The value gradually becomes noisy due to the lack of generated sequence with enough length.

The results demonstrate that there indeed exists such attention dilution issue.
Due to the autoregressive nature, LLMs’ attention to the initial prompt is gradually diluted as generating more code. LLMs tend to attend to code generated by itself.
Our finding is supported by another study [] which investigates the self-attention dilution of transformers in a more general scenario.

### A.2 Optimal Weighting Values

**Table 3: Optimal ${\omega}$ for each model and benchmark**
| Model | HumanEval | HumanEval+ | MBPP | MBPP+ | Average |
| --- | --- | --- | --- | --- | --- |
| CodeGen-Mono (350M) | 1.20 | 1.20 | 1.35 | 1.35 | 1.28 |
| DeepSeek-Coder (1.3B) | 1.05 | 1.05 | 1.20 | 1.20 | 1.13 |
| DeepSeek-Coder (6.7B) | 1.28 | 1.28 | 1.25 | 1.25 | 1.26 |
| CodeLlama (7B) | 1.60 | 1.60 | 1.20 | 1.20 | 1.40 |
| DeepSeek-Coder (33B) | 1.35 | 1.35 | 1.30 | 1.30 | 1.33 |
| Average | 1.30 | 1.30 | 1.33 | 1.33 | 1.28 |

Table 3 reports optimal weighting values ${\omega}$ that are used in our main results (Table 1).
We observe the average value of $0.28$ can be used to effectively improve performance across all benchmarks for all LLMs.

### A.3 Examples

Figure 5 presents two examples comparing the code generated by models alone and the models augmented using Spa.

In the first example, CodeLlama (7B) overlooks the specified condition "upper vowels." In contrast, Spa enhances the model’s focus on the intended purpose. The code initializes all the upper vowels in the first line and correctly refers to it later.

In the second example, DeepSeek-Coder (1.3B) erroneously sorts the list by string names instead of integers. When using Spa, the model demonstrates improved recognition of the required procedures, aligning more closely with the task description. The code correctly sorts and reverses the list. Then the integer list is mapped to the string list.

Figure: Figure 5: Examples of generated code by LLMs alone (left) and using Spa (right).
Refer to caption: extracted/5804348/figures/example.png

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 5
presentation: 2
questions: See the weaknesses section.
soundness: 2
strengths: - The overall structure is clear and easy to follow.
- The authors conduct experiments on several open Code LLMs like CodeGen, DeepSeek-Coder and Code Llama with different model sizes.
- The authors further provide ablation studies to validate the effectiveness of their methods from different perspectives.
summary: The authors propose a new method to improve the code generation quality of LLMs by enhancing the attention mechanism. To show the effectiveness of their method on HumanEval(+) and MBPP(+) datasets, they use 1/5 tasks in the datasets to set the hyperparameters and the other 4/5 tasks to evaluate the performance.
weaknesses: The contributions of the paper could be very limited, and the efficacy of the research is still questionable without further experiments.
Several weaknesses include:
- The motivation of the paper is not strong enough. For example, the "attention dillusion" phenomenon is not a new phenomenon and has been discussed in the literature [1-3]. It appears that the lack of attention is quite common in the existing LLMs, not just in code genaration. The authors should provide more motivation as to why they only focus on the code generation task.
- The efficacy of SPA is not very convincing. Essentially, SPA requires tuning/searching the hyperparameters (e.g., anchoring strength) on each benchmark, which makes it impractical.
- Although there is an ablation study on cross-dataset evaluation, it is still not enough to validate the effectiveness of SPA, as HumanEval and MBPP are in the same algorithmic paradigm. Widely-used open-domain code benchmarks like BigCodeBench [4] should be used to further validate the effectiveness of SPA.
- Section 5.3 shows that the anchoring weight affects the performance of SPA significantly, and results in various performance across models and datasets. It is unclear whether the proposed method can generalize without hyperparameter tuning.
- The evaluated Code LLMs are quite outdated. The authors should include more recent models like StarCoder2 [5] and DeepSeek-Coder-V2 [6].
- As pointed in A.5, SPA additionally takes 2 to 3.5 times longer than regular inference, with extra memory usage. This limitation further limits the usability of SPA.
- The current evaluation only focuses on Python-only and function-level code generation, and it is unclear how the proposed method can generalize to other programming languages and code generation tasks.


[1] Shi, F., Chen, X., Misra, K., Scales, N., Dohan, D., Chi, E. H., ... & Zhou, D. (2023, July). Large language models can be easily distracted by irrelevant context. In International Conference on Machine Learning (pp. 31210-31227). PMLR.

[2] Liu, N. F., Lin, K., Hewitt, J., Paranjape, A., Bevilacqua, M., Petroni, F., & Liang, P. (2024). Lost in the middle: How language models use long contexts. Transactions of the Association for Computational Linguistics, 12, 157-173.

[3] Zhang, Q., Singh, C., Liu, L., Liu, X., Yu, B., Gao, J., & Zhao, T. Tell Your Model Where to Attend: Post-hoc Attention Steering for LLMs. In The Twelfth International Conference on Learning Representations.

[4] Zhuo, T. Y., Vu, M. C., Chim, J., Hu, H., Yu, W., Widyasari, R., ... & Von Werra, L. (2024). Bigcodebench: Benchmarking code generation with diverse function calls and complex instructions. arXiv preprint arXiv:2406.15877.

[5] Lozhkov, A., Li, R., Allal, L. B., Cassano, F., Lamy-Poirier, J., Tazi, N., ... & de Vries, H. (2024). Starcoder 2 and the stack v2: The next generation. arXiv preprint arXiv:2402.19173.

[6] Zhu, Q., Guo, D., Shao, Z., Yang, D., Wang, P., Xu, R., ... & Liang, W. (2024). DeepSeek-Coder-V2: Breaking the Barrier of Closed-Source Models in Code Intelligence. arXiv preprint arXiv:2406.11931.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: 1. In Table 1, some SPA-optimal models perform slightly behind the SPA-turned models, for example, pass@10 of DeepSeek-Coder on HumanEval+ and MBPP. Are there any reasons for this? 

2. The authors used 15% of sampled data from a benchmark to tune the anchoring strength. If picking the 15% data was a random process, how stable is SPA with different sample sets?

3. As shown in Figure 4, the examined LLMs share similar trends regarding the performance of different anchoring strength values on the same dataset, is it possible to share/reuse anchoring strength among different LLMs?    

4. Many existing studies optimize the prompts to make LLMs more focused on the key information of the provided prompts. Are there specific reasons that authors do not compare SPA to these prompt-optimization-based approaches?

5. Does SPA perform significantly differently on short prompts compared to long prompts?
soundness: 3
strengths: + a new approach to improve LLMs for code generation 
+ bring significant performance improvement
summary: The paper proposes SPA, an approach to improve code generation in LLMs by addressing attention dilution, where models lose focus on the initial prompt during extended generation. The authors demonstrate the limitations of current LLMs in maintaining prompt relevance over generated sequences, potentially leading to inaccuracies. SPA amplifies the attention on selected prompt tokens, significantly improving performance across different LLMs and benchmarks.
weaknesses: - generalizability is limited to code generation with Python
- missing prompt optimization-based baselines
- the randomness in the process of tuning anchoring strength should be explored
- the impact of prompt length is unknown

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 2
questions: 1. It is recommended to include experimental comparisons with other LLM-based optimization approaches, rather than solely comparing with baselines. 
2. The authors mention identifying and anchoring the most informative tokens in longer prompts, thereby excluding trivial information. However, I didn't see any methods related to identifying fine-grained informative tokens  in the paper. 
3. The SPA method amplifies the influence of specific parts of the prompt. Will this approach change the model's behavior and compromise the initially correct output?
soundness: 2
strengths: 1. SPA, as a method that requires no training, can be applied to various models, demonstrating its broad applicability. 
2. The method has been validated across multiple code generation models, with experimental results showing consistent performance improvements among various models.
summary: This paper proposes that as the generated code tokens increase, LLMs often dilute their self-attention on the initial prompt, which is one of the fundamental reasons for inaccuracies in the code generation. To solve this issue, the paper introduces Selective Prompt Anchoring (SPA), a model-agnostic approach that amplifies the influence of selective prompts. The results demonstrate that, after tuning on a few dozen instances, SPA improves Pass@1 on new tasks by up to 7.6\%.
weaknesses: 1. This paper proposes to optimize the results by enhancing attention to anchored text. Similar methods have already been explored in natural language contexts, and it is recommended to include them in the "Related Work" section. 
2. The authors should specify which specific information should be selected as anchored text in section 3.5, or provide a method for segment identification.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: **1** What are the advantages of SPA over Anchor-LLM with improved attention mechanism? (I looked at the open source code and it seems that SPA only supports greedy search at the moment)

**2** How well does SPA generalize to other programming languages besides python?

**3** Why there is no consideration on how to choose the optimal anchored-text?
soundness: 3
strengths: **1 Rigorous Theoretical Foundation:** The authors provide a detailed theoretical proof for the concept of Selective Prompt Anchoring (SPA), demonstrating how it can approximate the contextual contribution of anchored text to the output logits. This theoretical underpinning adds depth to the methodology and supports the validity of the approach.

**2 Comprehensive Experimental Validation:** The paper backs the theoretical contributions with extensive experimental validation across multiple LLMs of varying sizes. The consistent performance improvements observed on different benchmarks and models showcase the robustness and generalizability of the SPA approach.
summary: The authors identify that LLMs tend to dilute their self-attention on the initial prompt as more code tokens are generated, leading to inaccuracies in the generated code. To address this, they propose a novel approach called Selective Prompt Anchoring (SPA), which amplifies the influence of selected parts of the initial prompt, referred to as "anchored text," during the code generation process.

Key contributions :

**1 Identification of Attention Dilution:** The authors conduct an empirical study revealing that LLMs' attention to the initial prompt diminishes as more code is generated, which they term as "attention dilution."

**2 Proposal of Selective Prompt Anchoring (SPA):** SPA is introduced as a model-agnostic method to optimize LLMs' attention by amplifying the contextual contribution of selective prompt text towards each generated token.

**3** SPA calculates the logit distribution difference with and without the anchored text, approximating the contextual contribution of the anchored text to the output logits. It then creates an augmented logit distribution by linearly combining the original logit distribution and the logit difference.
weaknesses: **1** The authors assume that self-generated tokens by the model are potentially incorrect, which leads to the idea that anchored text should only come from the initial prompt. However, it's possible that some self-generated tokens could also be considered as anchored text, which SPA does not account for currently.

**2** It's not clear how SPA compares to other state-of-the-art methods which focused on anchor-attention[1] improvements . More comprehensive benchmarking against a broader range of methods could strengthen the paper's claims.

**3** The selection of anchored text is crucial for the effectiveness of SPA, yet the paper does not provide a method for identifying the most informative tokens within the prompt. The authors conduct experiments in a general manner, but a more nuanced approach to selecting anchored text could potentially improve results.

**4** The experiments are conducted on HumanEval and MBPP datasets, which may not fully represent the complexity and diversity of real-world programming tasks. Testing SPA on more challenging datasets like OpenEval[2] or BigCodeBench[3] could provide a better understanding of its performance under more demanding conditions.

**5** The paper does not address the generalization of SPA across different programming languages. It's unclear if the same experimental conclusions would hold for languages other than those tested in the paper.

[1] [Anchor-based Large Language Models](https://aclanthology.org/2024.findings-acl.295) (Pang et al., ACL Findings 2024)

[2] [Chain-of-thought in neural code generation: From and for lightweight language models](https://ieeexplore.ieee.org/abstract/document/10634302/) (Yang et al., TSE 2024)

[3] https://bigcode-bench.github.io/

Meta-review: I read the paper, reviewer comments and author rebuttal. This paper was borderline and overall I feel can be improved by another round of reviews. Please consider some suggestions below for future submission:

- The paper currently feels very specific to couple of benchmarks and it is not clear if the results are generalizable. Authors kindly did more experiments with BigBench in the rebuttal period but found much less improvement compared to the two main benchmarks in the paper. This is especially concerning with the additional overhead required in using the approach for each coding task.

- It is not clearly justified why the studied issue of attention dilution is specific to code generation settings. It is also not sure whether this phenomenon is only specific to older class of models.

- The gap compared to existing prompt optimization baselines is limited while the proposed approach requires extra information (full logits) and cannot work with closed source LLMs unlike baselines like ReAct etc.

Decision: Reject