BERM: Low-Overhead Prompt-Injection Detection via In-Situ Benign Representation Modeling

1 Introduction

Large language models (LLMs) power mission-critical services such as chatbots, code assistants and search engines, answering billions of queries every day. Yet their remarkable flexibility also exposes a rapidly evolving threat: prompt injection. In these exploits, attackers cleverly embed harmful directives within apparently benign prompts, tricking models into revealing private information or bypassing built-in safety restrictions. According to the OWASP, prompt injection represents the most critical security vulnerability for LLM deployments[OWASP, 2023]. Current defense strategies typically involves a difficult trade-off between robustness and model performance, a central dilemma highlighted in recent security standards [Vassilev et al., 2025].

∗Corresponding author.

The overhead-robustness dilemma. Security teams implementing LLM protections routinely face a wrenching tradeoff: strengthening defenses inevitably slows down response times. Lightweight keyword filters are fast but crumble under sophisticated attacks; meanwhile, robust modelbased guards-typically DeBERTa classifiers [He et al., 2021], follow an external-guard paradigm that adds ≥ 16 ms latency and is often limited to a 512-token context window. These technical limitations critically undermine their real-world deployability where user experience demands minimal response times and prompts routinely exceed these context windows.

Our approach: In-Situ Benign Representation Modeling (BERM). We address this dilemma with BERM, a novel framework that casts prompt injection as an out-ofdistribution shift in the host LLM's own representation space. By harnessing the hidden-state representations extracted insitu during prefill, BERM bypasses the need for external guards. In their raw form, these representations are geometrically entangled (AUROC ≈ 0.59). The learned projector then re-embeds them in a more compact manifold where benign and injected prompts become separable. Figure 1 depicts BERM pushing the latency-performance frontier outward relative to prior guard models.

BERM leverages a joint contrastive learning scheme fusing (i) an unsupervised, cluster-based objective modeling the benign manifold and (ii) a supervised objective separating known attacks, allowing BERM to learn a robust notion of normalcy from as few as 300 injections.

Contributions.

1. Framework. BERM pairs strong prompt-injection detection with near-zero overhead. It improves overall F1 by +5.2 pp and cuts inference latency by ≥ 12×.
2. Learning objective. We propose a joint contrastive loss tailored for prompt-injection detection, achieving high sample efficiency by fusing unsupervised benign modeling with supervised ID-OOD separation.
3. Analysis tool. We develop the Whitened Intrinsic Distance Estimator (WIDE), a post-hoc diagnostic that audits representation quality and benign-injection overlap.

Road-map. Section 3 details BERM and WIDE; Section 4 reports experiments on four corpora and three backbones; Section 5 discusses limitations and future work.

2 RelatedWork

Prompt Injection Attacks. Prompt injection plants adversarial text into an LLM's context, prompting the model to override, ignore or leak the original prompt [Perez and Ribeiro, 2022]. Direct attacks are authored by the user [Willison, 2023]; indirect attacks can be embedded in third-party content fetched by a tool or search engine, taking effect once that content is appended to the prompt [Greshake et al., 2023]. Techniques include fake completions [Willison, 2023], multi-strategy chains [Yi et al., 2025] and gradientsearched suffixes [Liu et al., 2024]. Recent large-scale evaluations find that even state-of-the-art models can be jailbroken reliably-often with near-perfect success-and that such prompts transfer across model families, highlighting the seriousness of the threat [Shen et al., 2024].

Prompt Injection Defense. Existing defenses can be grouped into three lines of work: (i) Architecture-level fixes separate control from data, e.g., StruQ's dual channels [Chen et al., 2025], instruction hierarchies [Wallace et al., 2024], or task-specific de-instruction tuning [Piet et al., 2024]. These require weight access and compute-intensive retraining, and may constrain free-form usage. (ii) Selfsupervision adds meta-queries or reminder tokens so the model audits itself [Phute et al., 2024; Xie et al., 2023], straightforward to integrate but consumes context and lengthens latency. Hybrid pipelines that cascade detectors with architectural separation often inherit both cost and brittleness [Suo, 2024]. (iii) Detection filters route each prompt through a smaller surrogate classifier [Blueteam AI, 2024; ProtectAI.com., 2024; Shang et al., 2023] or a second LLM

