How I'm training a prompt injection detector

1 point

a year ago

I’ve been experimenting with different classifiers to catch prompt injection. They work well in some cases, but not in other. The issue is they are mostly trained for conversational agents. But for agents consuming other content they fall short. So, noticing different cases where I’ve had issues with them, I’ve decided to train one myself.

What data I use?

Public datasets from hf: jackhhao/jailbreak-classification, deepset/prompt-injections

Custom:

- collected attacks from ctf type prompt injection games,

- added synthetic examples,

- added 3:1 safe examples,

- collected some regular content from different web sources and documents,

- forked browser-use to save all extracted actions and page content and told it to visit random sites, - used claude to create synthetic examples with similar structure,

- made a script to insert prompt injections within the previously collected content

What model I use? mdeberta-v3-base Although it’s a multilingual model, non-english languages is still something to improve on in next iterations of the classifier.

Where do I train it? Google colab. Easy and no need to burn local machine.

Will be keeping track where the model falls short. I’d encourage you to try it out and if you notice where it fails, please let me know and I’ll be retraining it with that in mind.

Here is the trained model: https://huggingface.co/proventra/mdeberta-v3-base-prompt-injection

Also, it's becoming clear that different use cases and content types call for specific training. If you have a use case where you think that would be needed, please feel free to reach out https://www.proventra-ai.com/contact