Selective Prompt Anchoring for Code Generation

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.

The 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.

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).

Thank 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.

I worry that MBPP and HumanEval are in the memorization regime...

Thank you for the suggestion. We conducted a new experiment on LiveCodeBench (10/1/2024-2/1/2025). The results show that SPA remains effective.

It's not clear ... Is this rerunning included in the 1.27x overhead?

The 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!

In 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.

[1] Chen et al. Teaching large language models to self-debug. ICLR 2024.

[2] Zhang et al. Self-Edit: Fault-aware code editor for code generation. ACL 2023.

"For simplicity, we demonstrate by making the entire user prompt x as the anchored text."...

Sorry 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.

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 ω."

Ei(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.

In what way is the attention being steered and not just the final logits?

As 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.

Can you provide intuition for the code domain why the last layer "represent the most accurate attention distribution"?

According 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.

[3] Wan et al. What Do They Capture? A Structural Analysis of Pre-Trained Language Models for Source Code. ICSE 2022.

[4] Kou et al. Do Large Language Models Pay Similar Attention Like Human Programmers When Generating Code? FSE 2024.

It would be good to move how to select omega in the main body.

Thank you for the suggestion. We will move Appendix J.1 to Section 5.

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.

This 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.

inference time is also analyzed, but maybe better to discuss a little bit about potential memory overhead or precomputation requirements

There 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.

Suppose:

• vocabulary_size = 50,000
• token_embedding_dimension = 4,096
• logit_size = 2 bytes

The logit overhead will be 50,000 * 4,096 * 2 ≈ 390 MB.

In 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.

We will discuss this in the paper.

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.

Indeed, 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.

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.

How would SPA perform with more dynamic anchored text selection strategies that change during the generation process?

It 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).

Furthermore, 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.

We will discuss these ideas as future work in the paper.

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)?

SPA 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.

In 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.