in zero-shot mode [Meta, 2024]. The low-overhead runner-up of this paradigm-PIGuard [Li et al., 2025],relies on markedly more training data.

Out of Distribution (OOD) Detection. Our work frames prompt-injection detection as a fine-grained case of Outof-Distribution (OOD) detection. Foundational workssupervised [Hendrycks et al., 2019; Hendrycks et al., 2020] and unsupervised like ContraOOD [Zhou et al., 2021] achieved strong OOD detection by outlier exposure and contrastive learning to separate coarse-grained domain shifts (e.g., news vs. movie reviews). However, they are less effective for prompt injection. The distinction between benign and injected inputs is fine-grained and often hinges on adversarial phrasing within an otherwise in-distribution context. Within highly homogeneous domains like mathematical reasoning, Wang et al. [2025] leverage the volatility of this layer-wise hidden-state trajectory as a powerful OOD signal. This finding corroborates that a LLM's internal processing dynamics encode crucial cues for distinguishing in- from outof-distribution inputs.

3 Methodology

We propose BERM (Benign Representation Modeling), a framework that detects prompt-injection attacks by modeling the in-situ hidden state representations of a frozen large language model. Prompt injections can steer a model away from its intended instruction-following behavior. Our core intuition is that this deviation manifests as a geometric signature in the in-situ representation space, making it distinguishable from benign queries. BERM consists of an offline representation-learning stage and an online lightweight scoring stage. The methodology is built on three pillars:

(1) benign-manifold discovery via clustering, (2) joint contrastive learning that compactly encodes the benign manifold, and (3) whitened-distance auditing by Whitened Intrinsic Distance Estimator.

3.1 BERMBackbone

Let fθ denote a frozen host LLM. For a prompt x, we take the token-level hidden states h(x) ∈ RT×D from the final layer. The final layer is directly upstream of next-token prediction, since its hidden states feed the language-model head used for token generation [Ethayarajh, 2019]. For a prompt x of length T, we compute a prompt-level embedding by mean-pooling token-level hidden states, following common sentence-embedding practices [Reimers and Gurevych, 2019; Gao et al., 2021]:

This D-dimensional vector is the sole input to the BERM projector.

3.2 BenignManifold Discovery

GMM Clustering. Given benign prompts Did = {xi}Ni=1 , we extract their embeddings {˜zi} and model them with a Gaussian mixture of C components. Each sample produces a soft assignment p(c | ˜zi); we collapse it to the MAP label yclus i = argmaxc p(c|˜zi).

3.3 Joint Contrastive Learning

Training only on labeled injection examples tends to overfit to dataset-specific attack templates. In contrast, generic unsupervised anomaly detection can miss fine-grained malicious instructions embedded in otherwise benign prompts. We address this gap with a joint objective. The unsupervised cluster-contrast loss (Lunsup) tightens the benign manifold to support outlier detection, while the supervised ID-OOD contrast (Lsup) repels known attacks from benign clusters.

Projector Network. A three-layer residual MLP gϕ : RD → RD maps

Unsupervised Cluster Contrast. An instance of the InfoNCE loss [van den Oord et al., 2018], minimises

for mini-batch B ⊆ Did,with temperature τ = 0.15.

Supervised ID-OOD Contrast. With a labeled set Dood = {(x, ℓ)}, ℓ ∈ {0, 1}, we optimize a NT-Xent objective [Chen et al., 2020], treating all benign embeddings as positives and all malicious ones as negatives.

Binary Classification Head. A linear layer qψ : RD → R2 returns logits oi; we minimize cross-entropy

Joint Objective. Training minimizes a weighted combination of the three terms:

