PIShield: Detecting Prompt Injection Attacks via Intrinsic LLM Features

TL;DR

Prompt injection poses a critical risk to LLM-integrated applications. PIShield exploits intrinsic residual-stream features from injection-critical layers in Transformer-based LLMs and trains a simple linear classifier to detect contaminated prompts, achieving strong detection while incurring minimal overhead. Across five diverse benchmarks and eight attacks, PIShield outperforms 11 baselines and remains robust under strong adaptive threats. Its efficiency and generalizability suggest practical applicability for real-world defenses against prompt injection, with potential extensions to multi-modal settings.

Abstract

LLM-integrated applications are vulnerable to prompt injection attacks, where an attacker contaminates the input to inject malicious prompts, causing the LLM to follow the attacker's intent instead of the original user's. Existing prompt injection detection methods often have sub-optimal performance and/or high computational overhead. In this work, we propose PIShield, a detection method that is both effective and efficient. Our key observation is that the internal representation of the final token in a prompt-extracted from a specific layer of the LLM, which we term the injection-critical layer-captures distinguishing features between clean and contaminated prompts. Leveraging this insight, we train a simple linear classifier on these internal representations using a labeled set of clean and contaminated prompts. We compare PIShield against 11 baselines across 5 diverse benchmark datasets and 8 prompt injection attacks. The results demonstrate that PIShield is both highly effective and efficient, substantially outperforming existing methods. Additionally, we show that PIShield resists strong adaptive attacks.

Figures (11)

• Figure 1: Overview of LLM-integrated applications.
• Figure 2: Overview of the transformer from the residual stream vector perspective. $E$ denotes the embedding layer. $O$ denotes the output layer. $r^i_j$ represents the residual stream vector for token $p^i$ from layer $j$.
• Figure 3: PCA analysis for residual stream vectors of target instruction with clean (blue color) and contaminated (red color) target data for different LLMs. We randomly sample 1,000 clean and 1,000 contaminated prompts from the training data. The first (or second) row corresponds to injection-critical (or non-injection-critical) layers.
• Figure 4: Effectiveness of PIShield using residual stream vectors from different layers, where FPR and FNR are averaged over 5 datasets and 8 attacks. We find that the injection-critical layers consistently exist for different LLMs. As the results show, the residual stream vector for these layers can be used to effectively detect prompt injection attacks.
• Figure 5: Effectiveness of PIShield using residual stream vectors from different layers for different prompt injection attacks. The FPR and FNR are averaged over 5 datasets. The LLM is Llama-3-8B-Instruct. As the results show, the injection-critical layers are consistently effective in detecting different prompt injection attacks.