Skip to main content

Ensemble Prompt Injection Detection

Project description

PromptForest: Fast, Calibrated & Reliable Injection Detection

PyPI Downloads Apache License Open In Colab

This project is actively in development! If you have any questions or want to help out, you can reach out via issues or discussions, or send me an email.

PromptForest is an ensemble-based prompt injection detection system designed for high-throughput production environments. By aggregating predictions from multiple lightweight expert models, it achieves best-in-class calibration and competitive accuracy while maintaining significantly lower latency than comparable monolithic models.


1. Introduction

Large Language Models (LLMs) are vulnerable to prompt injection attacks where adversarial inputs manipulate the model into executing unintended commands. Traditional defenses often rely on single, large BERT-based classifiers which suffer from two key limitations:

  1. High Latency: Large transformer models introduce significant overhead in critical paths.
  2. Overconfidence: Single models tend to be overconfident in their predictions, even when incorrect (high Expected Calibration Error).

PromptForest addresses these issues through Ensemble Diversity. We combine the signals from diverse architectures (DeBERTa, XGBoost, etc.) to produce a discrepancy-weighted risk score. This allows for:

  • Robustness: Disagreement between models flags ambiguous inputs.
  • Efficiency: The total parameter count is <50% of the leading SOTA monolithic detector.
  • Calibration: Lower confidence on incorrect predictions reduces catastrophic failures.

2. Methodology

PromptForest utilizes a voting ensemble of three distinct lightweight models:

  1. Meta Llama Prompt Guard 86M: Meta's prompt injection detector, already very calibrated for its parameter count (ECE lowest for its parameter class in candidate models).
  2. Vijil Dome: Vijil's ModernBERT finetune, has the highest parameter efficiency in all candidate models.
  3. PromptForest-XGB: A custom XGBoost model trained on embedding features.

Predictions are aggregated using a weighted voting mechanism to form a sophisticated consensus mechanism that outperforms individual voters in calibration metrics.

2.1 Voting Method

We have found that using a Weighted Soft Voting approach is the most simple and effective voting method, ensuring that accurate models have more influence on the final result without drowning out weaker model's voices. We implement Soft Weighted Voting ($S_{w}$) as follows:
$$S_{w} = \frac{\sum_{i=1}^{n} (w_i \cdot p_i)}{\sum_{i=1}^{n} w_i}$$

2.2 Decision Threshold

PromptForest currently uses a simple decision threshold, which flags prompts as malicious when the weighted score exceeds the 0.5 threshold, which can be tuned.

2.3 Uncertainty Scoring

We measure model "uncertainty" using the Standard Deviation (σ) of all model predictions, scaled and capped at 1.0:
$$U = \min(2\sigma, 1.0)$$

3. Performance Evaluation

We benchmarked PromptForest against a wide range of detection models, with the SOTA model Qualifire Sentinel v2 being the best performing competitor model. You can view our benchmarking code in the benchmark folder.

3.1 Reliability & Calibration

PromptForest demonstrates superior calibration, meaning its reported confidence scores better reflect the true probability of correctness. Lower ECE (Expected Calibration Error) indicates more trustworthy probability outputs.

PromptForest Calibration Graph

Metric PromptForest (Ensemble) Qualifire Sentinel v2 (Baseline) Improvement
Accuracy 0.901 0.973 -
ECE (Lower is better) 0.070 0.096 +27%
Avg Conf. on Failure 0.642 0.760 +16%
Parameter Count ~237M ~600M 60% fewer

While Sentinel v2 achieves higher raw accuracy, PromptForest is significantly less confident when it makes mistakes (0.642 vs 0.760), making it safer for "Human-in-the-Loop" fallback systems.

3.2 Latency Benchmark

We compared the end-to-end inference time (including tokenization) on a standard Macbook CPU node on 3000 prompts of varying length. Our latency was significantly lower, despite our server overhead, showing that PromptForest is adequately optimised for real-world speed.

System Mean Latency (ms) P95 Latency (ms)
Qualifire Sentinel v2 ~225.77ms ~430.31ms
PromptForest ~141.07ms ~257.19ms

4. Quick Start

PromptForest is designed for immediate deployment.

# Install valid package
pip install promptforest

# Start the optimized inference server
promptforest serve --port 8000

The server automatically handles model downloading, caching, and ensemble orchestration.

5. Downstream Integration

PromptForest outputs the following JSON, where entries like is_malicious or confidence could be used in downstream tasks.

{
  "is_malicious": true/false,
  "confidence": float,
  "uncertainty": float,
  "malicious_score": float,
  "max_risk_score": float,
  "details": {
    "llama_guard": float,
    "vijil": float,
    "xgboost": float
  },
  "latency_ms": float
}

To configure the models or the output, please use a config.yaml.

6. Models & Attribution

PromptForest ensembles the following open-weights models:

Provider Model License
Meta Llama Prompt Guard 86M Llama Community
Vijil Vijil Dome Apache 2.0
Appleroll PromptForest-XGB Apache 2.0

7. Disclaimer

PromptForest is an evolving research project. It is not a standalone security solution and should be used as part of a defense-in-depth strategy alongside input sanitization (e.g., NeMo Guardrails) and system prompting.

License

Apache 2.0. See LICENSE and THIRD_PARTY_LICENSES for details.

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

promptforest-1.0.0.tar.gz (655.7 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

promptforest-1.0.0-py3-none-any.whl (677.3 kB view details)

Uploaded Python 3

File details

Details for the file promptforest-1.0.0.tar.gz.

File metadata

  • Download URL: promptforest-1.0.0.tar.gz
  • Upload date:
  • Size: 655.7 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.11.9

File hashes

Hashes for promptforest-1.0.0.tar.gz
Algorithm Hash digest
SHA256 0524d4ad50287347110e977a272d808f548f4a93f05a4c9ee6fe7fea747c7587
MD5 cbd751ab9293f72f21a627b5d6610761
BLAKE2b-256 23e7cd4bee90a34aab28c1a94f9d8ac8ab4f241620205e87e340c694a5f72e64

See more details on using hashes here.

File details

Details for the file promptforest-1.0.0-py3-none-any.whl.

File metadata

  • Download URL: promptforest-1.0.0-py3-none-any.whl
  • Upload date:
  • Size: 677.3 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.11.9

File hashes

Hashes for promptforest-1.0.0-py3-none-any.whl
Algorithm Hash digest
SHA256 7e31d97df68ad303033ee772f73b9aab9ade65b778d5d83a3a87e22e16d8cab2
MD5 ab18e3e6dd4784494c2ee6a72d1947cf
BLAKE2b-256 4472d4bc240cf2c9b614069958c5be833fd93e1c676959903191a0036691f315

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page