Unless noted otherwise, we set α:β:θ = 1:3:1.

3.4 Inference Path and Latency

During LLM prefill we extract the in-situ prompt representation ˜ztest and score it with:

With D=5120 the projector imposes ≤0.06G FLOPs overhead, delivering a 12× latency reduction compared to the leanest guard baselines (Table 1)

3.5 WIDE: Whitened Intrinsic Distance Estimator

To assess the representation quality of BERM and other prompt-injection detectors, we introduce WIDE (Whitened Intrinsic Distance Estimator), a lightweight post-hoc analysis tool. WIDE takes a set of embeddings and summarizes how tightly benign samples cluster by measuring distances to a fitted benign manifold in a whitened space. Rather than acting as a detector by itself, WIDE serves as a diagnostic: it quantifies separability and generalization capacity of the representation space itself. Given benign (in-distribution) embeddings Zid = {zi}Ni=1 ⊂ RD , WIDE assigns each test embedding ztest an anomaly score based on its whitened distance to the benign manifold.

Step 1: Benign Manifold Statistics. We start by summarizing the benign embeddings with their second-order statistics. Specifically, we estimate the sample mean

and the sample covariance matrix Σ, which describes the manifold's shape, orientation, and variance along its principal axes.

Step 2: Denoising via Principal Component Analysis. High-dimensional embeddings can contain noisy or uninformative directions. We eigendecompose the covariance matrix, ΣQ = QΛ, where Q contains the eigenvectors and Λ is the diagonal matrix of corresponding eigenvalues. We retain the top-k principal components that together explain a ratio η (e.g., 99.5%) of the total variance, thereby projecting the data onto a lower-dimensional denoised subspace prior to whitening.

This transformation rotates a deviation vector into the principal component space and re-scales it, amplifying deviations along compact, low-variance dimensions of the benign manifold while suppressing them along high-variance ones. The parameter β ∈ (0, 1] controls this suppression strength, while ε > 0 ensures numerical stability.

Step 4: Anomaly Scoring. For any test embedding ztest, we compute the final anomaly score as its Mahalanobis-like distance in this transformed space:

A large score s intuitively indicates that ztest is improbably far from the benign manifold, particularly along dimensions where benign samples show little variation.

WIDE as a Diagnostic Tool. To translate geometric insights into actionable diagnostics, we derive two metrics from the anomaly score distributions shown in Fig. 3: WhitenAUROC and Whiten-FPR95.

• Ideal Representation (Fig. 3a): Benign and malicious score distributions are clearly separated (WhitenAUROC ≈ 1.0). This signifies a high-quality embedding space where the model has learned a compact and generalizable manifold for benign prompts. As shown in Table 1, our BERMmodel trained on the full Safe Guard Prompt Injection dataset achieves this on in-distribution data (Whiten-AUROC=0.999).

• Overfitting Diagnosis (Fig. 3b): A key diagnostic is the divergence between in-distribution and out-ofdistribution test performance. If the Whitened distances for the benign training set are low, yet the distance distribution of the benign test set largely overlaps with that of the injection test set, this provides a strong quantitative signal of overfitting.
• Predicting Performance Bottlenecks: WIDE metrics expose geometric constraints in representation spaces that shape classifier capabilities. As Fig. 3c demonstrates, collapsed Whiten-AUROC values reveal substantial overlap between benign and malicious patterns, exposing that effective detection requires first untangling these representations at the embedding level. This geometric diagnosis uncovers which representation spaces preserve sufficient discriminative integrity to support reliable security decisions.

Beyond BERM. WIDE transcends its origins as a BERM analysis companion and functions as a general geometric probe for representation spaces. By calculating whitened distances from test samples to an authentic in-distribution manifold, WIDE exposes the representation quality of any embedding architecture: (a) WIDE distills manifold geometry into a singular diagnostic metric whose values presage the feasibility of linear decision boundaries; (b) per-sample distance visualizations could expose manifold collisions and potential representation flaws.

