Selective Prompt Anchoring for Code Generation
Authors: Yuan Tian
Tianyi Zhang
Venue: ICML 2025 poster (2025)
Keywords: Large Language Models (LLMs)
Prompt
Code Generation
Self-attention
Taylor expansion
Logits
Anchoring
Primary area: deep_learning->attention_mechanisms
TL;DR: A model-agnostic & training-free approach to adjust the influence of selected 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.

Selective Prompt Anchoring for Code Generation
Yuan Tian Tianyi Zhang
Abstract
this challenge, various prompting methods have been ex- plored to enhance the inference process of code LLMs, e.g., generation (Du et al., 2024), chain-of- thoughts (Le et al., 2024; Suzgun et al., 2022), self celf-planning and debugging (Jiang et al., 2024; Chen et al., 2024), etc.
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 atten- tion dilution issue is an important reason for code generation errors. To mitigate this issue, we pro- pose Selective Prompt Anchoring (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 sults demonstrate that SPA enhances Pass @1 1 by up 12.9%, consistently outperforming SOTA methods in all settings. Our code is available at https://ithub..om/magg Tiaannian/electvee. Prompt-Anchoring.
Despite these efforts, little is known about why LLMs fail to generate correct code. In this work, we seek to bridge the knowledge gap by investigating the attention pattern of LLMs during code generation. We analyzed five code LLMs and found that as more code tokens were generated, LLMs paid less attention to the user prompt. This caused LLMs to gradually deviate from the user intent, thereby leading to code generation errors. Furthermore, we found that more generated tokens led to worse code generation performance, demonstrating their struggle with long-term attention.
To mitigate this limitation, we propose Selective Prompt Anchoring (SPA), model-agnostic approach that optimizes LLMs' attention by amplifying the contextual impact of the user prompt. SPA is inspired by the anchoring effect (Furn- ham & Boo, 2011) 1) psychology, which refers to the phe- nomenon where humans can be influenced by specific infor- mation provided before decision-making. In SPA, we refer to this information anchored text, group of selected prompt tokens that should receive higher attention from the model than others. Figure illustrates the workflow of SPA. Given the anchored text, SPA creates an original embedding matrix ( as well as a masked embedding matrix by re- placing the embeddings corresponding to anchored text with mask embeddings (2). We mathematically show that the anchored text's 's contextual impact can be approximated by the difference between the logit distribution generated from the original prompt and the prompt with the anchored text masked (3). To amplify the impact of anchored text dur- ing code generation, SPA multiplies this logit distribution difference by a hyperparameter called anchoring strength and then adds the original logit distribution (5).
1. Introduction
Large language models (LLMs) have emerged as power- ful programming assistants. They have demonstrated un- precedented capabilities in interpreting natural language descriptions of programming tasks and generating source code (Guo al., 2024; Rozière et al., 2024). Despite this great progress, LLMs still produce incorrect solutions to challenging tasks or generate code that does not fully meet user expectations (Wang et al., 2024; Ning et al., 2023).
To improve the performance of LLMs on coding tasks, many efforts have been made to develop high-quality training data (Li al., 2023; Guo et al., 2024; Wei al., 2023) and design new domain-specific training objectives (Niu al., 2022; Chakraborty et al., 2022). However, these approaches require tremendous computational resources. To address
We evaluate SPA on six benchmarks using six code LLMs. The benchmarks cover different programming languages and task difficulty levels, while the LLMs vary in size and code generation performance. SPA enhances Pass 1 by up to 12.9% across all settings, outperforming four SOTA prompting methods and one SOTA attention steering method. Notably, with SPA, a smaller version of DeepSeek-
Department of Computer Science, Purdue University, West Lafayette, IN, USA. Correspondence to: Yuan Tian <tian211 purduu..ddduu edu>>, Tianyi Zhang <tianyi@ Silany@@purdu.... adrdu..edu>.
Proceedings of the nnnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 20255 by the author(s).

---

Selective Prompt Anchoring for Code Generation
You are an programiinng assistant. Please write python code for me.
Tokenize & Embedding
>>>
Test Cases
count upper(s88cEE")
count count_upperr"aee "popeerabccdegg"
Anchored Text
You
0.3
4.5
0.2
0.1
0.5
0.2
0.8
>>> Task
Given string s, count the number of uppercase vowels even indices.
Append
Code
Prompt
Tast
0.7
0.1
0.8
Oode
0.1
0.11
0.1
0.5
0.3
00.
0.7
[PAD]
[PADDIN]
Original logits
Original embedding matrix
Anchoring Strength
LLM
0.3
4.5
om
0.1
4.2
0
04
0.1
-0.2
Took
0.1
-0.3
.0.
0.5
-0.1
0.8
Augmented logits
(M988)))
[MASK]
Sample
Masked logits
Logits difference
(PAD]
[PADDIN]
Token
Masked embedding matrix
Figure 1. The Workflow of Selective Prompt Anchoring (SPA).
Coder (6.7B) outperform its larger counterpart (33B).
Percentage of attention to user prompt
1.0
CodeGen-350M-Mono
DeepSeek 1ordeeek-ooder-1.3BB 1.3--e--1.---nstuucc
2. Attention Analysis of Code LLMs
0.9
CodeLlama-7B
We conduct an empirical study to investigate the attention dilution phenomenon in code LLMs. Following prior stud- ies (Zhang et al., 2022; Galassi al., 2021). we obtain self-attention scores from the last layer in LLMs,, which has been shown to represent the most accurate attention distribution (Kou et al., 2024; Wan et al., 2022a). We calculate the percentage of attention on the user prompt.
PeepSeek-Coder-33BBe
0.7
0.6
Calculation details are provided in Appendix B.4. We also attention
0.4
experimented with an alternative gradient-bart calculation method (Selvaraju et al., 2016) and obtained similar results, detailed in Appendix B.2.
0.3
150
200
300
350
400
Number of Generated Tokens
We analyzed five code LLMs HumanEval (Chen et al., 2021) and LiveCodeBench (Jain et al., 2024). Figure 2 shows the shift of LLMs' attention on the user prompt dur-
Figure 2. Shift of LLMs' self-atiention the user prompt.
'nntiiiivlly, deeper layers capture representations with long- distance dependencies such the control flow, which mirrors how humans understand programs.
ing code generation. The result shows that as more tokens are generated, models' attention on the user prompt gradu- ally becomes less. Consequently, the code generation pro-
2

---

Selective Prompt Anchoring for Code Generation
creasing the impact of user intent and derive the augmented model output. In Section 3.3, derive and caleulate the augmented logits, which can be approximately represented by the linear combination of original and masked logits.
Table Length of correct vs. incorrect code generated by L.Mss.
EASY
MEDIUM
HARd
OVERALL
PASSED
294
475
400
390
FAILED
418
664
784
622
3.2. Attention Steering and Prompt Anchoring
SPA performs attention steering by scaling the impact of selected tokens to the output logits (E₃)). Suppose is the anchored text for which we want to adjust the impact. E, is an n x m input embedding matrix at step i, and E represents a n submatrix within E, corresponding to the æ. They are visualized below:
cess becomes increasingly influenced by tokens generated recent time steps, rather than the user prompt. This can problematic in two ways. First, any generation inaceuracy, such as creating
unhelpful variable in earlier steps, is likely to propagate and influence subsequent steps. Sec- ond, when generating long and complex code with many detailed requirements, the model may overlook some spec- ifications in the early stages, as the model tends to focus more on recent tokens. To further investigate whether this phenomenon affects performance, we experimented on Live- CodeBench (Jain et al., 2024), which provides difficulty levels for each task. Table 1 shows that, for tasks with the same difficulty level, the average length of incorrectly generated code is consistently longer than that of correctly generated code. This implies that more generation steps can indeed cause attention dilution and hinder accuracy. More details are included the Appendix A.
611
61k
61,++1
61m
621
C2k
02,
C2m
E
ent
enk
6n,6+1
erm
E
(2)
We construct two n x m matrices, X for user prompt and G, for generated code.
To construct X, we only retain i columns of E; that correspond to E", while setting remaining columns to zero. This matrix remains constant during token generation.
3. Approach
3.1. Overview
Given a user prompt v, code LLM fo generates tokens autoregressively. At step i, the input to fo is an n x embedding matrix E,, defined as:
Conversely, G; is formed by zeroing out the same k columns of E; that correspond to E while preserving all elements in the remaining columns.
E, = E" 1E,,,,, 62,.... ......e...........PA PAD]
(1)
They are visualized as follows:
where E is the submatrix of embeddings for tokens in user prompt the series c....., Ci-1 are embeddings of generated tokens, and PAD is a padding submatrix. The model outputs logits and transforms them into a probability distribution. Then, sampling method is applied select the next token.
en
612
f1k
0
0
E21
622
E2k
X
(3)
Cnl
6n2
Cnk
We propose Selective Prompt Anchoring (SPA) steer model attention by amplifying the contextual impact of the important tokens in the user prompt. In this work, we assume these important tokens are selected by users. In- spired by the anchoring effect (Furnham & Boo, 2011) in psychology, we call these important tokens "anchored text". For example, in the prompt, "Given a string, count the num- ber of uppercase vowels". the user may want emphasize "uppercase vowels" to ensure that the model does not for- get to define uppercase vowels (i.e., "A, E, O, U") in the generated code.
0
0
...
61,4+1
1m
0
0
62,x+1
62m
G,
(4)
Cn.4+1
X encapsulates the semantics for the anchored text x, while G; encapsulates the semantics for the remaining text, such as the generated code. The sum of X and G; reconstructs the original matrix E;:
E, X
55)
In Section 3.2, we first mathematically model how to steer model attention. Then we augment model attention by in-
Suppose we want to scale the semantic impact of the matrix
3

---

