Title: Context-Aware Hierarchical Learning: A Two-Step Paradigm towards Safer LLMs
Venue: NeurIPS 2025
Accepted: yes
Keywords: LLM

Large Language Model

Security

Prompt Injection Attack

Instruction Hierarchy
Abstract: Large Language Models (LLMs) have emerged as powerful tools for diverse applications. However, their uniform token processing paradigm introduces critical vulnerabilities in instruction handling, particularly when exposed to adversarial scenarios. In this work, we identify and propose a novel class of vulnerabilities, termed Tool-Completion Attack (TCA), which exploits function-calling mechanisms to subvert model behavior. To evaluate LLM robustness against such threats, we introduce the Tool-Completion benchmark, a comprehensive security assessment framework, which reveals that even state-of-the-art models remain susceptible to TCA, with surprisingly high attack success rates. To address these vulnerabilities, we introduce Context-Aware Hierarchical Learning (CAHL), a sophisticated mechanism that dynamically equilibrates semantic comprehension with role-specific instruction constraints. CAHL leverages the contextual correlations between different instruction segments to establish a robust, context-aware instruction hierarchy. Extensive experiments demonstrate that CAHL significantly enhances LLM robustness against both conventional attacks and the proposed TCA, exhibiting strong generalization capabilities in zero-shot evaluations while still preserving model performance on generic tasks. Our code is available at https://github.com/S2AILab/CAHL.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes a new class of vulnerabilities, Tool-Completion Attack (TCA), which insert malicious instructions in tool calling, making them look semantically legitimate. Then the authors introduce a Tool-Completion benchmark for evaluating LLMs against TCA and demonstrate that sota LLMs are performing poorly on this bench (>90% ASR). Finally, authors introduce a Context-Aware Hierarchical Learning (CAHL) for enforcing instruction hierarchies and reducing model vulnerabilities to TCA and other injection attacks.\n\nSTRENGTHS AND WEAKNESSES:\nOverall, a good paper with a novel approach to instruction hierarchies. I have serious concerns regarding the simplicity of the attack goal and the quality of writing in the introduction section, but considering the novelty of the method and importance of the setting, I would lean towards acceptance.\n\n\nStrengths:\n* The idea behind TCA and benchmark is solid. It is important to evaluate tool completion attacks specifically, as it’s one of the most realistic attack scenarios. It’s also an advantage that the proposed attacks are semantically blending with the rest of the prompt, unlike many security benchmarks. \n* The technical idea of CAHL is novel and interesting. The solution consists of Segment Summarization (a cross attention between token sequence and segment id sequence) which collects relevant segment representation from each token, and Contextual Propagation which exchanges this information between all token positions via self-attention. This is a nice improvement over ISE, and it was explained well by the authors.\n* The experiments are solid. The authors use non-instruct tuned Llama 3 8B to avoid introducing confounders from meta’s instruction and safety tuning. The authors explicitly report training details, and provide a lot of evaluations, including gcg attacks. I didn't spot any issues with the experimental setup, and the paper convinced me CAHL is a valid approach to hierarchies.\n\nWeaknesses:\n* The attack goal of the tool completion benchmark is to output “Hacked”. While this is a common practice for existing benchmarks (e.g., TensorTrust) to contain such attack goals, they are quite brittle in the sense that slight perturbations to the prompt can change ASR significantly. I appreciate the experiment in Appendix A.5., where another keyword was checked without significant changes to ASR, but I think the paper would benefit from including more attack goals in the benchmark.\n* Some sentences in the introduction are really hard to parse, e.g., “employing contextual propagation to dynamically equilibrate semantic comprehension”. As a result, from the introduction, it was very unclear what CAHL does specifically. Also some sentences used complex vocabulary, e.g., “incepting an inveracious belief state”. I got an impression that major parts of the intro were AI generated from the manually written rest of the paper. This was a major negative point for me. The reader should be able to get a good overview of the paper from the introduction without struggling with the language. Considering the good quality of the paper, I still recommend acceptance, but I strongly(!) encourage authors to rewrite the introduction. \n* It would be good to have more models than Llama 3.1 8B.\n* Minor: the paper doesn’t mention which dataset CAHL and baselines are tuned on. I presume it’s cleaned alpaca, since it’s cited but never used in the text, but it’s important to include this explicitly in the text.\n\nQUESTIONS:\nFeedback for camera ready: I’d recommend expanding discussion on the defense methods, and discuss some recent works that appeared after March 1.  “To mitigate the risks of prompt injection, two principal defense paradigms emerge: defensive prompt engineering and model training” is not entirely correct. There are also methods like detection filters, circuit breakers, and a few recent architectural-based approaches (of which you mention ISE). Consider citing [1] (follow-up on StruQ), [2] (follow-up on ISE), [3], [4] and other architectural approaches.\n\n\n[1] Ruiyi Zhang, David Sullivan, Kyle Jackson, Pengtao Xie, and Mei Chen. Defense against Prompt Injection Attacks via Mixture of Encodings. arXiv preprint arXiv:2504.07467, 2025.  \n\n[2] Egor Zverev, Evgenii Kortukov, Alexander Panfilov, Alexandra Volkova, Soroush Tabesh, Sebastian Lapuschkin, Wojciech Samek, and Christoph H. Lampert. ASIDE: Architectural Separation of Instructions and Data in Language Models. arXiv preprint arXiv:2503.10566, 2025.  \n\n[3] Praveen Venkateswaran and Danish Contractor. Spotlight Your Instructions: Instruction-following with Dynamic Attention Steering. arXiv preprint arXiv:2505.09241, 2025.  \n\n[4] Zihao Wang, Yibo Jiang, Jiahao Yu, and Heqing Huang. The Illusion of Role Separation: Hidden Shortcuts in LLM Role Learning (and How to Fix Them). arXiv preprint arXiv:2505.00626, 2025.  \n\n\nSome other questions I am curious about: \n\n\n* Can you please elaborate on the intuition behind the Segment Summarization? We have segment embeddings, queries and key-values from token embedding. Since you restrict attention via M mask to only attend within the segment, you’ll not get information like “how relevant is token x for each of the segments”. You solve it via contextual propagation, but why not allow segment embedding queries to attend to the whole sequence? What would happen if you did that? Would be interesting to see this as an ablation.\n* For StruQ, do you report “Naive Ignore Escape-S Completion-R” results on ID or OOD? If it’s either of them, what are the results on the other split of the dataset? Why didn’t you include them?\n* How diverse is the data generation for the benchmark? I’d appreciate it if you could attach 7-10 randomly sampled attacks from the benchmark to your response.\n* For the future iterations, if you have time, consider evaluating your method on other models,  e.g., Qwen 2.5 7b.\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 4, "ORIGINALITY": 4, "DATE": "2025-06-14"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper identifies a novel vulnerability in LLMs during tool use and introduces a new attack method called the Tool-Completion Attack (TCA). The TCA assumes that the attacker can arbitrarily modify dialogue history and third-party data, but has no access to the model’s parameters, predefined chat templates, training data, or system prompt. The attacker’s objective is to make the model output the phrase “Hacked” by constructing adversarial scenarios with semantically relevant context, followed by a malicious instruction. The authors also propose a benchmark to evaluate the susceptibility of different LLMs to the TCA systematically. To mitigate this attack, they introduce a training-time defense that enhances the model’s ability to detect semantic shifts by maximizing the discriminability of segment-wise semantic features while maintaining deep contextual interactions between user queries and response segments. Experimental results demonstrate that the proposed defense effectively reduces the impact of TCA across various LLMs.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n+ The paper is well-structured and comprehensive, presenting a novel attack alongside a corresponding benchmark and defense method.\n+ The proposed defense demonstrates promising results and shows potential for extension to related threats such as prompt injection attacks.\n\nWeaknesses:\n\n- The paper is kind of ambitious; it attempts to cover a new attack, benchmark, and defense within a single work. As a result, each component lacks sufficient depth and detailed analysis.\n- Important technical details are missing, for example, a clear description of the attack setup, how the proposed benchmark differs from existing ones, how its quality is evaluated, and specifics of the defense implementation.\n\nQUESTIONS:\n1. **Clarification on Attack Details:**\nThe attack methodology is not clearly explained. Specifically, how are the task completion prompts and scene-grounded objects generated? Are the adversarial scenarios transferable across different models, or does each target model require a uniquely crafted context and instruction pair?\n2. **Threat Model Concerns:**\nThe assumed threat model allows the attacker to modify the dialogue history. In this case, why can’t the attacker simply edit the history directly (e.g., “The user changes their mind and would like to change the amount from XXX to YYY”), which could also compromise the tool’s functionality? Would such direct edits be effective, and how would their performance compare?\nFurthermore, beyond triggering a specific output like “HACKED,” could the attack be extended to more practical goals tailored to specific tools—for example, altering the transfer amount in financial tools or injecting offensive language in emails? These fall under the same threat model but represent more realistic and impactful scenarios. Can the authors provide real-world use cases where this threat model is applicable?\n3. **Benchmark Evaluation:**\nHow does the proposed benchmark perform in diverse tool environments, such as email assistants, HR bots, or file management systems? Additionally, how is the quality of the benchmark evaluated, and what are its limitations in capturing the full range of possible attacks?\n4. **Defender’s Assumptions:**\nThe paper does not adequately discuss the threat model from the defender’s perspective. What assumptions are made about the defender’s capabilities and knowledge? Does the defender have access to the attack setup or any knowledge of the adversarial prompts? What is the defender’s access level to the target model?\n\nLIMITATIONS:\nyes.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-06-15"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper explores vulnerabilities in large language models (LLMs) caused by tool invocation mechanisms. It introduces a new vulnerability, the Tool-Completion Attack (TCA), which exploits the function-calling mechanism of LLMs to manipulate their behavior. The authors propose a novel defense mechanism called Context-Aware Hierarchical Learning (CAHL) to address these vulnerabilities. They also introduce a comprehensive benchmark to evaluate LLM robustness against such attacks.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\nS1: The identification of Tool-Completion Attack (TCA) is a significant contribution. It addresses a newly discovered threat and proposes a method to assess LLM vulnerabilities in real-world scenarios.\n\nS2: The work is timely, as it tackles the growing concerns regarding the security and integrity of LLMs, offering concrete strategies to enhance model safety in adversarial settings.\n\nS3: The paper includes a comprehensive evaluation using two different benchmarks (StruQ and Tool-Completion), with results showing significant performance improvement over existing methods.\n\nWeaknesses:\nW1: While CAHL performs excellently in trained models, the paper could benefit from a deeper analysis of its generalization capabilities in true zero-shot scenarios, where the model has never encountered a specific attack, i.e., the analysis of zero-shot performance is somewhat lacking.\n\nW2: Focusing on a single attack type may limit the generalizability of the findings. A broader range of attack vectors could further strengthen the paper’s relevance.\n\nW3: Some experimental results lack error bars or statistical significance metrics, reducing confidence in the stability of these results under different conditions.\n\nW4: The descriptions of CAHL and TCA are generally clear, but some key aspects of the methodology, particularly regarding Contextual Propagation and Segment Summarization, lack intuitive examples and detailed explanations. These concepts are somewhat abstract, and adding concrete examples and visualizations would help readers understand them more easily. Additionally, the interpretation of formulas and steps needs to be more detailed to aid in understanding the specific role of each part.\n\nAdditional Comments:\nFor the algorithm description, particularly for Semantic Propagation and Contextual Summarization, the current explanation may be too abstract for some readers. A more direct explanation of these concepts is recommended.\n\nIn Section 4, the term “Instructional Segment Embedding (ISE)” is used to describe the instruction hierarchy, while Section 5 and Section 6 refer to it as “instruction hierarchy”. Although related, it is recommended to standardize the terminology throughout the paper to improve consistency and readability.\n\nQUESTIONS:\nN/A\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 4, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work introduces a Completion-style Prompt Injection Attack, termed a Tool-Completion Attack which works by distracting the LLM from the result obtained from some Tool Call, thus encouraging it to naively execute subsequent malicious instructions. This appears to be a potent type of jailbreak, as evidenced by the results from Table 1. To address this vulnerability, the authors contribute both a benchmark and a method, CAHL (Context-Aware Hierarchical Learning), to measure and reduce ASRs associated with this type of attack.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n\n1. The results from Table 1 and Figure 2 effectively motivate the importance of this problem, by (1) demonstrating the susceptibility of recently released, highly-capable models to the described TCA and (2) highlighting the insufficiencies of current adversarial prompt classification techniques.\n\n2. The figures are clear and accessible.\n\n3. The threat model is well-posed.\n\nWeaknesses\n\n1. This type of attack repackages existing completion-style injections and the architectural change is a standard segment-level attention pattern layered on top of ISE. This makes the method contribution only incrementally novel.\n\n2. While the benchmark is a welcome contribution, it appears to already be saturated by the proposed method, somewhat limiting its usefulness in future study. Specifically, on the Adversarially trained split in Table 3, CAHL already attains an ASR between 1-2%. Future methods may struggle to show measurable gains.\n\n3. The proposed benchmark also does not include triggers like paraphrased payloads/instructions, JSON-schema tool outputs, or long contexts, which limits the realism of the evaluation as well as claims of practical robustness.\n\n4. In Appendix Section 3, CAHL appears to reduce ASR only very slightly (1-2%) against GCG attacks, which undercuts claims that CAHL is properly robust across attacks.\n\n5. Additional analysis of the results in Table 2 would help frame the contribution better. For example, why is it that CAHL reduces ASR on the Completion-R split so significantly compared to other baselines, while performance on the remaining splits across all the methods is more tightly clustered? Similarly, why does CAHL display a consistent Capability improvement over ISE?\n\n6. Some method details are unclear: (a) The paper doesn’t state the loss used to train the new attention modules. (b) The parameter count of the new modules is not shared. (c) The method appears to be trained on a single seed.\n\n7. Some portions of the presentation are unclear: (a) In both Figures 5 and 6, which token positions are associated with which parts of the input (i.e. system, user, tool, etc.)? (b) Some of the diction chosen to describe the method in the Introduction impedes understanding.\n\nQUESTIONS:\n1. Do the authors also merge system messages and user queries when implementing the ISE baseline? If so, this might inflate CAHL’s gains as it partially undermines ISE’s methodology of discriminating between system, user, and data content.\n\n2. How would the distribution of risk scores change if Prompt-Guard was finetuned on a portion of the Tool-Completion Benchmark?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 1, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper identifies a new prompt injection vulnerability in Large Language Models (LLMs), termed Tool-Completion Attack (TCA). This attack exploits the reliance of LLMs on external tool outputs to inject malicious instructions under the guise of legitimate follow-up. The authors also introduce a Tool-Completion Benchmark to systematically assess LLM robustness under TCA, showing that even state-of-the-art models like GPT-4o are highly vulnerable. To mitigate this, they propose a training method called Context-Aware Hierarchical Learning (CAHL), which structures input tokens according to instruction roles and uses contextual attention to model hierarchy.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n\nNovel vulnerability: The identification and formalization of TCA is a meaningful contribution to the growing literature on LLM safety.\nThorough evaluation: The authors test TCA on a wide range of models (GPT-4o, LLaMA-3.1, DeepSeek-R1, etc.), demonstrating high success rates of the attack (up to 99%).\n\nConcrete mitigation: CAHL is clearly described, with an architecture that builds on existing embedding strategies (e.g., ISE) and adds context propagation.\n\nEmpirical gains: CAHL consistently improves robustness across multiple benchmarks (StruQ and Tool-Completion), reducing attack success rate (ASR) while maintaining or slightly improving general capabilities.\n\nAblation and visualization: The paper includes thorough ablation studies and t-SNE visualizations to justify the benefits of the proposed approach.\n\nWeaknesses:\n\nLimited generalization testing: Experiments are mostly confined to supervised fine-tuning on narrow tasks and benchmarks. No exploration of CAHL’s behavior in open-ended real-world usage or RLHF-based pipelines.\n\nSimplistic attack metric: The main evaluation metric (presence of the keyword “Hacked”) is overly simplistic and may not fully capture nuanced model behavior or partial success/failure cases.\n\nFocus on a single attack vector: While TCA is shown to be dangerous, it is just one vector of prompt injection. The paper does not address how CAHL performs on other types of adversarial threats beyond what is benchmarked.\n\nUnclear deployment cost: The compute and integration overhead of CAHL in real-world LLM pipelines (e.g., OpenAI-style systems) is not assessed.\n\nQUESTIONS:\nQuestion 1: The reliance on exact string matching (e.g., keyword “Hacked”) to evaluate attack success is simplistic and may understate or overstate attack severity. Consider including qualitative examples, human evaluations, or more nuanced metrics (e.g., intent fulfillment, partial compliance) to capture model behavior more realistically.\n\nQuestion 2: The additional complexity introduced by CAHL (segment embeddings, contextual propagation) raises questions about computational and implementation overhead. Please provide guidance or empirical evidence on the training/inference cost and ease of integration into real-world LLM pipelines.\n\nLIMITATIONS:\nYES", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-23"}