3.6 Complexity

WIDE incurs an offline cost of O(ND2 + D3 ) operations with O(D2 ) memory footprint. During deployment, online inference incurs an O(kD) overhead, which is submillisecond on GPU for D ≤ 5120 in our setting. This computational frugality helps BERM maintains its real-time security posture without compromising system responsiveness.

3.7 Summary of Methodology

BERM optimizes a joint contrastive objective to sculpt a compact manifold of benign prompts. At inference time, each new prompt is scored with just one extra pass through a small projector head computation, adding only a tiny incremental cost. WIDE provides an geometric readout of how distinctly benign and injected prompts diverge, so we can assess representation quality decoupled from the classifier's decision logic.

4 Experiments

Our empirical analysis rigorously scrutinizes BERM across the robustness-efficiency spectrum. The goal is to quantify how much robustness can be achieved when both incremental compute and labeled injections are constrained. We therefore organized the evaluation to probe: (1) How does BERM compare with state-of-the-art guard models across diverse benchmarks? (2) How sample-efficient is BERM in low-data regimes? (3) What is the contribution of each component in our joint contrastive objective?

4.1 Experimental Setup

Datasets. Our evaluation draws on four public corpora that are commonly used in recent prompt-injection studies: SafeGuard (SG) [Shang et al., 2023], Jailbreak Classification (JC) [Shen et al., 2024], DeepSet Prompt Injections (DS) [Deepset, 2024], and Prompt-Injection-Safety (PS) [Jayavibhavnk, 2024]. SG and JC supply labelled training splits; DS and PS serve as test partitions and are solely used for outof-distribution (OOD) evaluation.

Protocol rationale. PIGuard classifier [2025], the latest over-defense mitigating DeBERTa, is trained on a corpus with 61 k benign + 15 k injections, and evaluated partially on private corpus, together contains >1.5 k non-English samples. We limit the study to publicly available, English-only corpora to keep all methods on equal footing and to focus on the Overhead-Robustness dilemma. We train BERM on SG (6k) to fairly compare with baselines trained on thousands of injections; BERM@300 on a random 300-sample subset to mirror low-label settings (§4.2).

Baselines. Five DeBERTa-based guards: Fmops [2024], SafeGuard [2023], ProtectAI-v2 [2024], Prompt-Guard-2 [2024],PIGuard [2025].

Metrics. Accuracy (A), Precision (P), Recall (R), F1, AUROC, FPR95, GFLOPs, and latency (ms). All results are averaged over five seeds (2022-2026).

Implementation. We instantiate BERM on three frozen backbones: Mistral-7B-v0.1 [Jiang et al., 2023], Llama-3.18B-Instruct [Meta AI, 2024], L1-32B [Wan et al., 2025]. 1 We extract in-situ hidden-state representations during the host LLM's native prefill stage, using the final-layer hidden states for their accessibility in most LLM inference stacks.

We use a fixed set of hyperparameters rather than exhaustive tuning: For the Bayesian GMM, we set the Dirichlet prior to 0.1, and run EM for 1000 steps, using 2 clusters for BERM@300, 5 clusters for BERM(@6k); Adam optimizer learning rate 5 × 10-3 , batch size 256; ResMLP hidden size 4096; InfoNCE τ = 0.15. We train BERM for 5 epochs with Adam. Training runs on two NVIDIA RTX-A6000 (48 GB) GPUs. Maximum wall-time for 5K prompts is 30 minutes.

4.2 Main Results

Table 1 benchmarks BERM with five premier stand-alone guard models. We instantiate BERM on three host LLMs: Mistral-7B-v0.1 [2023], Llama-3.1-8B-Instruct [2024], and L1-32B [2025], to validate the universality of benignrepresentation modeling across model families and scales.