Selective Prompt Anchoring for Code Generation
X by a value w.? We refer w as anchoring strength. We use (X, w) to represent the function that augments E, by scaling the influence of X by degree of w in the final logit.
Taylor expansion in Equation 8 after the first derivative to obtain an approximation, yielding:
F (e R Fo.s.z (0) + Fo.
(9)
w >> 1 indicates semantic amplification, meaning the model generates code with greater consideration of X. w indicates using the original embedding. E, is equivalent to EX( 1).
where the the integral part Jo dFo....t)) dt in Equation 7 is approximated by w Fo, (0). To calculate (0), 4 we mask the anchored text x using mask embeddings. Each LLM provides special tokens reserved for text masking, which almost has no semantic impact, e.g., <unk> for Code Llama (Rozière et al., 2024) and <pad> for DeepSeek- Coder (Guo et al., 2024). Each special token corresponds to a mask embedding. By replacing embeddings of x with mask embeddings, we get a masked input matrix E, It ablates the semantic impact of the anchored text x while ensuring that the positional encoding remains unaffected. Thus, we can get
de
0 I w < 1 indicates semantic diminishment, meaning the model generates code with less consideration of X. When w is O, the model does not consider X at all. w < 0 indicates a reversed semantic impact, meaning the model generates the code in the opposite manner. For example, if X corresponds to "uppercase", the model will instead consider "lowercase".
In this work, we focus on scaling up the impact of the anchored text x to mitigate the attention dilution issue. Let (w) represent the augmented logits caleulated by model fe at step i, where the impact of anchored text x is scaled by w. We can calculate the integral of the partial derivative of fe with respect to w from O to w.² Formally,
F (0) fo (E.pask).
(10)
To calculate Fe.t. (0), we use get an approximation. Assuming the interval of 1 0 is sufficiently small for Fo.i. we get:
methods to
Fo.s,x (0) ~ Fo Fo,
(11)
Fo = fe(E,(X, w) + G.)
(6)
1- 0
= Fo 1-2 (FFo..0()) dt,
Combining Equations 9, 12 and 13, we get the augmented logits by first-order approximation:
dl
Fo...((w) = Fo. (0) + w. (F), (o)),)(0))
(12)
where t is the variable of integration.
w fo (E, (X, 1) G;)
To reduce computational overhead, attention augmentation is activated when the original generation fails the test case, which is typically included in the code generation context. When no test case is available, attention augmentation is applied continuously throughout the generation. In the fol- lowing section, we explain how to calculate the augmented logits through approximation.
+ (1- fo, 0) + G.)
(13)
w (E.) + (1 - w) fo (Emask))
(14)
Based on the augmented logits Fo, (w) where the impact of the anchored text is scaled by the value w, certain sam- pling algorithm (e.g., greedy sampling) can be applied to select the token. We provide more discussion about approx- imation in Appendix C.
3.3. Augmented Logits by Approximation
We choose not to directly modify self-attention layers to steer the model attention, since this requires identifying which attention head in which layer to steer and the direct editing of attention values does not synchronize with other components like the feedforward layers, which can be fairly brittle and costly. Instead, we chose to simulate attention steering by logits manipulation, which is fast, reliable, and In the next section, we demonstrate that SPA achieves better performance with less computational overhead compared a SOTA method that directly modifies self-attention layers (Zhang al., 2024).
Given the computational complexities of LLMs, directly solving Je d(F....()) dt in Equation 7 is impractical. There- fore, we approximate it by employing the Taylor expansion:
di
w²
Fo.ix(() (w) = (0) w FF..., (0)x
Fos. ((0)
+
2
(8)
Since LLMs are inherently non- non-linear, higher-order deriva- tives of the logits function are non-zero. We truncate the
=Scaling the semantics of X by w is not equivalent to multiply- ing X by w. Multiplying the embedding of X by w does not simply improve the "semantic influence" of X. since the embedding of X also encodes other non- semantie information such as positional information. This is why we compute the difference hetween the logits when masking and unmasking X to cancel out noise and some of the non-semantic information (detailed in Section 3.3). fo is differentiable for hackpropagation.
*Fe... (0) does not mean setting the embedding vector zeros. Instead, it means setting w to zero, which replaces the original em- bedding for anchored text with the mask embedding that contains no semantic information. This mask embedding vector is non-zero. Conversely, zero embedding does not necessarily indicate the absence of semantic information.
4

---

Selective Prompt Anchoring for Code Generation
4. Experiment Setup
Prompting Methods. In addition PASTA, we also com- pare SPA to the mainstream prompting-based code genera- tion methods, including Self-Debugging (Chen et al., 2024), Self-Plamning (Jiang et al., 2024), ReAct (Yao et al., 2023), and Self-Fdit (Zhang et al., 2023b). Self-Debugging and Self-Edit leverage error messages from test cases to refine generated code. Self-planning generates a step-by-step plan before code generation. ReAct prompts an LLM to generate reasoning traces and action plans in an interleaved manner.
4.1. Benchmarks
EvmanEval (Chen et al., 2021) includes 164 Python tasks designed by OpenAl and now has become a widely-used benchmark for code generation.
MBPP (Austin et al., 2021) is another popular benchmark that includes 974 crowd-sourced Python tasks. Due to am- biguous task descriptions, the authors of MBPP created a sanitized version that included 427 tasks with clearer de- scriptions. We evaluate SPA on the sanitized version.
4.3. Evaluation Metries
Following prior work (Chen et al., 2021), we measure code generation performance using the Pass k: metric, which measures whether any of the top k candidates can pass all the test cases, For main results, we report Pass@) 1 using greedy sampling to generate a single code snippet. To demonstrate the generalizability of SPA, we further calculate Pass @ 10 using beam search in Appendix
Vala and MBPP+ (Liu al., 2023) improves the original HvmanEval and MBPP benchmarks with additional test cases to cover corner cases (Liu et al., 2024b).
HulanEval-X (Hendrycks et al., 202 extends the Hu- manEval benchmark to support more programming lan- guages such as Python, Java, JavaScript, C++, and Go. It aims to evaluate the multilingual code generation abilities.
4.4. Anchoring Strength Tuning
Bench (Zhuo et al., 2024) is more challenging benchmark for code generation that evaluates models' abili- ties to follow complex instructions and use tools, including 1,140 real-world Python tasks across 139 libraries.
The anchoring strength w serves as a hyperparameter in SPA. For each model and dataset, we use grid search to tune the anchoring strength w on 1/5 of the tasks, and evaluate Pass@1 1 of SPA on the remaining 4/5 of the tasks. This pro- cess is repeated across all five folds, with final performance metrics averaged across folds, We observe unimodal relationship between w and the performance. We provide more details for the tuning SPA in Appendix J.
LiveCodeBench (Jain et al., 2024) is a contamination free benchmark sourced from competitive programming plat- forms. The benchmark continues to evolve and add new code generation tasks. We conducted experiments latest release, which includes total of 1,055 tasks. The latest tasks span from October 1, 2024, to February 1, 2025.
the
5. Results
4.2. Baselines
5.1. Improvement over Base Models
Base Models. We select six representative open-source code LLMs: ModeGeen-Mooo350M (Nijkamp et al., 2023), CodeLlama-7B (Rozière et al., 2024), StarCoder2- 15B (Lozhkov et al., 2024), and DeepSeek- Coder-IIdtruuet 1.3B, 6.7B, and 33B (Guo et al., 2024). These models exhibit varying levels of performance in code generation. We include more setup details about models in Appendix H.
Table 2 shows SPA consistently improves Pass across all 6 benchmarks and 6 code LLMs. On HumanEval/Huu- manEval+ and MBPP/MBPP+, SPA enhances the base model with on average 5.5% absolute improvement and 14.5% relative improvement, achieving up to a 12.9% abso- lute improvement and a 42% relative improvement. These improvements are observed across models of varying sizes (350M-33B), original performance (15%,,)), and archi- tectures. Notably, with SPA, the smaller DeepSeek-Codee -odder (6.7B) outperforms its much larger 33B counterpart on Hu- manEval. This suggests optimizing model attention is a promising alternative to simply scaling up model size (Ka- plan et al., 2020) for performance improvement. On Big- CodeBench and BivehodeBench, while the absolute im- provements (average 1.57% and 1.90%) are more modest compared to other benchmarks, the relative gains (average 22.83% and 28.50%) remain significant. This is because
Attention Steering Baseline. PASTA (Zhang et al., 2024) is recent method designed to steer model attention for better instruction following. Unlike SPA, PASTA requires a model-specific and time-consuming profiling stage to iden- tify attention headers that are beneficial to performance. For each unique task, PASTA requires around 1000 training samples to identify attention headers that are beneficial to performance. By contrast, SPA is model-agnostic and only needs tuning a single hyperparameter with a few samples, which is fast and generalizable as detailed in Appendix J. PASTA internally edits transformers' self-attention during the feed-forward process, whereas SPA only edits the final logits based on mathematical approximation.
'ee tune SPA all tasks for better generalizability. while SPA is only activated for failed tasks during inference optimize computational efficiency and generation performance.
5

---

Selective Prompt Anchoring for Code Generation
Table Absolute (A) and Relative (1) Performance improvements Pass rates (%)
MODEL
SIZE
IlUMANEVAL
H[UMANEVAL+
MBPP
MBPP+
BIGCODEBENCH
LIVECODEBENCH
CODEGEN-MONE
(350M)
15.3
12.2
19.6
15.9
1.1 A+0.5
1.1 A40.4
+ SPA
20.1
(80% 1)
A+4.8
17.1
(40%
27.4
(40%
A+7.8
22.6
(429
446.7
1.6
(45%
1.5
(36R)
DEEPSEEK-
(1.38)
66.4
61.8
58.2
52.4
2.5
6.5
SPA
70.1
(0% 1)
4+3-7
67.7
(1095 1)
A+55.9
60.9
(6% 1)
A+2.7
52.4
(096 1)
3.4
(30% 1)
A+0.9
9.2
(4216 1)
A+22-
DEEPŠBEK-CORE
(6.7B)
75.6
70.2
67.0
58.5
12.7
7.8
+ SPA
88.5
(179 1)
12.0
79.9
A4.0.7
71.0 (6%
60.7
4.2.2
16.4 (209 1)
10.8
CODELLAMA
(7B)
33.6
28.2
50.9
40.8
3.4
3.8
+ SPA
44.0
10.4
(319 1)
36.0
4-7.
(2515 f)
54.3
4.4.4
(7%% 7)
44.0
4-3.2
(9%6
4.1
(6% 4.0.77
4.0 (3%
4.22
STARCODER2
(16B)
67.7
60.4
78.0
65.1
13.3
7.0
+ SPA
75.6
(1296 1)
A.07.0
65.6
(9% f))
A4552
82.0
(6% 13
4400000
69.1
(0% 1
A.4..0
14.3 (196.1)
A+-0
8.2
(1.756 m)
A4
DEEPŠEEK-CODEE
(33B)
81.7
77.1
73.4
63.2
18.9
11.9
+ SPAA
86.2
(0%6
79.3 (3%6.7)
79.4
A.1
70.3 1)
ATT.1
21.5 (1495
15.8 (33%
SPA leverages the ability of its base model. If the base model could solve a task but overlooks few important tokens in the prompt, SPA can help with this by adjusting the attention. If model lacks the ability to solve task, adjusting the model attention will not help much.
and 34% of the inference time. The performance superior- ity stems from two key factors. On the one hand, unlike prompting methods that add more information enforce generation workflows, SPA preserves the original prompt and does not involve additional LLM calls, thereby achiev- ing better time efficiency. On the other hand, SPA addresses the attention dilution issue that prompting methods do not explicitly handle, thereby achieving higher accuracy.
To demonstrate the generalizability of SPA, we show that SPA can consistently improve @ass@ 10 by, on average 3.42% and up to 7.9%, as detailed in Appendix F. We further evaluate SPA in scenarios where test cases are not available and show that it still significantly enhances perfor- mance, achieving an average 4.7% Pass @ I improvement on HumanEval, as detailed in Appendix K. We illustrate SPA's 's attention anchoring with two concrete examples in Appendix D. SPA improves performance by simply aligning attention to prompts, without introducing additional model parameters or context. We believe this success comes from its ability improve attention reliability. We provide thorough discussion in Appendix G.
Table 3. Comparison between SPA and SOTA methods
METHOD
APAsS@ I (%)
TIME (SEC)
BASE MODEL
0
7.7
PASTA
+1,2
48.8
SELF- DEBUGGING
+4.2
27.3
SELF-EDIT
+1.8
26.4
SELF-PLANNING
43.6
21.6
REACT
+1.3
28.8
SPA
+7.7
9.8
5.2. Comparison to SOTA Methods
PASTA. Table shows that SPA outperforms PASTA by achieving 5.9X higher Pass @1 improvement while using only 20% of the inference time." We include more experi- mental details and discussion in Appendix L. Compared to the base model, SPA increases the decoding time by a prac- tically negligible factor of 1.27. Further discussion about the computational cost of SPA can be found in Appendix E.
5.3. Evaluation on Different Programming Languages
To evaluate the generalizability across different program- ming languages, we further evaluate SPA using HumanEval- X (Hendrycks et al., 2021a), which includes five program- ming languages." 7 Table 4 demonstrates that SPA consis- tently improves Pass @ on HumanEval-X. with an aver- age increase of 7.9% for Python, 4.85% for Java, 6.5% for JavaScript, 3.65% for C++, and 5.2% for Go.
Prompting Methods. Table shows that SPA outperforms Self- Deff-Debugging. Self-Edit, Seif-Planning, and ReAct by achieving improvements that are 1.8X, 4.3X, 2.1X, and 5.9X higher, respectively, while only using 36%, 37%, 45%
5.4. Ablation Study of Anchored Text Selection
To investigate the impact of anchored text selection in code generation tasks, we conduct ablation study by masking
'For fair comparison, we average and add the latency of PASTA's model profiling and the SPA'S tuning the inference time for each task.
IIn the latest version, Rust caused issues when running cases, so we excluded Rust from the results.
6

---

Selective Prompt Anchoring for Code Generation
tire user prompt, performance deteriorates further to 4.76%. This suggests that removing less relevant context can help SPA better align LLM attention with the user intent defined in the natural language description.
Table 4. Evaluation on HumanFval-X with Different Languages.
MODEL
PY
JAVA
JS
C++
Go
CODEGEN (350M)
15.3
9.8
13.4
9.8
6.7
+SPA
+4.8
+3.0
+4.3
+3.7
+6.7
Interestingly, we find that anchoring using human-laheled important tokens achieves only 3.66% on HvmanEval, which is consistently less effective than all the automated experimental anchoring methods. It suggests that narrow- ing the anchored text more informative tokens does not necessarily improve performance. We think there are two plausible reasons. First, since LLMs need to attend to differ- ent context tokens at each decoding step, providing a narrow set of anchored tokens may have a negative impact and dis- tract the L.M in certain decoding steps. Second, previous studies such as Xiao al. (2024) show that even though some tokens, such separators and empty space, may not be semantically meaningful or informative, they provide important signals for LLMs to generate the right content (e.g., following the grammar rules). Thus, over-attending to the informative tokens but not the special tokens in the task description may disrupt the regular generation process. Nevertheless, we think this is challenging but interesting future direction to investigate. We believe our findings will open up new rescarch opportunities for the community on this topic.
DEEPSEEK (I.3B)
66.4
42.7
57.3
43.3
40.2
+SPA
+6.7
+3.0
+3.0
+2.4
+1.8
DEEPSEEK (6.7B)
75.6
48.8
65.2
49.4
45.7
+SPA
+12.9
+8.5
+11.6
+1.8
+7.3
CODELLAMA (7B)
33.6
22.0
29.3
22.0
20.1
+SPA
+10.4
+6.1
+8.5
+6.1
+6.7
STARCODER2 (15B)
67.7
22.0
29.3
22.0
20.1
+SPA
+7.9
+5.5
+7.9
+5.5
+6.1
DEEPSEEK (33B)
81.7
53.0
70.7
53.7
49.4
+SPA
+4.5
+3.0
43.7
+2.4
+2.4
different components in code generation prompts. We de- compose the code generation prompt into three components: (I) natural language description, (2) source code, and (3) test cases. We create automated anchored text selection methods by ablating the source code and test cases. 8 To explore whether anchoring
few more informative to- kens can enhance performance, we further create a condition using important tokens labeled by human programmers as the anchored text. We leverage the dataset from Kou et al. (2024), where multiple programmers manually identified critical tokens that models should attend to when solving HumanEval and MBPP programming tasks. Following our previous experimental setup, we tune SPA for each condition and benchmark, calculating the average Pass @1 1 improve- ment across all 6 experimental L.Mss.
5.5. Impact of Prompt Length
To further validate that SPA effectively addresses the at- tention dilution issue, we analyze both the original model performance and SPA's effectiveness across different prompt lengths, as shown in Table 6. We divide the HumanEval dataset into three equal-sized subsets (Short, Medium, and Long) based on the 33rd and 66th percentiles of prompt lengths. We find that the original code LLMs consistently achieve better performance shorter prompts compared to longer ones. The average Pass@1 is 74.3% for Short prompts, 49.2% for Medium prompts, and only 30.8% for Long prompts. This result echoes our empirical finding on attention dilution, as longer prompts can lead to more severe attention dilution issues.
Table 5. @ I Improvement (%) with Different Anchored Texts (HE: HumanEval, MB: MBPP, BCB: BigCodeBench, LCB: Live- CodeBench).
ANCHORED TEXT
HE
He+
MB
MB+
BCB
LCB
HUMAN ATTENTION (KOU ET AL., 2024)
3.66
3.04
2.58
2.11
N/A
N/A
NATURAL LANGUAGE
5.48
5.08
4.26
3.22
1.57
1.90
We further investigate whether SPA can effectively address the attention dilution issue by comparing performance im- provements across three subsets of tasks, Short, Medium, and Long. To ensure a fair comparison, the attention aug- mentation of SPA is activated for all tasks in each subset of this experiment. The average Pass@1 improvement of SPA is 1.1% for Short, 3.2% for Medium, and 1 11.7% for Long. It shows that SPA consistently achieves greater im- provements for longer prompts compared to shorter ones. This implies that SPA indeed improves performance by ad- dressing the attention dilution issue, as the improvement is more significant when attention dilution is more severe.
+ CODE
4.87
4.65
3.75
2.81
1.42
1.73
+ TEST CASES
5.11
4.89
4.05
3.11
1.50
1.82
CODE & TEST CASES
4.76
4.57
3.98
2.81
1.44
1.73
As shown in Table 5, anchoring the natural language descrip- tion alone achieves the best performance improvement. For example, on HvmanEval, anchoring the natural language description alone improves Pass 1 by 5.48%. Including source code (4.87% and test cases (5. 1..19 11.)) in the anchored text make the performance worse. When anchoring the en-
*We do ablate the natural language description since it explicitly represents the user intent.
7

---

Selective Prompt Anchoring for Code Generation
SPA enhances the base model's performance in all bench- marks except MMLU. While the performance improvement is not as significant as in code generation, we hypothesize that this is because different tasks have unique input and output patterns, leading to varying degrees of attention di- lution. For MMLU, the model only needs to generate a choice to answer the multiple-choice question. The output length is significantly shorter than in code generation and considerably less than in user prompts. Consequently, the model's attention is hardly diluted by self-generated tokens, making SPA not helpful. In contrast, TruthfulQA requires the model to generate a text analysis, provide reasoning, and answer the question. Therefore, SPA is more beneficial in addressing attention dilution and correcting the remaining 27% errors for "TRUTH." " (Truthfulness) and 41% errors for "INFO." " (Informativeness). Nevertheless, it remains an interesting future work to investigate how to further improve SPA for a broader range of generative tasks beyond code generation. More details are reported in Appendix I.
Table 6. SPA'S performance on tasks with different prompt lengths.
MODEL
SHoRt
MEDIUM LONG
CODEGEN (350M)
37.0
6.6
2.3
SPA
+1.7
+3.0
+4.3
DEEPSEEK- COEBER (I.3B)
81.8
45.5
29.6
SPA
-1.8
+14.5
+26.0
DREPSEEK-CODER (6.78)
87.3
69.1
44.4
SPA
+3.6
-3.6
+7.5
CODELLAMA (7B)
69.2
43.5
0
+ SPA
+2.6
+0.0
+10.0
STARCODER2 (15B)
85.5
49.1
40.1
SPA
-1.8
+5.5
+14.8
DEEPSEEK- CODER (33B)
85.5
81.8
68.5
+ SPA
+1.8
+0.0
+7.4
5.6. Impact of Anchoring Strength
In this section, we discuss how the anchoring strength w of SPA influences performance. We observe that different values of w yield different effects: when w is too low, per- formance gains are limited; when it too high, the LLM becomes biased, leading to performance degradation. As shown in Figure 3, = corresponds to the original gen- eration. As w increases, performance initially improves, reaches an optimum, and subsequently declines with further increases in w. This simple pattern applies to all the settings. We find w is both tunable and transferable. While the opti- mal value may vary slightly across models and benchmarks, it is generally model-dependent. A value of consistently yields robust improvements across all settings. Further details are provided in Appendix J.
6. Related Work
Code Generation. To enhance the performance of LLMs coding tasks, significant efforts have been dedicated to curating high-quality training data (Li et al., 2023; Guo et al., 2024; Tian et al., 2025; Wei et al., 2023) and design- ing domain-specific training objectives (Niu et al., 2022; Chakraborty et al., 2022). Furthermore, techniques such as instruction tuning (Wei et al., 2022), reinforcement learning with human feedback (Ouyang et al., 2022), and repository- level context modeling (Zhang et al., 2023a) have been explored to improve alignment, reasoning, and context un- derstanding abilities in code generation. However, these approaches often require significant fine-tuning effort. In the meantime, a line of work has focused on developing interactive methods (Di & Zhang, 2025; Tian et al., 2023; 2024) that incorporate real-time human feedback to guide and refine the model's 's output. Despite their effectiveness, such human-ine -n-th--oo- approaches remain limited due to the requirement of user intervention. To overcome these challenges, recent research has increasingly explored au- tomated and training-free free prompting methods (Chen et al., 2024; Zhang et al., 2023b; Suzgun et al., 2022: Le et al., 2024), which aim to automatically optimize prompts by integrating additional reasoning or contextual information. For example, Self- teff-Deuugging (Chen et al., 2024) enable I.IMS to debug code based on error messages and execu- tion results. Self-Planning (Jiang et al., 2024) allows LLMs to decompose tasks into subtasks and implement solutions step-by-step. ReAct (Yao et al., 2023) prompts an LLM to generate reasoning traces and action plans in an interleaved manner. Compared with these methods, SPA takes an or- thogonal approach that amplifies the influence of the user prompt to mitigate the attention dilution issue.
1.25
To make SPA more stable and eliminate the need of tuning, weighting strategy
we also propose a that adjusts each token's 's anchoring strength based on its original probability. Details are provided in Appendix M.
5.7. Evaluation on Other Generative Tasks
Table 7. SPA'S performance on different generative tasks ("TRUTH." and "INFo." are evaluation metries of TruthfulQA).
MODEL
TRUTH. INFO. GSMSK MMLU BooLQ
LLAMA-3.1
88.4
97.8
76.5
69.5
83.0
+ SPA
+3.1
+0.9
+1.1
+0.0
+0.1
While we focus on code generation in this work, we are interested in whether SPA can be applied to other gen- erative tasks, Thus, we evaluate SPA on other genera- tive benchmarks, including TruthfulQA (Lin et al., 2022), GSM8K (Cobbe et al., 2021), MMLU (Hendrycks et al., 2021b), and BoolQ (Clark et al., 2019). As shown in Table 7,
8

---

Selective Prompt Anchoring for Code Generation
HumanEvell
HumanEval-
NBPP
MBPP*
Cadootie--aana (350M)
DeepSeek Coder lndde--isuuutttt (1.35)
DeepSeck-Codeee Contr-tttut (6. (8.7B)
CodeLiama (7B)
DespSook- Coder-nstrdd- (33B)
33
0.9
0.
Pass@ln Rette
0.6
0.5
0.0
0.2
0.8
0.1
0.4
0.6
0.3
0.2
1.00 1.25
1.50
1.75
1.50
1.25
1.90
1.55
2.00
1.25
2.00
1.25
1.75
Anchoring Strength w
Figure 3. Analysis of Anchoring Strength
Attention Steering. TOAST (Shi et al., 2023a) tunes fea- ture selection module to redirect attention to task-relevant features. PASTA (Zhang et al., 2024) performs model pro- filing to identify beneficial attention headers and recalculate attention distributions across transformer layers. However, these attention-stecring approaches usually require exten- sive model adaptations and complex setup procedures. In contrast, SPA provides model-agnostic solution that math- ematically simulates attention steering via logit arithmetic.
anchored text and anchoring strength. For the selection of anchored text, one idea is to use LLMs to dynamically iden- tify relevant words or phrases corresponding to the current generation step. For tasks where NL may not be important compared to code (e.g., code translation), we can use static code analysis identify important code elements (c.g., func- tion calls and variable names heavily used in the code). For the selection of anchoring strength, one idea is to develop method to calculate the relevance of words and phrases to each generation step. Based on the relevance scores, the system can assign higher values to more relevant contexts while assigning lower values to less relevant ones.
Logit Arithmetic. There has been a growing body of re- search on performing arithmetic transformations on logits to enhance text generation, such as contrasting logits from mul- tiple LMs (Liu et al., 2024a; 2021; Dou et al., 2019; Zhao et al., 2024) and contrasting logits from different layers of model (Chuang et al., 2024; Gera et al., 2023). Unlike these methods, SPA contrasts logits from the same model by perturbing the input through masking, rather than providing additional context (Pei et al., 2023; Shi et al., 2023b; Malkin al., 2022) or changing to a completely new prompt (Sen- nrich et al., 2024). Furthermore, we provide mathematical approximation of semantic scaling arbitrary groups of embeddings. SPA is specifically designed to address the attention dilution issue in LLMs during code generation- phenomenon first observed in our work. By contrast, none of the existing works explored code generation or model attention. They primarily focus on enhancing coherence (Malkin et al., 2022), factuality (Shi al., 2023b; Chuang al., 2024; Sennrich et al., 2024; Leng et al., 2023), and controllability (Liu et al., 2021; Pei et al., 2023; Zhao et al., 2024).
In our current experiments, utilize existing test cases in the benchmark to determine whether to activate SPA. Al- though this is common practice in code generation (Chen et al., 2024; Zhang et al., 2023b), this may not be applicable to programming tasks where test cases are not available. This can be potentially addressed by prompting LLMs to generate test cases. Furthermore, we believe SPA will be beneficial in a self-improving pipeline for fairly complex tasks. By analyzing the errors in the initially generated code, LLMs can be prompted identify which instructions or requirements were not followed in the prompt. Then SPA can be used to amplify the influence of these ignored instructions.
8. Conclusion
We present SPA, model-agnostic approach to enhancing LLM code generation. Our study first identifies the atten- tion dilution phenomenon, where code LLMs increasingly overlook the prompt as generation progresses. To address this, SPA introduces a training- free, mathematically proven mechanism for controlling the influence of selected prompt tokens. Experiments demonstrate that aligning model atten- tion to user prompts using SPA significantly and consistently enhances the code generation performance of base LLMs.
7. Limitations & Future Work
In this work, we pre-define the method for selecting an- chored tokens and use a fixed anchoring strength when generating code. We consider this approach a baseline. Future work could explore how to dynamically select the

---

Selective Prompt Anchoring for Code Generation
Acknowledgements
Clark, C.,, Lee, K., Chang, M.-W., Kwiatkowski, T., Collins,
M., and Toutanova, K. Boolq: Exploring the surprising difficulty of natural yes/no questions, 2019.
We thank all the anonymous reviewers for their valuable and detailed feedback, which has significantly improved the quality of this paper. This work was supported in part by the National Science Foundation (NSF CAREER Grant 2340408).
Cobbe, K., Kosaraju, V., Bavarian, M., Chen, M., Jun, H.,
Kaiser, L., Plappert, M., Tworek, J., Hilton, J., Nakano, R., Hesse, C., and Schulman, J. Training verifiers to solve math word problems, 2021.
Impact Statement
Dettmers, T., Lewis, M., Belkada, Y., and Zettlemoyer, Llm. int((: 8-bit matrix multiplication for transformers at scale, 2022.
This paper presents work whose goal to advance the field of Machine Learning. There are many potential societal consequences of our work, none of which feel must be specifically highlighted here.
Di, Y. and Zhang, T. Enhancing code generation via bidi- rectional comment-level mutual grounding, 2025.
Dou, Z.-Y., Wang. X., Hu, J., and Neubig, G. Domain differential adaptation for neural machine translation. In Birch, A., Finch, A., Hayashi, H., Konstas, I., Luong, T., Neubig, G., Oda, Y., and Sudoh, K. (eds.), Proceedings of the 3rd Workshop on Neural Generation and Translation, pp. 59-69, Hong Kong, November 2019. Association for Computational Linguistics. doi:
References
Austin, J., Odena, A., Nye, M., Bosma, M., Michalewski, H., Dohan, D., Jiang, E., Cai, C., Terry, M., Le, Q., and Sutton, C. Program synthesis with large language models, 2021.
Chakraborty, S., Ahmed, T., Ding, Y., Devanbu, P. T., and Ray, B. Natgen: generative pre-training by "naturalizing" source code. In Proceedings of the 30th ACM joint eu- ropean software engineering conference and symposium on the foundations of software engineering, pp. 18-30, 2022.
Du, K., Rui, R., Chai, H., Fu, L., Xia, W., Wang, Y., Tang, R., Yu, Y., and Zhang, W. Codegrag: Extracting com- posed syntax graphs for retrieval augmented cross-lingual code generation, 2024.
Frantar, E., Ashkboos, S., Hoefler, T., and Alistarh, D. Gptq: Accurate post-training quantization for generative pre- trained transformers, 2023.
Chen, M., Tworek, J., Jun, H., Yuan, Q., de Oliveira Pinto, H. P., Kaplan, J., Edwards, H., Burda, Y., Joseph, N., Brockman, G., Ray, A., Puri, R., Krueger, G., Petrov, M., Khlaaf, H., Sastry, G., Míshkin, P., Chan, B., Gray, S., Ryder, N., Pavlov, M., Power, A., Kaiser, L., Bavar- ian, M., Winter, C., Tillet, P., Such, F. P., Cummings, D., Plappert, M., Chantzis, F., Barnes, E., Herbert- Voss, A., Guss, W. H., Nichol, A., Paino, A., Tezak, N., Tang, J. Babuschkin, I., Balaji, S., Jain, S., Saunders, W., Hesse, C., Carr, N., Leike, J., Achiam, J., Misra, V., Morikawa, E., Radford, A., Knight, M., Brundage, M., Murati, M., Mayer, K., Welinder, P., McGrew, B., Amodei, D., McCandlish, S., Sutskever, I., and Zaremba, W. Evaluating large language models trained on code. CoRR, abs/2107.00 .03374, 2021.
Furnham, A. and Boo, H. C. A literature review of the anchoring effect. The Journal of Docio-Economies 40(1):
35-42, 2011. ISSN 1053-5357. doi: https://doi.org/10.. (016/i.00ce..200..10..
Galassi, A., Lippi, M., and Torroni, P. Attention in natu- ral language processing. IEEE Transactions on Neural Networks and Learning Systems, 32(10):4291-4308 Oc- tober 2021. ISSN 2162-2388. doi: 10.1 1020 3019893.
Gera, A., Friedman, R., Arviv, O., Gunasekara, C., Szna- jder, B., Slonim, N., and Shnarch, E. The benefits of bad advice: Autocontrastive decoding across model lay- ers. In Rogers, A., Boyd-Graber, J., and Okazaki, N. (eds.), Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics ( (Volume I: Long Papers), pp. 10406-10420, Toronto, Canada, July 2023. Association for Computational Linguistics. doi: 10. 10.18863//1/20233aclle lon..000..0.
Chen, X., Lin, M., Schärli, N., and Zhou, D. Teaching large language models self-debug. In The Twelfth International Conference on Learning Representations, 2024.
Chiang, D. and Cholak, P. Overcoming theoretical limita- tion of self-attention 2022.
Guo, D., Zhu, Q. Yang. D., Xie, Z., Dong, K., Zhang, W., Chen, G., Bi, X., Wu, Y., Li, Y. K., Luo, F., Xiong, Y., and Liang, W. Deepseek-coder. When the large language model meets programming - the rise of code intelligence, 2024.
Chuang, Y.S.,, Xie, Y., Luo, H., Kim, Y., Glass, J., and He, P. Dola: Decoding by contrasting layers improves factuality in large language models, 2024.
10

---

Selective Prompt Anchoring for Code Generation
Hendrycks, D., Basart, S., Kadavath, S., Mazeika, M., Arora, A., Guo, E., Burns, C., Puranik, S., He, H., Song, D., and Steinhardt, J. Measuring coding challenge competence with apps, 2021a.
controlled text generation with experts and anti-experts. 2021.
Liu, A., Han, X., Wang, Y., Tsvetkov, Y., Choi, Y., and Smith, N. A. Tuning language models by proxy, 2024a.
Hendrycks, D., Burns, C., Basart, S., Zou, A., Mazeika, M., Song, D., and Steinhardt, J. Measuring massive multitask language understanding, 2021b.
Liu, J., Xia, C. S., Wang, Y., and Zhang, L. Is your code generated by chatgpt really correct? rigorous evaluation of large language models for code generation, 2023.
Huang, W., Ma, X., Qin, H., Zheng, X., Lvv, C., Chen, H., Luo, J., Qi, X., Liu, X., and Magno, M. How good are low-bit quantized Ilama3 models? an empirical study, 2024.
Liu, Y., Tantihaamthavom, C., Liu, Y., and Li, L. On the reli- ability and explainability of language models for program generation, 2024b.
Lozhkov, A., Li, R., Allal, L. B., Cassano, F., Lamy-Poiriee J., Tazi, N., Tang, A., Pykhtar, D., Liu, J., Wei, Y., Liu, T., Tian, M., Kocetkov, D.,, Zucker, A., Belkada, Y., Wang, Z,, Liu, Q., Abulkhanov, D., Paul, I., Li, Z.,, Li, W.-D., Risdal, M. Li, J., Zhu, J., Zhuo, T. Y., Zheltonozhskii, E., Dade, N. O. O., Yu, W., Krauß, L.., Jain, N., Su, Y., He, X., Dey, M., Abati, E., Chai, Y., Muennighoff, N., Tang, X., Oblokulov, M., Akiki, C., Marone, M., Mou, C., Mishra, M., Gu, A., Hui, B., Dao, T., Zebaze, A.. Dehaene, O., Patry, N., Xu, C., McAuley, J., Hu, H., Scholak, T., Paquet, S., Robinson, J., Anderson, C. J., Chapados, N., Patwary, M., Tajbakhsh, N., Jernite, Y., Ferrandis, C. M., Zhang, L., Hughes, S., Wolf, T., Guha, A., von Werra, L., and de Vries, H. Starcoder 2 and the stack v2: The next generation, 2024.
Jain, N., Han, K.. Gu. A., Li. W.-D., Yan, F. Zhang, T., Wang, S., Solar-Lezama, A., Sen, K., and Stoica, I. Live- codebench: Holistic and contamination free evaluation of large language models for code, 2024.
Jiang, X., Dong, Y., Wang, I.., Fang, Z., Shang, Q. Li, G., Jin, Z., and Jiao, W. Self-planning code generation with large language models. ACM Transactions on Software Engineering and Methodology, 33(7):1-30, 2024.
Kaplan, J., McCandlish, S., Henighan, T., Brown, T. B., Chess, B., Child, R., Gray, S., Radford, A., Wu, J., and Amodei, D. Scaling laws for neural language models, 2020.
Kou, B., Chen, S., Wang, Z., Ma, L., and Zhang, T. Do large language models pay similar attention like human programmers when generating code? Proc. ACM Softw. Eng., I(FSE), July 2024. doi: 10/3667/360000.
Malkin, N., Wang. Z., and Jojic, N. Coherence boost- ing: When your pretrained language model is not paying enough attention. In Muresan, S., Nakov, P., and Villavi- cencio, A. (eds.), Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume I: Long Papers). pp. 2214-8236, Dublin, Ireland, May 2022. Association for Computational Linguistics. doi: 10. 10.18665/11122222.cllooogg
Le, H., Chen, H., Saha, A., Gokul, A., Sahoo, D., and Joty, S. Codechain: Towards modular code generation through chain of self-revisions with representative sub-modules, 2024.
Nijkamp, E.,, Pang. B., Hayashi, H., Tu, L., Wang. H., Zhou, Y., Savarese, S., and Xiong, C. Codegen: An open large language model for code with multi-turn program synthe- sis, 2023.
Leng, S., Zhang, H., Chen, G., Li, X., Lu, S., Miao, C., and Bing, L. Mitigating object hallucinations in large vision- language models through visual contrastive decoding, 2023.
Ning. 7.,, Zhang, Z., Sun, T., Tian, Y., Zhang, T., and Li, T. J.-J. An empirical study of model errors and user error discovery and repair strategies in natural language database queries. In Proceedings of the 28th International Conference on Intelligent User Interfaces, IUI '23, pp. 633-649, New York, NY, USA, 2023. Association for Computing Machinery. ISBN 9798400701061. doi: 10. 1145/358 1113588661135840677
Li, R., Zi, Y., Muennighoff, N., Kocetkov, D., Mou, C., Marone, M., Akiki, C., Jia, L., Chim, J., Liu, Q. et al. Starcoder: may the source with you! Transactions on Machine Learning Research, 2023.
Li, S., Ning, X., Wang, L.. Liu, T., Shi, X., Yan, S., Dai, G., Yang, H., and Wang, Y. Evaluating quantized large language models, 2024.
Lin, S., Hilton, J., and Evans, O. Truthfulqa: Measuring how models mimic human falsehoods, 2022.
Ning. Z., Tian, Y., Zhang, Z.,, Zhang, T., and Li, T. J.- J. Insights into natural language database query errors: From attention misalignment to user handling strategies. ACM Trans. Interact. Intell. Syst., mar 2024. ISSN 1660 6455. doi: 10. 145/36501 14. Just Accepted.
Liu, A., Sap. M., Lu, X., Swayamdipta, S., Bhagavatula, C., Smith, N. A., and Choi, Y. Dexperts: Decoding-time

---

Selective Prompt Anchoring for Code Generation
Niu, C., Li, C., Ng, V., Ge, J., Huang, L., and Luo, B. Spt-code:
Shrikumar, A., Greenside, P., and Kundaje, A. Learning important features through propagating activation differ- ences. CoRR, abs/1704.02685, 2017.
pre-training for learning source code representations. In Proceedings of the 44th international conference on software engineering, pp. 2006-2018, 2022.
Suzgun, M., Scales, N., Schärli, N., Gehrmann, S., Tay, Y., Chung, H. W., Chowdhery, A., Le, Q. V., Chi, E. H., Zhou, D., and Wei, J. Challenging big-bench tasks and whether chain-of tf-hhuughtt can solve them, 2022.
Ouyang, L., Wu, J., Jiang, X., Almeida, D., Wainwright, C. L., Mishkin, P., Zhang, C., Agarwal, S., Slama, K., Ray, A., Schulman, J., Hilton, J., Kelton, F., Miller, I.., Simens, M., Askell, A., Welinder, P., Christiano, P., Leike, J., and Lowe, R. Training language models to follow instructions with human feedback, 2022.
Tian, Y., Zhang, 7.,, Ning, Z.,, Li, T. J.J., Kummerfeld, J. K., and Zhang, T. Interactive text-to-SQL genera- tion via editable step-by-step explanations. In Bouamor, H., Pino, J., and Bali, K. (eds.). Proceedings of the 2023 Conference on Empirical Methods in Natural Lan- guage Processing, pp. 16149-16166, Singapore, Decem- ber 2023. Association for Computational Linguistics. doi: 10.1 10.18653/v1/2023.emm www.:/222.mmmmmpmmmmee main..004
Pei, J., Yang, K., and Klein, D. PREADD: Prefix- adaptive decoding for controlled text generation. In Rogers, A., Boyd-Graber, J., and Okazaki, N. (eds.), Findings of the Association for Computational Linguistics: ACL 2023, pp. 10018-10031 Toronto, Canada, July 2023. Association for Computational Linguistics. doi: 10. 18653/11/22233. findings-acl.636.
Tian, Y., Kummerfeld, J. K., Li, J.JJ,, and Zhang, T. Sqlu- cid: Grounding natural language database queries with interactive explanations. In Proceedings of the 37th An- nual ACM Symposium User Interface Software and Technology, UIST 224, New York, NY, USA, 2024. Asso- ciation for Computing Machinery. ISBN 9798400706288. doi: 10. 00.114//654777..66663
Raganato, A. and Tiedemann, J. An analysis of encoder representations in transformer-based machine transla- tion. In Linzen, T., Chrupala, G., and Alishahi, A. (eds.), Proceedings of the 2018 EMNLP Workshop Black- boxNLP: Analyzing and Interpreting Neural Networks for NLP, pp. 287-297, Brussels, Belgium, November 2018. Association for Computational Linguistics. doi: 10.
Tian, Y. Lee, D., Wu, F., Mai, T., Qian, K., Sahai, S., Zhang, T., and Li, Y. Text-to-sql domain adaptation via human-llm collaborative data annotation. In Proceedings of the 30th International Conference on Intelligent User Interfaces, pp. 1398-1425, ACM, March 2025. doi: 10. 1145/370835..771009
Rozière, B., Gehring, J., Gloeckle, F., Sootla, S., Gat, I., Tan, X. E., Adi, Y., Liu, J., Sauvestre, R., Remez, T., Rapin, J., Kozhevnikov, A., Evtimov, I., Bitton, J., Bhatt, M., Ferrer, C. C., Grattafiori, A., Xiong, W., Défossez, A., Copet, J., Azhar, F., Touvron, H., Martin, L., Usunier, N., Scialom, T., and Synnaeve, G. Code Ilama: Open foundation models for code, 2024.
Touvron, H., Lavril, T., Izacard, G., Martinet, X., Lachaux, M.-A., Lacroix, T., Rozière, B., Goyal, N., Hambro, E., Azhar, F., Rodriguez, A., Joulin, A., Grave, E., and Lam- ple, G. Llama: Open and efficient foundation language models, 2023a.
Selvaraju, R. R., Das, A., Vedantam, R., Cogswell, M., Parikh, D., and Batra, D. Grad-cam: Why did you say that? visual explanations from deep networks via gradient-based localization. CoRR, abs/1610.02391, 2016.
Touvron, H., Martin, L., Stone, K., Albert, P., Almahairi, A., Babaei, Y.. Bashlykov, N., Batra, S., Bhargava, P., Bhosale, S., Bikel, D.,, Blecher, I.., Ferrer, C. C.,, Chen, M., Cucurull, G., Esiobu, D., Fernandes, J., Fu, J., Fu, W., Fuller, B., Gao, C., Goswami, V., Goyal, N., Hartshorn, A., Hosseini, S., Hou, R.. Inan, H., Kardas, M., Kerkez, V., Khabsa, M., Kloumann, I., Korenev, A., Koura, P. S., Lachaux, M.-A., Lavril, T., Lee, J., Liskovich, D., Lu, Y., Mao, Y. Martinet. X., Mihaylov, T., Mishra, P., Molybog. I. Nie. Y., Poulton, A., Reizenstein, J., Rungta, R., Saladi, K., Schelten, A., Silva, R., Smith, E. M., Subramanian, R., Tan, X. F., Tang. B., Taylor, R., Williams, A., Kuan, J. X., Xu, P., Yan, Z., Zarov, I., Zhang, Y., Fan, A., Kambadur, M., Narang, S., Rodriguez, A., Stojnic, R., Edunov, S., and Scialom, T. Llama 2: Open foundation and fine-tuned chat models, 2023b.
Sennrich, R., Vamvas, J., and Mohammadshahi, Miti- gating hallucinations and off-target machine translation with nurce-contrrative and language-contrastive decod- ing, 2024.
Shi, B., Gai, S., Darrell, T., and Wang, X. Toast: Transfer learning via attention steering, 2023a.
Shi, W., Han, X., Lewis, M., Tsvetkov, Y., Zettlemoyer, L., and tau Yih, S. W. Trusting your evidence: Hallucinate less with aaate decoding, 2023b.
12

---

Selective Prompt Anchoring for Code Generation
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L.,, Gomez, A. N., Kaiser, I.., and Polosukhin, I. Attention is all you need. CoRR, ttss/170600366. 2017.
Zhang, K., Li, Z.,, Li, J., Li, G., and Jin, Z. Self-edit: Fault- aware code editor for code generation. In Rogers, A., Boyd-Graber, J., and Okazaki, N. (eds.), Proceedings of the 61st Anmual Meeting of the Association for Computa- tional Linguistics (Volume Long Papers), pp. 769-787, Toronto, Canada, July 2023b. Association for Computa- tional Linguistics. doi: 10.1 1115533v1/022.acl--2ao
Voita, E., Talbot, D., Moiseev, F., Sennrich, R., and Titov, I. Analyzing multi-head self-attention: Specialized heads do the heavy lifting, the rest can be pruned, 2019.
Wan, Y., Zhao, W., Zhang, H., Sui, Y., Xu, G., and Jin, H. What do they capture? a structural analysis of pre- trained language models for source code. In Proceedings of the 44th International Conference on Software Engi- neering, ICSE '22, pp. 2377-2388, New York, NY, USA, 2022a. Association for Computing Machinery. ISBN 9781450392211. doi: 10.1 10.1145/351000033515555
Zhang, Q., Singh, C., Liu, L., Liu, X., Yu, B., Gao, J., and Zhao, T. Tell your model where to attend: Post-hoc attention steering for 1lms, 2024.
Zhao, X., Yang, X., Pang. T., Du, C.,, Li. Wang, Y.-X., and Wang, W. Y. Weak-to-strong jailbreaking on large language models, 2024.
Zhuo, T. Y., Vu, M. C.,, Chim, J., Hu, H., Yu, W., Widyasari, R., Yusuf, I. N. B., Zhan, H., He, J., Paul, I., Brunner, S., Gong, C., Hoang, T., Zebaze, A. R., Hong, X., Li, W.-D., Kaddour, J., Xu, M., Zhang, Z., Yadav, P., Jain, N., Gu, A., Cheng, Z., Liu, J., Liu, Q., Wang, Z., Lo, D., Hui, B., Muennighoff, N., Fried, D., Du, X., de Vries, H., and Werra, V. Bigcodebench: Benchmarking code genera- tion with diverse function calls and complex instructions. 2024.
Wan, Y., Zhao, W., Zhang, H., Sui, Y., Xu, G., and Jin, H. What do they capture? a structural analysis of pre-trained language models for source code, 2022b.
Wang, B. and Komatsuzaki, A. GPT-J-6B: A billion parameter autoregressive language model, 2021.
Wang, Z., Zhou, Z., Song, D., Huang, Y., Chen, S., Ma, L., and Zhang, T. Where do large language models fail when generating code?, 2024.
Wei, J., Bosma, M., Zhao, V. Y., Guu, K., Yu, A. W., Lester, B., Du, N., Dai, A. M., and Le, Q. V. Finetuned language models are zero-shot learners, 2022.
Wei, Y., Wang, Z., Liu, J., Ding, Y., and Zhang, I.. Magi- coder: Source code is all you need. arXiv preprint arXiv:2312.02120, 2023.
Xiao, G., Tian, Y., Chen, B., Han, S., and Lewis, M. Ef- ficient streaming language models with attention sinks, 2024.
Yao, S., Zhao, J., Yu, D., Du, N., Shafran, I., Narasimhan, K. R., and Cao, Y. React: Synergizing reasoning and acting in language models. In The Eleventh International Conference on Learning Representations, 2023.
Zhang, F., Chen, B., Zhang, Y., Keung, J., Liu, J., Zan, D., Mao, Y., Lou, J.-G., and Chen, W. RepoCoder: Repository-level code completion through iterative re- trieval and generation. In Bouamor, H., Pino, J., and Bali, K. (eds.), Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, 2471-2484, Singapore, December 2023a. Association for Computational Linguistics. doi: 10. 10118833/1/223. emnlp-main main. 151.
Zhang, K., Li, G., and Jin, What does transformer learn about source code?, 2022.
13

---

Selective Prompt Anchoring for Code Generation
Table of Contents
Performance Analysis by Different Generation Lengths
14
Calculation Details and Extended Discussion of LLM Attention
15
B.1
Calculation of Self-attention
15
B.2
Calculation of caadient-based Attention
15
B.3
Attention Misalignment
15
B.4
Attention to the User Prompt Ratio
15
Extended Discussion of Approximation in SPA
16
D Code Generation Examples
17
E
Computational Cost of SPA
18
E.1 Latency
18
E.2 Memory
18
E.3 Tradcoff Between Latency and Memory Usage via Parallelization
18
F @ass@ 10 & Beam Search with SPA
18
Hypothetical explanation for Attention dilution and SPA's effectiveness
19
Implementation and Deployment
20
H.I Imlementation of SPA
20
H.2 Model Deployment
20
H.3 Prompt Design
20
I Evaluating SPA on Other Generative Tasks
20
J Tuning of Anchoring Strength
J.1 Tuning Stability
21
J.2 Cross-dataset & Cross-Mode Tuning Evaluation
21
J.3 Comparision between Tuned SPA and Optimal SPA
21
Analysis of SPA's 's Performance without Test Case Availability
23
L Comparison PASTA
23
M Confidence-Mddulated Anchoring Strength
24
A. Performance Analysis by Different Generation Lengths
2033
324.4
81.85
Number of Characters
Number of Words
Number of Lines
1323
211.2
53.62
613
97.8
25.38
Failed
Passed
Failed
Passed
Failed
Passed
Figure 4. Length of correct vs. incorrect code generated by LLMs at three levels of granularity: character, word, and line.
To evaluate how code length affects model performance, we analyze the length distribution of correctly and incorrectly generated code. We use the same set of experimental LLMs (Nijkamp et al., 2023; Guo et al., 2024; Rozière et al., 2024) as in Section 2 and experiment on both HumanEval (Chen et al., 2021) and LiveCodeBench (Jain et al., 2024).
Particularly, we evaluate code length three granularities: characters, words (tokens), and lines. Figure demonstrates
14

---

Selective Prompt Anchoring for Code Generation
that incorrectly generated code is significantly longer than correctly generated code across different granularity levels. To mitigate the impact of the correlation between task difficulty and generation length, conduct further analysis at different task difficulty levels in LiveCodeBench. Table I shows that the average length of incorrectly generated code is consistently longer than that of correctly generated code, across all levels.
These results echo our findings about attention dilution-as as the generated sequence grows longer, the model's 's attention to the user prompt diminishes, leading to more errors.
B. Calculation Details and Extended Discussion of LLM Attention
B.1. Calculation of Self-attention
Most LLMs are based on the decoder of transformer (Vaswani et al., 2017), which has multiple self-attention layers. Roughly speaking, given an LLM fo and an input sequence of tokens to, ti,..., fn where & represents the ith token. The transformer calculates relevance scores between every pair of tokens. The self-attention score for a token 4 in the sequence can be roughly formulated as:
Σn= relevance(ts,tt) أ== relevance(!;, (g)
attention(4.)
(15)
where the relevance function approximates the computation among Q K, V in transformers (Vaswani et al., 2017). However, different layers have different attention distributions. According to a study (Wan et al., 2022b), deeper self-attention layers can better capture long-distance dependencies and program structure, so we calculate the attention by aggregating attention from multiple heads the last layer. Nevertheless, this still excludes the influence from the last forward layer.
B.2. Calculation of Gradient- based Attention
To validate the generalizability of attention dilution, we employ a gradient- attention calculation method (Selvaraju al., 2016; Shrikumar al., 2017). Compared using self attention layers in transformers, the gradient-based method can be generalized to different model architectures by treating the entire model as whole differentiable function. It computes the model's attention by calculating the gradients relative to each input token. Intuitively, token that induces larger gradient is considered more influential, suggesting that the model pays greater attention to it. Formally, the attention over the token f; is calculated by
attention((.) Ofo(to, .t.....tt.) tn) Ot,
(16)
As shown in Figure 5, we observe a similar declining pattern in the model's attention over the initial prompt, suggesting that attention dilution is a fundamental phenomenon that persists across different attention mcasurement approaches.
B.3. Attention Misalignment
Despite the success of the attention mechanism in LLMs, prior works found that language models often exhibit simple attention patterns (Raganato & Tiedemann, 2018; Voita et al., 2019). Furthermore, empirical study (Kou et al., 2024) found that given coding task, there often exists a misalignment between LLM attention and human attention. When generating code, LLMs often focus on parts in natural language deseriptions that are different from what human programmers focus on. Another work (Ning et al., 2024) shows that when the model's attention aligns more closely with human programmers' attention, the model generates more accurate SQL queries. Inspired by these findings, we hypothesize that a root cause of inaccuracy in --g-eerertted code stems from the suboptimal model attention.
B.4. Attention to the User Prompt Ratio
Based on these two methods to calculate LLMs' attention, we analyze how the attention of LLMs to the initial prompt shifts. Formally, given the prompt & and the following generated tokens tg, 1,..., f,-1, we calculate the attention to the user prompt ratio (x) over the initial prompt
attention(x)
a(x) attention(z) + Ei attention(t.)
(17)
15

---

Selective Prompt Anchoring for Code Generation
1.0
CodeGen-350M-Mono
Percentage of attention to the Initial prompt
DeepSeek-Coder11BBBnntttt
DeepSeek-Coder-6.7BB Inntstrrut
0.8
CodeLlama-78
Depp
0.6
0.4
0.2
0.0
50
100
150
200
250
300
350
400
Number generated tokens
Figure 5. Shift of LLMs' gradient-based attention the initial prompt. The gradient is calculated with respect the output logits.
Given that attention analysis requires open sourcing, we select five SOTA code LLMs with various sizes. We run the experiments on HumanEval (Chen et al., 2021), one of the most popular benchmarks for evaluating code generation models. We run five LLMs (Nijkamp et al., 2023; Rozière et al., 2024; Guo et al., 2024) on all 164 Humaneval tasks. Figure 5 and Figure show the gradient-based attention shift when generating the first 400 tokens. The value gradually becomes noisy due to the insufficient length of the generated sequence.
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 they generate more code. LLMs tend to attend to code generated by themselves. Our finding is supported by another study (Chiang & Cholak, 2022) which investigates the self-attention dilution of transformers in a more general scenario.
C. Extended Discussion of Approximation in SPA
In Section 3.3, Equation delivers the approximation by only keeping the first derivative in Equation 8, but is also feasible to calculate higher-order approximation. For example, if we want to keep the term involving the second-order derivative 2 2r Fo.s." (00), it can still be computed using diffe-difference methods:
Fo,i. "(0) Fg.i, (1) -2 (0) Fe.i.x (-1)
(18)
(1 0)2
Fo.i.z (-1) can be solved by Equation 12 where Fo, (0) and F (1) are the logits generated from the original input and the logits generated from masked input.
However, no matter how many terms keep in Equation 8, we find we can only represent Fo, (w) linear combination of F(0) and F11), weighted by unknown variable
In Section 5.6, our experiments reveal that w's impact on code generation performance follows an unimodal pattern pattern-initillly increasing, then decreasing. Due to its distribution simplicity, we argue that while higher-order approximation may yield a more reasonable performance distribution across different w values, it does not significantly affect the process of locating the optimal anchoring strength. Therefore, beyond its computational efficiency, the first-order approximation in SPA is adequate for calculating semantically accurate augmented logits.
16

---

Selective Prompt Anchoring for Code Generation
D. Code Generation Examples
Figure 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." contrast, SPA enhances the model's focus the intended purpose. The code initializes all the upper vowels in the first line and correctly refers it later.
In the second example, DeepSeek-Cone Coderr (1.3B) erroncously 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.
count_upperie
def count, uppen
Selective Anchored Text
Given string 5, count the number uppercase vowels even indices.
Given string count the number uppercase vowels even Indices.
Fon example:
For example:
count, upper aBCdEf") neturns
count_uppert "'aCCdE"')
count_upper( "abcdeegg) count_upper( 'dBe") returns
count, _uppttrpppr___ ('abcdeffg)) returns a count_upper( (BBEE)
vowels "AEIOU"
index, letter enurerate(s): if Index 8:
For in range(e, len(s), 2): (1))e))e)(( count
count +-
X
return
def by lmngth(arr::
def (a_le))h((arr::
Selective Anchored Text
Given an arnay of integers, sort the integers inclusive,,
Given array integers, sort integens that ane betseen and Inclusive, reverse the resuiting array, and then replace each digit by corresponding
resulting areay, and then replace each digit by
corresponding nase "One", "Two", "Three", "Four", "Five", "SSX", "Seven", "Eight", "Nine".
"One", "Tuo", "Three", "Four", "Five", "Six", "Seven", "Eight", "Nine".
For examole:
by_]angth([3,, 1, B, 2, s])
exexple:
returns ["eight", "three", "tue", "huo", "One") by_length(()
returns ["Eight", "Three", "Two", "Two", "One") by_length([)
(e_ -1, 55]) 55)) ['One']
by_length([I, -1, 23])
2
returns ('One')
nue, to word= {
'Bheee'
4: "Four", "Five", 6: "Six", "Seven', "Eight", 'Nine')
nun_to_naa
"One' "woo 'Three', Four", 5: "Five", 6: "Six",
'Eight' 9: "Kine')
volid_nuns nuns [num for nuv in are if <- nun 9]
If nus
valid_nure sort()
result. (nus, word (nue])
valid nuns, reverse()
(rint()
nomes [nun_to_ nonnameeeenaeooooeeeeeeeeeeeeees [nu]] nux in velid_rums)
result, (raverseerr)))
X
return result
Figure 6. Examples of generated code LLMs alone (left) and using SPA (right).
17

---

Selective Prompt Anchoring for Code Generation
E. Computational Cost of SPA
E.1. Latency
We analyzed the sources of inference overhead, which include running test cases, the two decoding processes, and the logit arithmetic operations. We observed that each attention augmentation in SPA results in an additional 0.6 times the base model's latency, but this increase occurs only when test cases fail. Running test cases incurs only very small cost (0.1s average) compared to the total inference time (9.6s). Thus, only activating SPA when test failures are detected can reduce the overall overhead. The more accurate the base model, the less overhead SPA introduces. On average, SPA increases the inference time to 1 1.27 times that of the original model. We believe this overhead is acceptable in practical development.
Table 8. Comparison of Tokens Per Second With and Without SPA.
MODEL
TOKEN/SECOND
CODEGEN (350M)
34.1
+SPA
22.7
DEEPSEEK-CODE (I.3B)
17.8
+SPA
14.4
DEEPSEEX-CODER (6.7B)
12.1
+SPA
10.2
CODELLAMA (7B)
14.5
+SPA
10.6
STARCODER (15B)
7.4
+SPA
6.2
SHEPSEEK-CODER (33B)
5.3
+SPA
4.6
E.2. Memory
SPA requires the storage of the logit generated by the masked prompt embeddings. Theoretically, the additional memory requirement, denoted as Moverhead: can be expressed as:
Moverhead = V x D x Stogit
(19)
where represents the vocabulary size, D is the token embedding dimension, and Stogi the size single logit value.
Consider vocabulary size 50, 000, a token embedding dimension 4, 096, and logit size Stogit 2 bytes. Then the additional memory overhead is calculated
50, 000 x 4, 096 2 bytes R 390 MB
(20)
In practice, this memory overhead can be significantly reduced because many low-ranked tokens hardly contribute to the results. We can calculate the augmented logits by considering only the top-ranked tokens and ignoring the rest. For example, we focus only on the top 100 logits, the overhead will dramatically decrease to 800 KB.
E.3. Tradeoff Between Latency and Memory Usage via Parallelization
Currently, SPA is implemented by sequentially computing the logits. However, the latency can be further reduced by parallelizing the logits generation from the original and masked embeddings, as these computations are independent (Figure 1). This optimization, however, comes at the cost of increased memory usage (double VRAM or memory for forwarding). Given that LLM inference is memory- memory-intensive, this introduces tradeoff between reduced latency and higher memory consumption.
F. Pass @ 10 & Beam Search with SPA
To further evaluate the generalizability SPA, we assess its Pass@ performance using beam search. While SPA produces augmented logits that could be used directly for beam search, we observed that directly sampling top beams from these
18

---

Selective Prompt Anchoring for Code Generation
augmented logits does not improve the performance. We hypothesize that this phenomenon occurs because while SPA successfully amplifies the influence of anchored text and improves the accuracy of top logits, it also amplifies noise in lower-ranked logits. This undermines the reliability of the overall probability distribution, thereby hindering the sampling process.
To address this issue, we retrieve top candidate tokens based on the augmented logits but use original probabilities to accumulate beam probability. This strategy ensures that important, potentially overlooked tokens are considered while maintaining reliable probabilities.
Table 9. Pass@1 @1 and Pass@ 10 (%) with and without using SPA.
HUMANEVAL
HUMANEVAL+
MBPP
MBPP+
MODEL
SIZE
@ASS@1 1
PASS@ 10 PASS@1 1 PASS@ 10 PASS@1 1
PASS@ 10 @ I
PASS@ 10
CODEGEN-MONO (350M)
15.3
36.6
12.2
33.6
19.6
47.7
15.9
42.4
+ SPA
20.1 (44.8)
39.0 (x2.4)
17.1 (+4.9)
38.5 (44.9)
27.4 (43.8)
55.6 (47.9)
22.6 (16.2)
43.3 (+0.9)
TERSSEK-CODDES (I.3B)
66.4
73.3
61.8
68.7
58.2
67.0
52.4
63.7
+ SPA
70.1 (43.7)
73.3 (10.0)
67.7 (45.9)
69.3 (10.6)
60.9 (42.7)
68.8 (+1.8)
52.4 (10.0)
64.3 (10.6)
DEEPSEEK-CONE -CODER (6.7B)
75.6
84.0
70,2
77.9
67.0
79.8
58.5
70,2
+ SPA
88.5 (412.9)
86.4 (+2.4)
79.9 (49.7)
82.8 (14.9)
71.0 (14.0)
87.7 (47.9)
60.7 (v2.2) 73.9 (43.7)
CODELLAMA
(7B)
33.6
58.0
28.2
48.9
50.9
61.0
40,8
49.0
+ SPA
44.0 (-10.4)
64.7 (+6.7)
36.0 (+7.8)
54.4 (+5.5)
54.3 (+3.4)
65.3 (+4.3)
44.0 (+3.2) 52.0 (+3.0)
DEEPSEEK- -CDERR (33B)
81.7
88.5
77.1
80.2
73.4
86.8
63.2
75.8
+ SPA
86.2 (+4.5)
89.7 (+1.2)
79.3 (+2.2)
83.2 (+3.0)
79.4 (+6.0)
89.2 (+2.4)
70.3 (+7.1)
80.1 (+4.3)
As shown in Table 9, SPA consistently improves Pass@ 10 by, on average, 3.42% and up to 7.9% when using beam search. While the improvements not pronounced as those seen with Pass @ 1. we anticipate that future work could develop beam search algorithms specifically optimized for SPA's 's unique logit distribution characteristics.
G. Hypothetical explanation for Attention dilution and SPA'S effectiveness
SPA is motivated by a recent study (Kou et al., 2024) and our empirical observations demonstrating the attention dilution issue. Our experiment results in Section 5 echo our observation and confirm the existence of attention dilution during code generation. Here we propose detailed explanation for this phenomenon based on our knowledge and hypotheses. We believe it stems from two limitations in regular autoregressive decoding: (1) Distraction and (2) Error propagation.
Distraction. When transformer generates a token, its correctness depends on two abilities: (1) whether the model attends to the correct context, and (2) whether the model can derive the correct token based on this context. SPA aims to improve the first ability. Suppose we have a perfect transformer. For each generated token, should only attend to relevant prior tokens and ignore irrelevant ones. However, no model is perfect. For each prior token, there is a chance the to it. As the model generates more tokens that compete for attention, it becomes increasingly challenging for the model to accurately distribute its attention. The model has more chances attend to irrelevant tokens, making its attention increasingly unreliable.
model incorrectly identifies and attends
In contrast, the user prompt is persistently relevant throughout the generation since it represents the user's intent. While eeff-generated tokens are also important context, they are less persistently related than task descriptions in code generation. Amplifying the impact of the task description by SPA essentially enhances attention reliability, thereby mitigating distraction.
Error propagation. During code generation, the model may generate irrelevant code tokens. However, the autoregressive nature of LLMs assumes that all previously generated tokens correct. For example, if the model introduces an irrelevant variable declaration, subsequent generations may take it into account and continue to generate irrelevant code. Although the model can still generate correct code behavior in later generations, this assumption of correctness makes it difficult to identify errors. As a result, the error can propagate and accumulate, leading to a higher probability of errors in later generations.
SPA mitigates this issue by reinforeing attention to the user prompt while downplaying reliance on self-generated tokens.
19

---

Selective Prompt Anchoring for Code Generation
This optimizes the attention distribution based the trustworthiness of different contexts, thereby increasing accuracy.
H. Implementation and Deployment
H.1. Implementation of SPA
SPA is a model-agnostic algorithm and our implementation does not rely on specific models. All six models in our paper are built upon the Huggingface Transformer library, which offers APIs to directly access and edit token emheddings and logits. Particularly, the SPA generator inherits the Huggingface Transformers generation API. We leverage the hook to modify the logit calculation within the original generation pipeline. The API works for any LLM in the Huggingface model collections with native hyperparameters such TEMPERATURE. We have released PyPI library for developers quickly test SPA.
H.2. Model Deployment
We downloaded and deployed LLMs from Huggingface. To expedite evaluations, we apply 8-bit quantization (Frantar et al., 2023; Dettmers et al., 2022) to all models. Prior studies (Li et al., 2024; Huang et al., 2024) have demonstrated that this approach has very little impact on LLM performance. We set the Temperature to 0 and the beam to 1 for greedy decoding in all experiments, except for the one described in Appendix F. All experiments were conducted on a 64-bit Ubuntu 22.04 LTS system, equipped with an AMD EPYC 7313 CPU, eight NVIDIA A5500 GPUs, and 512GB of memory. The experiments ran for approximately seven weeks.
H.3. Prompt Design
We use the original task descriptions from the datasets as prompts for the text-completion models, CodeLlama and CodeGen- Mono. For the remaining models, we format the prompts using the official chat template from HuggingFace. All experiments are conducted zero-shot setting.
Evaluating SPA on Other Generative Tasks
To evaluate the generalizability of SPA beyond code generation, we experiment SPA on other generative tasks, including TruthfulQA (Lin et al., 2022), GSM8K (Cobbe et al., 2021), MMLU (Hendrycks et al., 2021b), and BoolQ (Clark et al., 2019).
TruthfulQA (Lin et al., 2022) a benchmark designed to measure models' ability to avoid generating false or misleading information, requiring models to answer questions while remaining truthful. GSM8K (Cobbe al., 2021) tests mathematical reasoning capabilities through grade school math word problems that require multi-step solutions. MMLU (Hendrycks et al., 2021b) evaluates models across 57 subjects, including elementary mathematics, US history, computer science, law, and more, comprehensively testing both breadth and depth of knowledge. It provides multiple-choic questions for LLMs to identify the correct answers. BoolQ (Clark et al., 2019) consists of naturally occurring yes/no questions from web queries, testing reading comprehension and binary classification abilities. We use Llama I-Intrucct-8 as our base model since the other models in our study are specifically fine-tuned for code tasks.
Table 10. Evaluating SPA on Different Generative Tasks,
TRUTHFUL.Q QA TRUTHFUL GSM8K MMLU BoOLQ HUMANEVAL (TRUTH.)
MODEL
(INFO.)
88.40%
97.78%
76.5%
69.5%
83%
63.4%
SPA
3.14%
+ 0..1%%
1.10%
0%
0.03%
9.76 %
As shown in Table 7, while SPA provides improvements across these tasks, the gains are significantly smaller compared
http:///s//ee..iii.ccassss////ee classes/text gewerationnoooo
20

---

Selective Prompt Anchoring for Code Generation
code generation tasks. We attribute this performance difference to the unique input-output pattern of different generative tasks.
For MMLU, the model only needs to generate choice to answer the multiple-choice question. The output length is significantly shorter than in code generation and considerably less than in user prompts. Consequently, the model's 's attention is hardly diluted by self-generated tokens, making SPA not helpful. In contrast, TruthfulQA requires the model to generate a text analysis, provide reasoning, and answer the question. Therefore, SPA is more beneficial in addressing attention dilution and correcting the remaining 27% errors for Truthfulness and 41% errors for Informativeness. In contrast, code generation typically lengthy and can easily lead to attention dilution. Moreover, code generation prompts serve as persistent instructions, requiring the LLM to maintain focus throughout the generation process. This differs from tasks like translation, where there is no inherent need to consistently anchor attention to specific components of the input.
Therefore, we believe SPA is especially suitable for enhancing LLMs' attention for code generation tasks, Nevertheless, it remains an interesting future work to investigate how to further improve SPA for a broader range of generative tasks beyond code generation. More details are reported in Appendix
J. Tuning of Anchoring Strength
As demonstrated in Section 5.6, different anchoring strengths w lead to different performance and follow a simple unimodal pattern. In this section, we report more evaluation of tuning SPA. To ensure fair comparisons, SPA is activated for all tasks in all experiments, unlike during inference time where SPA is only activated for failed tasks. Moreover, we tune SPA on all the tasks for better generalizability.
J.1. Tuning Stability
We evaluate the stability of tuning anchoring strength by tuning SPA on five exclusive subsets of each dataset for each model, as shown in Table 11. The average variance across subsets is 0.0046 for HumanEval/HHumanEvaa and 0.0026 for MBPP/MBPP+, demonstrating the tuning stability.
J.2. Cross-dataset & Cross-model Tuning Evaluation
We investigate the transferability of this hyperparameter across different models and datasets. Firstly, we conduct a cross-datasel evaluation between Humanniaa/HuumanEvale and MBPP/MBPP+, which have distinet prompt formats. We tune w on HumanEval+ and evaluate Pass@ @ on MBPP and MBPP+, and vice versa (denoted as SPAgross-dataser We calculate average Pass@) 1 improvements the original and plus versions across all baseline models. Secondly, we perform cross-model evaluation by tuning on one model and evaluating Pass @ 1 on the remaining four. For each model, we compute the average @ass@1 improvements across all the other models, for HumanEval/HuumnnEyak and MBPP/MPP++ respectively (denoted as SPhorose-model) Similar to Section 5, SPA represents tuning within the split partial dataset, while SPA" represents tuning within the entire dataset.
As shown in Table 12, we find the anchoring strength w tuned on one model hardly transferred to another. However, w tuned on one dataset can be transferred another with reduced but still effective performance. These observations suggest that the anchoring strength is highly model-dependent and partially task-dependent.
We further investigate whether it is possible to find a universal anchoring strength that works for most scenarios. One potential value is the average of the tuned anchoring strengths across benchmarks for each model. We apply this fixed anchoring strength to all settings, denoted as SPA-preset. As shown in Table 13, although the generation accuracy decreases compared to using a tuned strength, SPA with the fixed strength can still outperform the baselines. It implies that SPA is effectively deployable in new scenarios once reasonable value is set.
J.3. Comparison between Tuned SPA and Optimal SPA
To better understand the tuning efficiency, we tune SPA on each complete dataset to obtain the optimal anchoring strength (denoted as SPA*). As shown in Table 14, we find that SPA* * achieves 20% higher performance improvements compared to
The "plus" versions of HumanEval and MBPP share identical prompts with their base counterparts, so we only tune once on the plus version.
21

---

Selective Prompt Anchoring for Code Generation
Table 11. Tuned Anchoring Strength Different Subsets.
MODEL
SUBSET
HUMANEVAL/4
MBPP/+
CODEGEN-MONO (350M)
SUBSET)
1.05
1.30
SUBSET2
1.10
1.35
SUBSET3
1.20
1.25
SUBSET4
1.30
1.35
SUBSET;
1.25
1.35
COMPLETE
1.20
1.35
SEEPSEEK-CODER (1.3B)
SUBSET,
1.05
1.20
SUBSET2
1.05
1.15
SUBSET3
1.10
1.15
SUBSET
1.00
1.20
SUBSETs
1.05
1.25
COMPLETE
1.05
1.20
CHEPSEEK-CODERR (6.7B)
SUBSET,
1.30
1.30
SUBSET2
1.25
1.20
SUBSET3
1.30
1.25
SUBSET4
1.20
1.20
SUBSET;
1.35
1.25
COMPLETE
1.28
1.25
CODELLAMA (7B)
SUBSET
1.55
1.25
SUBSET2
1.55
1.20
SUBSET3
1.50
1.20
SUBSET4
1.65
1.25
SUBSET;
1.65
1.15
COMPLETE
1.60
1.20
DEEPSEEK-CODEL (33B)
SUBSET
1,25
1.25
SUBSET2
1.30
1.30
SUBSETa
1.40
1.30
SUBSET+
1.35
1.40
SUBSET;
1,35
1.20
COMPLETE
1.35
1.30
Table Pass@ improvements (%) based on cross-dataset tuning.
DATASET
SPAcross dataset
SPAcross model
SPA
SPA"
HUMANEVAL/+ MBPP/
2.01
00.9
4.36
5.11
2.50
+ 0.37
2.86
3.57
Table 13, Comparison between SPA and SOTA methods.
METHOD
APASS (%)
TIme (SEC)
BASE MODEL
7.7
PASTA
48.8
DEBUGDEBGGG
+4.2
27.3
SELF-EDIT
+1.8
26.4
SELF-PLANNING
+3.6
21.6
REACT
+1.3
28.8
SPA-HREDDST
+5.0
9.8
SPA
+7.7
9.8
the tuned SPA on average.
Table 15 reports optimal anchoring strength values w. We observe that the average value of 1.28 can be used to effectively improve performance across all benchmarks for all L.MSs.
22

---

Selective Prompt Anchoring for Code Generation
Table 14. Pass@1 Improvements (%) Comparision between Tuned SPA and Optimal SPA (SPA is applied all tasks).
MODEL
SIZE
HUMANEVAL
HUMANEVAL+
MBPP
MBPP+
CODEGEN-MONS (350M)
15.3
12.2
19.6
15.9
SPA
18.3 (+3.0)
16.0 (+3.8)
24.9 (+5.3)
20.6 (+4.7)
SPA*
18.3 (+3.0)
16.0 (+3.8)
24.9 (+5.3)
20.6 (+4.7)
DEEPSEEK-CODER (1.3B)
66.4
61.8
58.2
52.4
SPA
(+3.1)
66.4 (44.6)
59.1 (+0.9)
52.4 (40.0)
SPA"
71.0 (+4.6)
66.4 (+4.6)
61.7 (+3.5) 53.4 (+1.0)
DEEPSEEK-CONE (6.7B)
75.6
67.0
58.5
+ SPA
83.2 (+7.6)
75.6 (+5.4)
69.6 (+2.6)
60.2 (+1.7)
SPA"
84.0 (+8.4)
76.3 (+6.1)
72.2 (+5.2) 61.1 (+2.6)
CODELLAMA
(7B)
33.6
28.2
50.9
40.8
SPA
40.5 (+6.9)
33.6 (+5.4)
52.9 (+2.0)
43.1 (42.3)
SPA"
41.2 (+7.6)
35.9 (+7.7)
$2.9 (+2.0)
43.1 (+2.3)
DEEPSEEK-CON CCDEER (33B)
81.7
77.1
73.4
63.2
SPA
84.7 (+3.0)
77.9 (40.8)
77.2 (+3.8)
68.5 (45.3)
SPA"
85.5 (43.8)
(+1.5)
77.2 (+3.8) 68.5 (45.3)
Table 15. Optimal Anchoring Strength (w) 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.38)
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
DREPSHRK-CODER (33B)
1.35
1.35
1.30
1.30
1.33
Average
1.30
1.30
1.33
1.33
1.28
K. Analysis of ' Performance without Test Case Availability
To analyze SPA'S 's effectiveness in scenarios without available test cases, we conduct additional experiments where SPA is applied to all code generation tasks (S) As shown in Table 17, SPA can still achieve significant improvements across all benchmarks and models. For example, on HumanEval, SPA improves DeepSeek- (6.7B)'s Pass@ by 7.6% (from 75.6% to 83.2%).
Table 16. Comparison PASTA in terms of Pass @ 1.
HUMANEVAL
HUMANEVAL+
MBPP
MBPP+
BIGCODEBENCH
PaSTA
+1.22
+1.22
+1.17
+0.94
+0.1
SPA
+7.32
+4.88
+4.22
+3.51
+0.4
L. Comparison to PASTA
Unlike mainstream methods that optimize prompts, SPA enhances performance by optimizing model attention. We compare SPA with another attention steering method, PASTA (Zhang et al., 2024). PASTA's implementation only supports LLAMA (Touvron et al., 2023a), LLAMA2 (Touvron et al., 2023b), and GPT-J (Wang & Komatsuzaki, 2021). We find that adapting PASTA to new models requires sigmificant effort, so we only evaluate on LLAMA- -B, LLAMA2-7B, and GPT-J-6B.
23

---

Selective Prompt Anchoring for Code Generation
Table 17. Absolute (A) and Relative (1) Performance improvements in @ rates (%) SPA all indicates SPA is activated when original generated code fails test cases.
is applied to all tasks, while SPA
MODEL
SIZE
HUMANEVAL
HUMANEVAL+
MBPP
MBPP+
BIGCODEBENCH
CODeGEN-MONO
(350M)
15.3
12.2
19.6
15.9
1.1
SPA
18.3
A43.0
16.0
443.8 (3156 T)
24.9
445..
20.6
1.4
4-0.3
(20% 1)
(279 ")
(30% 1)
(27% 1
SPA
20.1
A4448
17.1
A+44.9
27.4
A4778 (409)
22.6
6,6.7
1.6
A+-0.5 (45%
(31%
(10% t)
(42% 1)
DEEPSEBK-CODER
(1.3B)
66.4
61.8
58.2
52.4
2.5
SPA
69.5
A+311
66.4
59.1
A+0.9
52.4
3.3
A+0.8
(5% 1)
(796 1)
(2% 1)
(0% 1)
(32% 1)
SPA
70.1
A43.77
67.7 (10% 1)
A460.9
60.9
A-277
52.4
0.000
3.4 (30% 1)
A+0..9
(6% f)
(0%6 12
(0% 13
DEEPSEEK CRERSEEK-CODER
(6.7B)
75.6
70.2
67.0
58.5
12.7
-SAA
83.2
A-
75.6
5.4.
69.6
A+22.6 (496 1)
60.2
14.2
A41..
(10% 1)
(896 t
(39) 1)
(12% 1)
+PPA
88.5
A--12.0
79.9
A449.7
71.0 (6% 1)
A444.0
60.7
16.4
A-43.77
(17% 1)
(14% f)
(496 t)
(29% f)
CODELLAMA
(7B)
33.6
28.2
50.9
40.8
3.4
+ SPA
40.5
A.66.
33.6
A+5.4
52.9
A+2.0 (4% 1)
43.1
2.3.3
3.8
A+0.4
(219% 1)
(19% 1)
(6% 1)
(12% 1)
SPA
44.0
A+100.4
36.0
4+7.8
54.3
A+3.4
44.0
4.1 A++0.7 (21% 1)
(31% 1)
(28% t)
(7% 1)
(8% 1)
STARCODER2
(16B)
67.7
60.4
78.0
65.1
13.3
+ -SAA
72.1
A+44.4
63.6
+-3
80.9
A22299
67.6
A-2.5
14.1
0.8
(6% f)
(6% ()
(496 13
(4%
(6% ft)
+ -PA
75.6
A-77.9
65.6 (9% f)
4.2
82.0
4444.0
69.1
14.3
A.10.0
(12% 1)
(6%6 1)
(0% 1)
(8%6 f)
DEEPSHEK-CODER
(33B)
81.7
77.1
73.4
63.2
18.9
+ SPA
84.7
A++3.0 (496
77.9
0460.6 (196
77.2
4-3..8 (6%6 1)
68.5 A 5.3
20.7
L+118
(8%
(1.0% 11)
+ SPA
86.2
A-44.5 (6% 1)
79.3
24222 (3% 1)
79.4
A466.0 (8% 1)
70.3
A-7.1
21.5
A42266
(11% 1)
(14% 1)
Table 16 demonstrates that SPA consistently outperforms PASTA on five benchmarks, achieving 4X higher Pass@I on average. We attribute this to two hypothetical reasons. First, internally editing model attention can be sensitive and exhibit unexpected behaviors. Second, identifying effective attention headers may not be stable or generalizable across different tasks. In contrast, SPA augments the final logits without internally editing the model's feed-forward process, making it model-agnostic. Furthermore, SPA only introduces a single hyperparameter, and the tuning is stable (discussed in Appendix J.1).
M. Confidence-Moddlltted Anchoring Strength
To make SPA more stable and eliminate the need for tuning, we propose confidence-moodulated weighting method that adjusts each token's anchoring strength based on its original probability respectively.
When the LLM is highly confident about a specific token, the corresponding logit is adjusted less, resulting in lower anchoring weight applied to the logit difference. Conversely, if the LLM shows lower confidence, a relatively higher weight is applied, leading greater adjustment of the logit value. The confidence level is measured using the original probability distribution generated by the LLM.
Formally, let Pt represent the predicted probability of token t within the original probability distribution for the next token. The anchoring strength Wt for token t is defined as:
Wi = A pt),
(21)
where À is fixed coefficient that controls the overall anchoring strength. This results in a token-wisse anchoring strength vector w E R over the vocabulary V. The vector w is then element-wise multiplied with the original logit difference vector obtain a reweighted logit difference (Figure
24


Reviews:
[{"review_id": "vGzgNFI4Mb", "reviewer": "Reviewer_8eZp", "summary": "This paper identifies an attention dilution problem in code generation using LLMs, where models pay decreasing attention to the user prompt as more code tokens are generated. To address this issue, the authors propose Selective Prompt Anchoring (SPA), a model-agnostic approach that amplifies the contextual impact of user prompts during generation. SPA works by calculating the difference between logit distributions from original and masked versions of anchored text, then scaling this difference by a hyperparameter. Experiments across multiple benchmarks (HumanEval, MBPP, etc.) and models demonstrate consistent improvements, with Pass@1 increasing a bit.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "`While correlation between incorrect code and longer generation length is shown, this doesn't necessarily show causation, e.g, \"code is longer\" maybe just mean the question itself is harder.`\n\nThis is a good point! We investigated this using LiveCodeBench, which provides the difficulty level for each task. For tasks with the same difficulty level, we still observed a significant difference between correct solutions and incorrect solutions. We appreciate your insightful feedback and will include this new result in the appendix.\n\n|         |   Easy   |  Medium  |   Hard   |\n|---------|:--------:|:--------:|:--------:|\n| Passed  |   294    |   475    |   400    |\n| Failed  |   418    |   664    |   784    |\n\n\n\n`inference time is also analyzed, but maybe better to discuss a little bit about potential memory overhead or precomputation requirements`\n\nThere is no precomputation requirement. For memory overhead, SPA needs to store the logits computed from masked prompt embeddings. Theoretically, the extra memory is equal to vocabulary_size * token_embedding_dimension * logit_size.\n\nSuppose:\n\n- vocabulary_size = 50,000\n- token_embedding_dimension = 4,096\n- logit_size = 2 bytes\n\nThe logit overhead will be 50,000 * 4,096 * 2 ≈ 390 MB.\n\nIn practice, this memory overhead can be further reduced because low-ranked tokens do not contribute significantly. We can only augment the logits of a few top tokens. For example, if we consider the top 100 logits, the overhead will dramatically drop to 800 KB.\n\nWe will discuss this in the paper.\n\n`My understanding is that the SPA method amplifies the attention of specific tokens of the prompt. Now, would this approach, to some extent, change the model’s original behavior, and like, possibly affecting outputs that were initially correct? It seems plausible that it could introduce some bias, but I suppose the actual impact would depend on the context and how the method is applied.`\n\nIndeed, as shown in Figure 6, different anchoring strengths ω of SPA would have different impacts. If the strength is too low, the performance improvement is low; if it is too high, the LLM starts becoming biased, leading to a decline in performance. Therefore, the actual impact depends on how we set this hyperparameter. Fortunately, we observe that it is simple to tune a balanced value where performance significantly improves. Please check more details in Appendix J. We will clarify this in the paper.\n\n`A limitation is the selection of \"anchored text.\" This approach appears to be somewhat fixed and not \"clever\", and constrained to HumanEval-related tasks and may not generalize well to more complex programming scenarios (e.g. in some other code tasks, NL might not be that important compared to code\") . While the authors explore various selection strategies in Section 5.4, a more intelligent approach for identifying optimal anchored text across diverse programming contexts would significantly strengthen the practical applicability of SPA.`\n\n`How would SPA perform with more dynamic anchored text selection strategies that change during the generation process?`\n\nIt is a very interesting future work to investigate how to intelligently select anchored text. One idea is to use LLMs to first select important words or phrases to anchor on before code generation. For tasks where NL may not be important compared to code (e.g., code translation), one idea is to use static code analysis to identify important code elements (e.g., function calls and variable names heavily used in the code). \n\nFurthermore, the anchoring strength, which determines the degree of attention steering, can also change dynamically at different steps. One idea is to develop a method to calculate the relevance of words and phrases in the user prompt to each decoding step. Based on the relevance scores, SPA can dynamically assign higher values to more relevant contexts while assigning lower values to less relevant ones.\n\nWe will discuss these ideas as future work in the paper.\n\n`The paper shows that SPA is particularly effective for longer prompts. Would it be suitable for more complex programming tasks like super long prompts (even whole repos) and sometimes, random prompts (hard to define which parts are more important)?`\n\nSPA is suitable for super long prompts since specific instructions like \"do not use lambda expressions in the generated code\" are likely to be buried in the long prompt and not followed by the LLM. SPA will help improve the influence of such specific instructions. \n\nIn fact, it may be more helpful to use SPA in a self-improving pipeline for fairly complex tasks. Based on the errors of the initially generated code, we prompt the LLM to identify which instructions or requirements are not followed in the prompt and then use SPA to improve the influence of these instructions/requirements.", "claims_and_evidence": "The primary claims regarding attention dilution and SPA's effectiveness are well-supported. The authors provide persuasive empirical evidence demonstrating attention patterns across multiple models (Figures 2-4). The performance improvements are consistently observed across diverse models and benchmarks.\n\nHowever, the claim that attention dilution is the reason for code generation errors could be better substantiated. While correlation between incorrect code and longer generation length is shown, this doesn't necessarily show causation, e.g, \"code is longer\" maybe just mean the question itself is harder.", "methods_and_evaluation_criteria": "- The evaluation is comprehensive, using established benchmarks (HumanEval, MBPP, and variants) along with multilingual code generation (HumanEval-X) and more complex real-world tasks (BigCodeBench). \n- The comparison against both attention steering (PASTA) and prompting methods (Self-Debugging, Self-Edit, etc.) provides a well-rounded assessment. \n- inference time is also analyzed, but maybe better to discuss a little bit about potential memory overhead or precomputation requirements", "theoretical_claims": "The mathematical derivation of SPA in Section 3 is generally sound.", "experimental_designs_or_analyses": "The experimental design is thorough with appropriate controls and ablations. \n\nA limitation is the selection of \"anchored text.\" This approach appears to be somewhat fixed and not \"clever\", and constrained to HumanEval-related tasks and may not generalize well to more complex programming scenarios (e.g. in some other code tasks, NL might not be that important compared to code\") . While the authors explore various selection strategies in Section 5.4, a more intelligent approach for identifying optimal anchored text across diverse programming contexts would significantly strengthen the practical applicability of SPA.", "relation_to_broader_scientific_literature": "- The paper effectively situates SPA within both code generation and attention steering literature. The connection to the psychological concept of the \"anchoring effect\" provides an interesting interdisciplinary perspective.\n\n- The work advances upon prior attention steering methods like TOAST and PASTA by offering a more model-agnostic approach with lower computational overhead. It also contributes to the growing body of research on training-free methods for improving LLM performance.\n\n- A notable constraint is that SPA appears most effective for code generation tasks specifically. The paper's own experiments in Section 5.6 show limited effectiveness on other generative tasks (e.g., MMLU), suggesting task-specificity that restricts its broader applicability in the LLM literature.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "See above comments", "other_comments_or_suggestions": "See above comments", "questions_for_authors": "- My understanding is that the SPA method amplifies the attention of specific tokens of the prompt. Now, would this approach, to some extent, change the model’s original behavior, and like, possibly affecting outputs that were initially correct? It seems plausible that it could introduce some bias, but I suppose the actual impact would depend on the context and how the method is applied.\n- How would SPA perform with more dynamic anchored text selection strategies that change during the generation process?\n- The paper shows that SPA is particularly effective for longer prompts. Would it be suitable for more complex programming tasks like super long prompts (even whole repos) and sometimes, random prompt (hard to define which parts are more important)?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "R5HzpT0RoT", "reviewer": "Reviewer_Nvi5", "summary": "This work identifies attention dilution as a cause of code performance worsening as the context (generated code) increases. They subsequently propose a solution based on attention steering to upweight relevant tokens. This effectively shortens the effective context of the model and they show substantial performance (12.9%) improvement.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "`I worry that MBPP and HumanEval are in the memorization regime... `\n\nThank you for the suggestion. We conducted a new experiment on LiveCodeBench (10/1/2024-2/1/2025). The results show that SPA remains effective.\n\n| Model                        | LiveCodeBench             |\n|-----------------------------|---------------------------|\n| Codegen-mono (350M)         | 1.1                       |\n| +SPA                        | 1.5 (+0.4) (36%↑)         |\n| DeepSeek-Coder (1.3B)       | 6.5                       |\n| +SPA                        | 9.2 (+2.7) (42%↑)         |\n| DeepSeek-Coder (6.7B)       | 7.8                       |\n| +SPA                        | 10.8 (+3.0) (38%↑)        |\n| CodeLlama (7B)              | 3.8                       |\n| +SPA                        | 4.0 (+0.2) (5%↑)          |\n| StarCoder2 (16B)            | 7.0                       |\n| +SPA                        | 8.2 (+1.2) (17%↑)         |\n| DeepSeek-Coder (33B)        | 11.9                      |\n| +SPA                        | 15.8 (+3.9) (33%↑)        |\n\n`It's not clear ... Is this rerunning included in the 1.27x overhead?`\n\nThe inference overhead includes running test cases, the two forward passes, and the logit arithmetic operations. We did a further analysis of these different sources of overhead and found that the overhead for running test cases only incurs a very small overhead (0.1s on average) compared to the entire inference process (9.6s). Thus, only performing steering in case of test failures actually reduces the overhead. We did not perform a parallel decoding but we think this is a great idea and will discuss how this will further reduce the overhead in future work. Thank you!\n\nIn our experiments, we use test cases in the benchmark instead of prompting the LLM to generate new test cases. Leveraging existing test cases is a common practice in code generation, such as [1] and [2]. But we agree that if there are no test cases available, prompting the LLM for test generation could be a worthwhile solution even with the additional overhead.\n\n[1] Chen et al. Teaching large language models to self-debug. ICLR 2024.\n\n[2] Zhang et al. Self-Edit: Fault-aware code editor for code generation. ACL 2023. \n\n`\"For simplicity, we demonstrate by making the entire user prompt x as the anchored text.\"...`\n\nSorry for the confusion. We initially intended to use anchoring the entire prompt as an example scenario to illustrate how the attention steering mechanism works in Section 3.3. We will change this writing and present the anchored text more generally and formally.\n\n`Ei(X, ω) is not formally defined so It's not clear why \"Scaling the semantics of X by ω times is not equivalent to multiplying X by ω.\"`\n\nEi(X,ω) represents the augmented embedding matrix of Ei (Line 142 after Eq. 5), where the semantic influence of X on the generated output is scaled by ω. Simply multiplying the embedding of X by ω does not simply improve the “**semantic influence**” of X since the embedding of X also encodes other non-semantic information such as the positional information. This is why we compute the difference between the logits when masking and unmasking X to cancel out noises and some of the non-semantic information. We will clarify this in the paper.\n\n`In what way is the attention being steered and not just the final logits?`\n\nAs shown in Section 3.3., SPA can mathematically simulate attention steering via logit arithmetics. We chose not to directly modify self-attention since it is brittle and costly. For example, PASTA requires an expensive model profiling stage to identify usable attention headers to steer. Furthermore, only modifying self-attention layers does not account for the impact of other components like the feedforward layers and may cause an adversary effect given that self-attention and other components are trained together. As shown in Table 2, SPA is more computationally efficient and also achieves better performance than PASTA.\n\n`Can you provide intuition for the code domain why the last layer \"represent the most accurate attention distribution\"?`\n\nAccording to [3], while lower attention layers capture local dependencies such as the variable relationship within an expression, deeper layers capture more abstract representations with long-distance dependencies such as the control flow. Deeper layer attention distribution mirrors how humans understand programs by integrating information across the project. [4] confirmed this by showing attention distributions in the last layer produce the highest alignment between model and human.\n\n[3] Wan et al. What Do They Capture? A Structural Analysis of Pre-Trained Language Models for Source Code. ICSE 2022.\n\n[4] Kou et al. Do Large Language Models Pay Similar Attention Like Human Programmers When Generating Code? FSE 2024.\n\n`It would be good to move how to select omega in the main body.`\n\nThank you for the suggestion. We will move Appendix J.1 to Section 5.", "claims_and_evidence": "The consistent improvement is impressive especially if omega can be consistently selected without bootstrapping on a subset of problems.\nThe claims seem well substantiated.", "methods_and_evaluation_criteria": "The motivation of the proposed method is difficult to understand but boils down to a simple reweighing of the logits based on masking vs no masking of the prompt text (anchoring text).", "theoretical_claims": "\"For simplicity, we demonstrate by making the entire user prompt x as the anchored text.\" As figure 1 indicates the choice of the anchored text is important this simplification does not help and makes it more confusing as to what the Masking does. It seems that the approach is to mask out the anchored text and see the impact on the log probabilities. \n\nEi(X, ω) is not formally defined and so It's not clear in the derivation why \"Scaling the semantics of X by ω times is not equivalent to multiplying X by ω.\".", "experimental_designs_or_analyses": "It's not clear why exactly there's a 1.25X inference overhead. Is this simply a result of applying omega to the original embeddings? Figure 1  and eq 14 seems to indicate a parallel decoding. Is this included in the 1.25X overhead. \n\nLine 175 come out of nowhere. Why is it that you only run this steering when there are failed test cases? If this consistently leads to improved performance it seems natural to always run with the steering. Additionally, this would implicitly avoid any errors introduced by incorrectly setting omega -- leading to upside only. In practice, you would be incurring at least 2x the overhead by having to generate and then run test cases, before rerunning. Is this rerunning included in the 1.27x overhead?", "relation_to_broader_scientific_literature": "This work proposes a paired decoding technique that's especially effective for coding applications. The derivation is missing some details but it's a neat approach to steering models to pay more attention to relevant code segments. This work relates well to localization for coding agents. In that literature, it's understood that the precision of localization has a large impact on the agent's performance. This work seems to indicate clever attention steering can reduce the effective context and implicitly does localization.", "essential_references_not_discussed": "Previous work has shown attention blocks tend to focus on the recent tokens but also on early tokens (corresponding to system prompt): \n\nStreamingLLMs: https://arxiv.org/pdf/2309.17453\n\nIn settings where the model does indeed focus on the early tokens mostly, the dilution may not be of that much of a concern. However, it would be interesting if it was a code model specific phenomena.\n\nOther work also looks at paired decoding capturing the differential signal based on the prompt: \n\ncontext-aware decoding: https://arxiv.org/pdf/2305.14739", "other_strengths_and_weaknesses": "I worry that MBPP and HumanEval are in the memorization regime. It would support the argument to demonstrate the same trends on LiveCodeBench or other coding tasks to show this behavior also occurs when the model is operating on uncontaminated data.\n\nI really like that the attention steering is a very general approach and should apply also to vision and other non-code language tasks. This initial work provides strong evidence to explore the impact in other domains.", "other_comments_or_suggestions": "It would be good to have the discussion on how to select omega in the main body of the text and move parts of the derivation to the appendix. Empirically, it seems the choice of omega is important.", "questions_for_authors": "Can you provide intuition especially for the code domain why the last layer \"has been shown to represent the most accurate attention distribution\" (line 86)? \n\nIn what way is the attention being steered and not just the final logits? This may be clarified by clarifying Ei(X, ω).", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "umUn0CZ4uK", "reviewer": "Reviewer_fXfw", "summary": "This paper proposes a method for improving LLM's performance at test time. The method helps the LLM keep the focus on the task, avoiding dilution as the number of tokens generated grows for complex tasks such as coding. The method anchors a part of the prompt that specifies the task to accomplish by amplifying its corresponding logits before next token generation. The method compares the logits generated by the LLM with the orginal token embedding matrix and those generated by an embedding matrix where the part corresponding to the anchor has been masked; the difference in logit strength is used to amplify the output logits. The anchor part of the prompt is defined by special tokens. The empirical evaluation of the method comprises 5 code benchmarks (HumanEval, BigCodeBench, etc.), 6 LLM open-source base models (DeepSeek, CodeGEN, etc.) and 5 SOTA methods (ReACT, PASTA, etc.). Results shows consistent improvements (13%) across all benchmarks compared to the baseline and better performance than SOTA methods. Several ablation and parameter studies are performed to determine the best part of the prompt to mask, the influence of the prompt length, the effect of using other programming languages and the tuning of the anchoring strength.\n\n## update after rebuttal:\nAfter carefully reviewing the author's rebuttal and considering the other reviews, i raise my score to Accept.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "`The method applies masking to the token embeddings and arithmetic to the logits, so it needs to be inserted within the LLM model. No discussion is made on how complex that is. The modifications would also be model-dependent. It is not discussed how hard it is to implement these modifications for the chosen models. Are there commonalities that can be abstracted ? Overall, this approach is an integral part of an open-source model code, and cannot just be applied on top of an existing models in an easy way.`\n\nThe implementation is not model-dependent. All six models in our paper are built upon the Huggingface Transformer library, which offers the APIs to directly access and edit token embeddings and logits. So we only need to implement one single decoding method for all these models using these APIs based on the SPA algorithm. Please check Lines 2859-3019 for the implementation of this decoding method in our anonymous repository (https://anonymous.4open.science/r/Selective-Prompt-Anchoring-3693/weighted_utils/weighted_text_utils.py). We currently cannot integrate SPA with closed-source models like GPT-4 since these models only provide a text-based prompt and response API. However, SPA will be applicable once they offer APIs to their logits and token embeddings. \n\n`The hyperparameter that controls the strength of the logit amplification is model and dataset -dependent. It needs to be tuned as it exhibits a clear maxima. This effect is well discussed in appendix J and the authors show that, while the parameter is strongly model-dependent, the method can still be effective with a fixed value across the benchmark datasets. Unfortunately, the authors do not compare directly in a table the fixed approach with the other SOTA methods. This would be desired, as the real question is whether the proposed approach beats SOTA for a deployable model (dataset-independent).`\n\nThank you for the suggestion. We conducted an additional experiment to evaluate a fixed anchoring strength across different benchmarks. Specifically, we got this fixed value by averaging the tuned anchoring strengths across benchmarks for each model. In the table below, we update Table 2 by adding a new condition (SPA-fixed). It shows that SPA is effectively deployable in new scenarios once a reasonable value is set. We will add more discussion about the practical deployment of SPA in Appendix J.3.\n\n| Method          | Improvement |\n|----------------|:-------------:|\n| PASTA          | + 1.2        |\n| Self-Debugging | + 4.2        |\n| Self-Edit      | + 1.8        |\n| Self-Planning  | + 3.6        |\n| ReAct          | + 1.3        |\n| SPA-fixed      | + 5.0        |\n| SPA            | + 7.7        |", "claims_and_evidence": "Yes, the effects of anchoring part of the prompt to avoid dilution are shown clearly on several empirical studies.", "methods_and_evaluation_criteria": "Yes, the evaluation methods are well suited for the application. The code-generation benchmarks used are well-known and effective.", "theoretical_claims": "The mathematics of the method are clearly explained, including finite-difference approximation used to reduce computational complexity.  There are no theoretical proofs for bounds though.", "experimental_designs_or_analyses": "The experimental design is fine, including the evaluation metrics (pass@k). The choice of 6 open-source code LLMs is extensive and representative: CodeGen-Mono-350M (Nijkamp et al.,2023), CodeLlama-7B (Rozi`ere et al., 2024), StarCoder2-15B (Lozhkov et al., 2024), and DeepSeek-Coder-Instruct-1.3B, 6.7B, and 33B (Guo et al., 2024). The choice of SOTA methods to compare to is fine (although i may be unaware of some recent proposed methods).", "relation_to_broader_scientific_literature": "The works on code-generation, attention and logit arithmetic are well cited.", "essential_references_not_discussed": "I'm not aware of essential references not discussed in the paper.", "other_strengths_and_weaknesses": "Strengths:\n- the method is general, model agnostic and can be applied beyond coding tasks. The intuition behind the approach is sound and well explained.\n- the method is fast as it does only adds 60% over the base model and does not used much extra memory. It is significantly faster than the SOTA methods compared against (> 3x).\n- The paper is well written and organized. It is easy to follow.\n- Code is provided on anonymous github\n\nWeaknesses:\n- The method applies masking to the token embeddings and arithmetic to the logits, so it needs to be inserted within the LLM model. No discussion is made on how complex that is. The modifications would also be model-dependent. It is not discussed how hard it is to implement these modifications for the chosen models. Are there commonalities that can be abstracted ? Overall, this approach is an integral part of an open-source model code, and cannot just be applied on top of an existing models in an easy way.\n- The hyperparameter that controls the strength of the logit amplification is model and dataset -dependent. It needs to be tuned as it exhibits a clear maxima. This effect is well discussed in appendix J and the authors show that, while the parameter is strongly model-dependent, the method can still be effective with a fixed value across the benchmark datasets. Unfortunately, the authors do not compare directly in a table the fixed approach with the other SOTA methods. This would be desired, as the real question is whether the proposed approach beats SOTA for a deployable model (dataset-independent).", "other_comments_or_suggestions": "see above.", "questions_for_authors": "no questions.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The reviewers broadly agreed that this work presents an interesting approach that significantly improves the performance of the models it is applied to. The motivation is well-supported by empirical evidence and the writing is clear. While they note that the approach is somewhat invasive and may have potential, unobserved downsides, the gains in performance and thorough evaluation clearly support its net positive value. Given this, we recommend its acceptance.
Decision: Accept (poster)