Thanks for your insightful comments! We reply to each question (Weakness) below.

Question 1 (Weakness 2)

Thanks for introducing Anchor-LLM. While their names are similar, we find that SPA and Anchor-LLM are different approaches that address different challenges. Anchor-LLM aims to compress semantic information into anchor tokens, thus reducing the need for KV cache. They aim to decrease the time and space requirements for LLM decoding, while not compromising performance a lot. By contrast, SPA introduces a mechanism to amplify/reduce the influence of certain tokens in the prompt, thereby controlling the LLM decoding direction. We found amplifying attention over the original prompt can consistently improve code generation performance. Therefore, SPA and Anchor-LLM are not directly comparable. We think they can be integrated together to achieve both computationally efficient and more accurate decoding.

Question 2 (Weakness 4 & 5)

Thank you for your suggestion. We have conducted additional experiments to evaluate SPA on other programming languages via HumanEval-X [1]. Below is the results. We will add the results in paper.

[1] Qinkai Zheng, Xiao Xia, Xu Zou, Yuxiao Dong, Shan Wang, Yufei Xue, Lei Shen, Zihan Wang, Andi Wang, Yang Li, Teng Su, Zhilin Yang, and Jie Tang. 2023. CodeGeeX: A Pre-Trained Model for Code Generation with Multilingual Benchmarking on HumanEval-X. In Proceedings of the 29th ACM SIGKDD Conference on Knowledge Discovery and Data Mining (KDD '23). Association for Computing Machinery, New York, NY, USA, 5673–5684. https://doi.org/10.1145/3580305.3599790

Question 3 (Weakness 1 & 3)

Weakness 1

Thank you for your question. SPA is mainly designed to address the attention dilution issue with respect to the original prompt in code generation task. Our empirical study revealed that the model typically over-attends to recently generated tokens. Consequently, increasing attention to recent self-generated tokens will hinder addressing attention dilution issue.

Furthermore, self-generated tokens can be incorrect. Overly attending to these tokens may lead to error propagation in subsequent steps (Appendix A.7).

Weakness 3

As the model generates different tokens, its attention dynamically changes at each step. Precisely locating the "most informative" tokens at all steps is extremely challenging. A recent study [2] has shown that attention can be overly distributed to the first or special tokens (a phenomenon called "attention sink"). Furthermore, determining how the model distributes its attention to sub-tokens is complex. Overly micro-managing specific tokens can easily lead to poor performance. For instance, if we incorrectly steer the model's attention to the wrong words in just 5% of cases, the final generated code may be incorrect. Therefore, in our approach, we pursue a balanced strategy. SPA anchors the natural language (NL) instruction in the code generation prompt. We chose this method because, although it may slightly reduce precision, the NL instruction remains consistently relevant to all generated code tokens. Thus, negative influence of less relevant tokens in the prompt can be counteracted by most other tokens. We promise to add this discussion in paper.

[2] Xiao, G., Tian, Y., Chen, B., Han, S., & Lewis, M. (2024, April 7). Efficient streaming language models with attention sinks. In The Twelfth International Conference on Learning Representations.

Thank you for your follow-up questions! We are glad to answer them.

(1)

SPA can support any decoding strategy, such as beam search or nucleus sampling. This is because SPA only augments logits of the last layer of the model, and it is irrelevant to the decoding strategy (Section 3.3). In our current experiment, we've used beam search to calculate pass@10 in Table 1. The results demonstrate that SPA can also improve beam search accuracy compared to the original model. We've provided a more detailed discussion of beam search in Appendix A.6. You can find our code implementation for beam search + SPA at Line 3202 in https://anonymous.4open.science/r/Selective-Prompt-Anchoring-3693/weighted_utils/weighted_text_utils.py.

(2)

We've conducted additional experiments on BigCodeBench. Please check the results below. While the absolute improvements aren't as big as in simple benchmarks, the relative improvements remain comparable. For example, although the absolute improvement for CodeGen-Mono-350M is 0.3%, SPA enhances its performance by 27% relative to the original 1.1% performance. This is because SPA only adjusts the attention of the code generation model and therefore still relies on the model's innate capability of code generation. In other words, if a model could solve a task but misses a few tokens or requirements in the prompt, SPA can help with this by adjusting the attention. If a model is very poor and doesn't possess the capability to solve a task, adjusting the model attention won't help much. We promise to include these new results and discussion in the paper.

Thanks for your insightful comments! We reply to each weakness and question below.

Weakness 1

Thanks for the reviewer’s comments. We'd like to emphasize that our approach differs significantly from existing methods. Current methods such as [1, 2] require extensive adaption to existing models to manipulate the model attention. For example, [1] requires an model profiling stage to identify attention headers to be adjusted for improvement. During inference, it recalculates attention distribution for each layer and selected headers. In contrast, SPA only requires tuning 1 hyper-parameter for optimal performance. During inference, SPA can adjust attention by simply computing logits difference. Furthermore, [1] requires user input to steer model attention, while SPA can automatically improve code generation performance by amplifying the influence of original prompt. [2] tunes a feature selection module to redirect the attention to task-relevant features. In contrast, SPA is model-agnostic and applicable to different model architectures. We promise to improve the related work section by including a more comprehensive comparison between SPA and existing approaches.

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

[2] Shi, B., Gai, S., Darrell, T., & Wang, X. (2023, July 11). Toast: Transfer learning via attention steering. arXiv.org. https://arxiv.org/abs/2305.15542

Question 1

We replicate 4 popular LLM-based optimization code generation approaches [3-6] and directly compare SPA against them. We report the average improvements across all experimental models on HumanEval below. SPA not only outperforms these approaches in terms of Pass@1 improvements but also cost significantly less time.

* Self-debugging leverages error messages from test cases, while SPA doesn’t require any test case.

[3] Chen, X., Lin, M., Schärli, N., & Zhou, D. (2023, October 5). Teaching large language models to self-debug. In The Twelfth International Conference on Learning Representations.

[4] Xue Jiang, Yihong Dong, Lecheng Wang, Zheng Fang, Qiwei Shang, Ge Li, Zhi Jin, and Wenpin Jiao. 2024. Self-Planning Code Generation with Large Language Models. ACM Trans. Softw. Eng. Methodol. 33, 7, Article 182 (September 2024), 30 pages. https://doi.org/10.1145/3672456

[5] Yao, S., Zhao, J., Yu, D., Du, N., Shafran, I., Narasimhan, K., & Cao, Y. (2023, March 10). React: Synergizing reasoning and acting in language models. In The Eleventh International Conference on Learning Representations.

[6] Zhang, K., Li, Z., Li, J., Li, G., & Jin, Z. (2023, July). Self-Edit: Fault-Aware Code Editor for Code Generation. In Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics.

Question 2 (Weakness 2)

As the model generates different tokens, its attention dynamically changes at each step. Precisely locating the "most informative" tokens at all steps is extremely challenging. A recent study [7] has shown that attention can be overly distributed to the first or special tokens (a phenomenon called "attention sink"). Furthermore, determining how the model distributes its attention to sub-tokens is complex. Overly micro-managing specific tokens can easily lead to poor performance. For instance, if we incorrectly steer the model's attention to the wrong words in just 5% of cases, the final generated code may be incorrect. Therefore, in our approach, we pursue a balanced strategy. SPA anchors the natural language (NL) instruction in the code generation prompt. We chose this method because, although it may slightly reduce precision, the NL instruction remains consistently relevant to all generated code tokens. Thus, negative influence of less relevant tokens in the prompt can be counteracted by most other tokens. We promise to add this discussion in paper.

In this work, we evaluate anchoring different components in code generation prompt in Section 5.4. We created four experiment baselines by not anchoring test cases and code. Our results indicate that "anchoring NL instruction alone in the prompt" achieves the best performance.

[7] Xiao, G., Tian, Y., Chen, B., Han, S., & Lewis, M. (2024, April 7). Efficient streaming language models with attention sinks. In The Twelfth International Conference on Learning Representations.

Question 3

Thanks for the reviewer’s suggestion. This is an interesting question! We conducted additional experiments to study the number of tasks that SPA compromises initially correct output V.S. number of tasks that SPA rectifies initially incorrect output. Below are the average results for each model on HumanEval.

The overall ratio of "compromised" to "rectified" tasks is approximately 1:2. We find this result interesting. While SPA successfully corrects many incorrect generations, it also compromises some initially correct generated code. This result suggests that the current significant improvements of SPA can be further enhanced. For example, we could use test cases to decide when to trigger SPA (when any test case fails).

To validate this design, we conducted additional experiments testing a setting where SPA is triggered only when the code fails to pass the test cases. The result shows that SPA serves effectively as an error correction approach, significantly improving Pass@1 (%) performance in this scenario. We will include these new results and design details in the paper.

Thanks for your reply. The author's answers to question 2 and weakness 2 are not convincing. I'll maintain the score.

Question 4 (Weakness 2)

We replicate popular prompt optimization-based code generation approaches [2-5] and directly compare SPA against them. We report the average improvements across all experimental models on HumanEval below. SPA not only outperforms these prompt-based approaches in terms of Pass@1 improvements but also cost significantly less time.

* Self-debugging leverages error messages from test cases, while SPA doesn’t require any test case.

Notably, compared to most prompt-optimization approaches, SPA enhances code generation performance at the logits level, which can be easily integrated with them to form a more advanced pipeline. We will add the results and discussion in the paper.

[2] Chen, X., Lin, M., Schärli, N., & Zhou, D. (2023, October 5). Teaching large language models to self-debug. In The Twelfth International Conference on Learning Representations.

[3] Xue Jiang, Yihong Dong, Lecheng Wang, Zheng Fang, Qiwei Shang, Ge Li, Zhi Jin, and Wenpin Jiao. 2024. Self-Planning Code Generation with Large Language Models. ACM Trans. Softw. Eng. Methodol. 33, 7, Article 182 (September 2024), 30 pages. https://doi.org/10.1145/3672456

[4] Yao, S., Zhao, J., Yu, D., Du, N., Shafran, I., Narasimhan, K., & Cao, Y. (2023, March 10). React: Synergizing reasoning and acting in language models. In The Eleventh International Conference on Learning Representations.

[5] Zhang, K., Li, Z., Li, J., Li, G., & Jin, Z. (2023, July). Self-Edit: Fault-Aware Code Editor for Code Generation. In Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics.

Question 5 (Weakness 4)

We appreciate the reviewer's insightful question about prompt length. We divided the HumanEval dataset into three subsets (Short, Medium, and Long) based on the 33rd and 66th percentiles of prompt lengths (Average length is 451 characters). The effectiveness of SPA across these divisions is as follows:

We find the results interesting. While LLMs consistently perform better on short prompts, SPA is consistently more effective when handling longer prompts. This finding confirms that SPA can effectively address attention dilution issue observed in our empirical study. It also suggests that SPA is particularly beneficial when dealing with lengthy prompts. We promise to include these new results and discussion in paper.

Thanks for your insightful comments! We reply to each weakness and question below.

Weakness 1

Thank you for your suggestion. We have conducted additional experiments on HumanEval-X [1] to show the generalizability of SPA to other programming languages. We will add the results in paper.

[1] Qinkai Zheng, Xiao Xia, Xu Zou, Yuxiao Dong, Shan Wang, Yufei Xue, Lei Shen, Zihan Wang, Andi Wang, Yang Li, Teng Su, Zhilin Yang, and Jie Tang. 2023. CodeGeeX: A Pre-Trained Model for Code Generation with Multilingual Benchmarking on HumanEval-X. In Proceedings of the 29th ACM SIGKDD Conference on Knowledge Discovery and Data Mining (KDD '23). Association for Computing Machinery, New York, NY, USA, 5673–5684. https://doi.org/10.1145/3580305.3599790

Question 1

This is because we SPA is tuned based on Pass@1, whose optimal values may slightly differ from the optimal values of Pass@10. We promise to clarify this in the paper.

Question 2 (Weakness 3)

To demonstrate the tuning stability, we evenly split the original dataset into 5 subsets. Below we show the hyper-parameters tuned on each subset as well as the hyper-parameters tuned on the entire dataset. As demonstrated in Figure 4, the tuned hyper-parameters are near the optimal ones. It is stable because the hyperparameter-performance distribution follows a relatively simple unimodal pattern. We promise to better clarify this in paper.

Question 3

We have conducted cross-model experiments in Section 5.2. The results in Table 2 demonstrate that the anchoring strength tuned on one model can be effectively transferred to another.

Furthermore, we've found that setting a universal anchoring strength can enhance performance across all models (as detailed in Section 5.3 and Appendix A.3). For example, when we set the default value to 1.2, we observe consistent performance improvements as shown below. We promise to better clarify this in paper.

Thanks for your insightful comments! We reply to each weakness below.

Weakness 1

Thank you for mentioning these works. They are related but the phenomena described in these works are different from the phenomena in our findings. Specifically, the first work [1] found that LLMs can be distracted by irrelevant information. However, in our work, the code generation prompts do not include any irrelevant information. The function header and natural language instruction are consistently relevant to the generated code. The second work [2] found that LLMs struggle to attend to the middle of long contexts. However, code generation prompts are usually short and appear at the beginning. The third work [3] points out that model attention should align with human intention in human-AI communication. However, code generation prompts directly represent user intentions in our experiments. Our work is the first to confirm the existence of an attention dilution issue in code generation tasks.

It's also worth noting that the first and second papers [1, 2] are empirical studies and didn't propose any technical solutions. The third paper [3] proposes a technical approach that requires user input to steer model attention. In contrast, SPA automatically amplifies the influence of original prompt to address attention dilution issue and doesn't require any user input. Furthermore, it requires extensive model profiling to identify attention headers to be adjusted for improvement. During inference, it needs to recalculate attention distribution for each layer and selected headers. In contrast, SPA doesn't require any model profiling and can directly adjust attention based on logits difference.

We will cite these papers recommended by the reviewer and clarify the differences in the paper.

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

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

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

Weakness 2

Unlike approaches requiring multiple hyperparameter tuning, SPA only needs to tune one hyperparameter (anchoring strength) for optimal performance. This hyperparameter follows a simple pattern (Figure 4) and is easy to tune. In practice, we can set a default value of 1.2, which improves performance across various models and benchmarks. Please see the results in our response to Weakness 4.

Additionally, our approach is model-agnostic. It can be applied to other architectures. In contrast, the existing attention steering method such as [3] requires a model profiling stage, which is not only model-specific but time-consuming in practice.

Weakness 3

Following the reviewer’s suggestion, we've conducted additional experiments on BigCodeBench. Please check the results below. While the absolute improvements aren't as big as in simple benchmarks, the relative improvements remain comparable. For example, although the absolute improvement for CodeGen-Mono-350M is 0.3%, SPA enhances its performance by 27% relative to the original 1.1% performance. This is because SPA only adjusts the attention of the code generation model and therefore still relies on the model's innate capability of code generation. In other words, if a model could solve a task but misses a few tokens or requirements in the prompt, SPA can help with this by adjusting the attention. If a model is very poor and doesn't possess the capability to solve a task, adjusting the model attention won't help much. We promise to include these new results and discussion in the paper.

Weakness 4

Yes! In practice, we can set a default value of 1.2 (Section 5.3 & Appendix A.3), which improves performance across various models and benchmarks. Below we shows Pass@1 rates when the anchoring effect is set to the default 1.2. We will include this result in paper.

Weakness 5

Following the reviewer’s suggestion, we've conducted additional experiments on DeepSeek-Coder-V2 (16B) and StarCoder2 (15B). Please see the results in the table below. It shows can SPA can consistently improve recent code LLMs.

Weakness 6

In practice, we argue that the additional overhead is negligible for developers. On average, SPA took 15.4 seconds to complete a HumanEval task, which is comparable to the original model's 9.6 seconds. During the rebuttal, we've also integrated KV-cache and flash attention to speed up SPA. The overhead is reduced to only 1.6 times that of the original model. To illustrate, we report the decoding speed with and without SPA on our machine.

We will provide the implementation details and more discussion about the inference time in appendix.

Weakness 7

Thank you for your suggestion. In addition to BigCodeBench which includes Multiple-E [5], we have also experimented with HumanEval-X [6] to show the generalizability of our approach to other programming languages. We will add the results in paper.

[5] Cassano, F., Gouwar, J., Nguyen, D., Nguyen, S., Phipps-Costin, L., Pinckney, D., Yee, M.-H., Zi, Y., Anderson, C. J., Feldman, M. Q., Guha, A., Greenberg, M., & Jangda, A. (2022, December 19). MULTIPL-E: A scalable and extensible approach to benchmarking neural code generation. arXiv.org. https://arxiv.org/abs/2208.08227

[6] Qinkai Zheng, Xiao Xia, Xu Zou, Yuxiao Dong, Shan Wang, Yufei Xue, Lei Shen, Zihan Wang, Andi Wang, Yang Li, Teng Su, Zhilin Yang, and Jie Tang. 2023. CodeGeeX: A Pre-Trained Model for Code Generation with Multilingual Benchmarking on HumanEval-X. In Proceedings of the 29th ACM SIGKDD Conference on Knowledge Discovery and Data Mining (KDD '23). Association for Computing Machinery, New York, NY, USA, 5673–5684. https://doi.org/10.1145/3580305.3599790