Accuracy and OOD robustness. On the four-corpus aggregate, BERM (8B) reaches an Overall F1 of 86.29%, outperforming the previous best PIGuard by +5.2pp. The margin persists on the Unforeseen split (78.18% F1), suggesting that BERM does more than track dataset-specific templates and can generalize to unseen corpora. Scaling the host to L132B still delivers 82.89% Overall and 77.45% Unforeseen F1, indicating that BERM transfers to substantially larger (RL-tuned) backbones.

1 Code available at https://github.com/CytAI/BERM

Low and scalable overhead. Guard baselines require ≈ 60GFLOPs and 16ms per prompt, whereas BERM adds only 0.05 GFLOPs and 0.35 ms, a 400-1200× reduction in (incremental) compute and a 12-47× reduction in latency. The projector cost is governed by the hidden dimension (D≤ 5120) rather than host parameter count. The same overhead applies to BERM (32B), keeping throughput essentially unchanged as the host scales up.

Sample efficiency. The few-shot variant BERM@300 uses only 300 injected prompts and still secures 83.58% Overall F1 on L1-32B, surpassing PIGuard, which relies on 15k injections. This mirrors the trend in Fig. 4 and supports the view that explicitly modeling benign structure can reduce the amount of injection samples needed (see §4.3).

4.3 Few-Shot Sample Efficiency

A cardinal advantage of our joint-learning framework is its capacity to exploit scarce malicious data, a common constraint in production environments. To investigate this, we evaluate BERM's performance when trained with a varying number of malicious samples, ranging from just 100 up to the full 6k set. We probe this by sweeping the injection sample size from 100 to 6k. Figure 4 plots the F1 trajectories for both in-distribution (ID) and out-of-distribution (OOD) test sets.

The results underscore BERM's acute sample efficiency. The in-distribution performance (solid line) demonstrates a steep learning curve, clearing 94% F1-score with only 100 samples and steadily improving thereafter. Crucially, the OOD performance (dashed line), a proxy for generalization to unforeseen attacks, initializes at a remarkably high baseline of nearly 80% F1-score. This indicates that our core approach of modeling the benign manifold confers potent intrinsic robustness, even before extensive exposure to attack patterns.

As the number of malicious samples increases, we observe a slight, transient dip in OOD performance. This is a familiar trade-off: with more labeled attacks, the classifier can lean toward patterns specific to the seen injections, which slightly tightens the decision rule around the training distribution and yields a mild overfitting effect. However, the performance quickly stabilizes to a plateau, consistent with the joint contrastive objective acting as a regularizer that preserves crossdataset generalization.

Our BERM@300 model, trained on just 300 injection samples, already achieves 82.3% F1-score overall and 78.2% on unforeseen datasets (Table 1). This notably represents more than 95% of the model's eventual OOD performance and exceeds PIGuard [2025], which is trained with 15k injections. This few-shot experiment suggest that BERM can deliver strong detection performance under tight labeling budgets, while retaining its low runtime cost.

4.4 Out-of-Distribution (OOD) Robustness

We stress-test the model's adaptability via a leave-onedataset-out regime, training on a solitary dataset and probing performance on the unforeseens.

Protocol. Under this protocol, a detector is fit on a single corpus and evaluated on (i) on its in-distribution split (IN) and (ii) on all OUT-of-distribution corpora held out from training. We report the same metrics as in §4.1: WIDE-AUROC (AUROC), WIDE-FPR95 (FPR95), and the downstream F1. Higher AUROC/F1 and lower FPR95 indicate stronger robustness. We evaluate BERM (Llama-3.1-8B) against the corpus-specific detectors released with each dataset, ensuring identical training splits.

Findings. BERM exhibits a consistent lead over datasetspecific baselines on held-out corpora, regardless of the training source. Trained on SG, BERM boosts F1 by +23.1 pp on average (+38.5 pp on DS), while lowering FPR95. We see the same overall trend when training on JC, even though AUROC on DS drops slightly, BERM still yields +18.7 pp higher F1 in aggregate. Such consistent gains validate that BERM's joint objective captures signals that transfer across datasets rather than keying on corpus-specific artifacts.

4.5 Ablation Study

To dissect the distinct contribution of each module within the BERM framework, we conduct a rigorous four-level ablation study on our best-performing setting with Llama-3.18B-Instruct . Table 3 reports the results. For variants with the final classifier ablated, F1 is no longer well-defined; in those cases we rely on WIDE (see Section 3.5) as a consistent and geometrical probe. This allows us to standardize evaluation across all variants using AUROC and FPR95 derived from whitened distances. The study systematically dissect BERM to make the contribution of each component explicit.

The Foundation: Raw Representations. We begin with the most basic variant, (c) Raw LLM representation, where we apply WIDE directly to the host LLM's raw hidden states without any BERM components. On this raw space, WIDE

yields an AUROC of just 0.590, barely above random chance. In other words, while the vanilla representation geometry may carry some signal, it does not separate benign and injected prompts in a way that makes detection straightforward.

First Layer of Gain: Unsupervised Clustering. Our first addition, (b) GMM only, applies Gaussian Mixture Model clustering to the raw representations prior to scoring. This simple, unsupervised step of identifying and separating submanifolds within the benign distribution provides a notable performance lift. The AUROC improves by +0.132 to 0.722. This suggests that benign prompts exhibit heterogeneous structure in the representation space, and modeling this structure with clusters can improve separability.

The Core Contribution: joint contrastive learning. The most significant performance leap occurs when we introduce our core innovation, the joint contrastive learning projector. In configuration (a) w/o Classifier, the projector is optimized using both the unsupervised cluster-contrastive and supervised ID-OOD contrastive terms. This step dramatically transforms the representation space, catapulting the AUROC from 0.722 to 0.968-a gain of +0.246. This demonstrates that the contrastive learning objective is the dominant factor in BERM's success: it actively sculpts the embedding space, projecting the entangled raw representations into a highly structured manifold where the distance between benign and malicious clusters is maximised.

Final Polish: The Classification Head. Finally, the Full BERM model adds the binary classification head supervised by a cross-entropy loss. While the primary goal of this component is to provide a direct classification output aligned with the prompt-injection task, its training signal also sharpens the learned representations. This brings the AUROC to a nearperfect 0.999 and the FPR95 to 0.000, representing a final gain of +0.031 in AUROC over configuration (a).

Overall Impact. Starting from the nearly inseparable raw representations of a LLM (AUROC 0.590), BERM's components progressively structure and separate the representation space. The final framework achieves an overall improvement of +0.409 in AUROC and -0.851 in FPR95, projecting seemingly elusive LLM representations into separable, robust detection signals. Across these ablations, performance degrades monotonically as we remove components, underscoring that each stage contributes to the final separability observed in Table 3.

5 Discussion

Robust AI safety defenses can emerge from learning better representations. The large performance gaps in our ablation study (Table 3) show that BERM's gain comes primarily from learning a contrastive projector that re-embeds the LLM representations into a space where benign manifolds and malicious prompts become easier to separate. Besides being robust and symbiotically efficient, BERM integrates cleanly into existing inference stacks. It reuses hidden states already produced during prefill and avoids an auxiliary guard model call. Our future work will expand evaluation to broader realworld settings to validate BERM's effectiveness in these scenarios. We are committed to accelerate the adoption of efficient, representation-based defenses in the broader academic and industrial communities.

6 Conclusion

In this work we present BERM, a light-weight framework that detects prompt-injection attacks by modeling the intrinsic geometry of a host LLM's benign in-situ representations. BERM achieves state-of-the-art accuracy and out-ofdistribution robustness while reducing computational cost to a level suitable for latency-sensitive deployments. Looking ahead, these findings suggest that in-situ representation modeling is a promising building block for LLM safety, especially as models and attack surfaces continue to evolve.

Acknowledgments

This work was supported by the National Natural Science Foundation of China under Grant No. 62276260 and the National Key Research and Development Program of China (No. 2024YFE0210600).