ONNX export of the multilingual-e5-large-wjb-threatfeed_v1 model — a fine-tuned sentence-transformers model for detecting duplicate vulnerability submissions (jailbreak and prompt injection attacks) in the 0din threat feed.
It maps prompts to a 1024-dimensional dense vector space optimized for semantic similarity comparison of attack prompts.
This model achieves a +59.5% F1 improvement over the OpenAI text-embedding-3-large baseline on duplicate detection, and is the best-performing model in the series.
Model Details
Model Description
Model Type: Sentence Transformer (two-stage fine-tuned), exported to ONNX
Format: ONNX (compatible with onnxruntime, tract-onnx, and other ONNX runtimes)
Embedding Pipeline
Input Text → Tokenizer → ONNX Model → Mean Pooling → L2 Normalization → Embedding
The ONNX model contains only the transformer backbone. Mean pooling and L2 normalization must be implemented in application code (see usage examples below).
Model Inputs
The ONNX model requires 3 inputs:
input_ids: Token IDs from tokenizer
attention_mask: 1 for real tokens, 0 for padding
token_type_ids: All zeros for single-sentence embeddings
ONNX Verification
The ONNX export produces near bit-for-bit identical embeddings to the native sentence-transformers model (0.000001 max difference across all test sentences).
Intended Use
This model is designed for:
Duplicate detection in AI security vulnerability reports (jailbreak/prompt injection attacks)
Semantic similarity comparison of attack prompts that may use different surface-level techniques but target the same underlying vulnerability
Embedding generation for LSH-based similarity search in vulnerability management systems
Edge/server deployment via ONNX runtime without requiring PyTorch
The model is trained to recognize semantic equivalence between attack prompts even when they use different jailbreak tactics (e.g., role-playing, encoding, academic framing) to elicit the same harmful behavior.
1import numpy as np
2import onnxruntime as ort
3from tokenizers import Tokenizer
45# Load model and tokenizer6session = ort.InferenceSession("onnx/model.onnx")7tokenizer = Tokenizer.from_file("tokenizer.json")8tokenizer.enable_padding(pad_id=1, pad_token="<pad>")9tokenizer.enable_truncation(max_length=512)1011# Tokenize12texts =["First attack prompt","Second attack prompt"]13encodings = tokenizer.encode_batch(texts)14input_ids = np.array([e.ids for e in encodings], dtype=np.int64)15attention_mask = np.array([e.attention_mask for e in encodings], dtype=np.int64)16token_type_ids = np.zeros_like(input_ids)1718# Run ONNX inference19outputs = session.run(None,{20"input_ids": input_ids,21"attention_mask": attention_mask,22"token_type_ids": token_type_ids,23})24token_embeddings = outputs[0]# [batch, seq_len, 1024]2526# Mean pooling27mask = attention_mask[:,:, np.newaxis].astype(np.float32)28embeddings =(token_embeddings * mask).sum(axis=1)/ mask.sum(axis=1)2930# L2 normalization31norms = np.linalg.norm(embeddings, axis=1, keepdims=True)32embeddings = embeddings / norms
3334# Cosine similarity35similarity = np.dot(embeddings[0], embeddings[1])36print(f"Similarity: {similarity:.4f}")
Rust (tract-onnx)
rust
1usetract_onnx::prelude::*;2usetokenizers::Tokenizer;34// Load model and tokenizer5let model =tract_onnx::onnx()6.model_for_path("onnx/model.onnx")?7.into_optimized()?8.into_runnable()?;9let tokenizer =Tokenizer::from_file("tokenizer.json")?;1011// Tokenize12let encoding = tokenizer.encode("Attack prompt text",true)?;13let input_ids:Vec<i64>= encoding.get_ids().iter().map(|&x| x asi64).collect();14let attention_mask:Vec<i64>= encoding.get_attention_mask().iter().map(|&x| x asi64).collect();15let token_type_ids:Vec<i64>=vec![0i64; input_ids.len()];1617// Run inference, then apply mean pooling + L2 normalization18// (see full Rust implementation at github.com/0din-ai)
Training Details
This model was trained using a two-stage fine-tuning approach:
Stage 1: WildJailbreak Pre-training
Pre-trained on public synthetic data to learn jailbreak semantics.
Purpose: Calibrate the model for real-world duplicate detection on production vulnerability data
Stage 2 Hyperparameters
Parameter
Value
Epochs
50 (early stopped)
Batch size
8 (per device) x 4 gradient accumulation = 32 effective
Learning rate
1e-5
LR scheduler
Linear
Warmup ratio
0.1
Weight decay
0.01
FP16
True
Early stopping patience
10
Eval steps
50
Seed
1
Evaluation Results
Duplicate Detection Performance
Evaluated on 55 human-labeled vulnerability pairs (10 duplicates, 45 non-duplicates) from a corpus of 3,749 vulnerabilities. Best F1 score at each model's optimal threshold:
Model
Best F1
Threshold
Precision
Recall
OpenAI text-embedding-3-large (baseline)
0.462
0.80
1.000
0.300
Finetuned V1 (WildJailbreak only, e5-small)
0.500
0.50
0.333
1.000
Finetuned V2 (WJB + threat feed v1, e5-small)
0.526
0.70
0.556
0.500
Finetuned V3 (WJB + threat feed v2, e5-small)
0.556
0.75
0.625
0.500
Finetuned V4 (WJB + threat feed 10k, e5-small)
0.600
0.70
0.600
0.600
Finetuned Base V1 (e5-base)
0.696
0.70
0.615
0.800
This model (Large V1)
0.737
0.80
0.778
0.700
Threshold Analysis (This Model)
Threshold
Precision
Recall
F1
TP
FP
FN
TN
0.50
0.250
0.900
0.391
9
27
1
18
0.55
0.310
0.900
0.462
9
20
1
25
0.60
0.346
0.900
0.500
9
17
1
28
0.65
0.391
0.900
0.545
9
14
1
31
0.70
0.500
0.800
0.615
8
8
2
37
0.75
0.615
0.800
0.696
8
5
2
40
0.80
0.778
0.700
0.737
7
2
3
43
0.85
1.000
0.400
0.571
4
0
6
45
0.90
1.000
0.200
0.333
2
0
8
45
Key Findings
+59.5% F1 improvement over the OpenAI text-embedding-3-large baseline (0.737 vs 0.462)
Best in series: Continues the scaling trend: e5-small (0.600) → e5-base (0.696) → e5-large (0.737).
Highest precision at optimal threshold: 0.778 precision with only 2 false positives, compared to 0.615 for e5-base at its optimal threshold.
Precision-recall tradeoff vs e5-base: Trades a small amount of recall (0.700 vs 0.800) for a significant precision gain (0.778 vs 0.615), resulting in a better-balanced F1.
Higher optimal threshold (0.80): The larger model produces more confident and well-separated similarity scores, allowing a higher decision threshold while maintaining strong performance.
Strong recall at lower thresholds: Maintains 0.900 recall across thresholds 0.50–0.65, indicating very few true duplicates are missed at permissive thresholds.
Note: The evaluation dataset is small (55 pairs, 10 positive). With only 10 true duplicates, each TP/FP change causes large metric swings. Results should be interpreted with caution.
Limitations
Small evaluation set: Only 55 human-labeled pairs (10 duplicates). Results should be taken as directional rather than definitive.
LLM annotation bias in training data: Stage 2 training data was annotated by a single LLM (Gemini 2.5 Pro), which may affect calibration.
Model size: ~560M parameters with 1024-dim embeddings. The ONNX model is ~2.1GB.
Domain-specific: Optimized for jailbreak/prompt injection duplicate detection. Performance on general semantic similarity tasks is not evaluated.
Single-turn only: This model was only trained on single-prompt jailbreaks and should not be used to process multi-turn conversations. In the future, we plan to release models that can handle multi-turn jailbreak scenarios.
Citation
BibTeX
Sentence Transformers
bibtex
1@inproceedings{reimers-2019-sentence-bert,
2 title = "Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks",
3 author = "Reimers, Nils and Gurevych, Iryna",
4 booktitle = "Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing",
5 month = "11",
6 year = "2019",
7 publisher = "Association for Computational Linguistics",
8 url = "https://arxiv.org/abs/1908.10084",
9}
ContrastiveLoss
bibtex
1@inproceedings{hadsell2006dimensionality,
2 author={Hadsell, R. and Chopra, S. and LeCun, Y.},
3 booktitle={2006 IEEE Computer Society Conference on Computer Vision and Pattern Recognition (CVPR'06)},
4 title={Dimensionality Reduction by Learning an Invariant Mapping},
5 year={2006},
6 volume={2},
7 number={},
8 pages={1735-1742},
9 doi={10.1109/CVPR.2006.100}
10}
WildJailbreak
bibtex
1@article{jiang2024wildteaming,
2 title={WildTeaming at Scale: From In-the-Wild Jailbreaks to (Adversarially) Safer Language Models},
3 author={Jiang, Liwei and Bhatt, Kavel and Phute, Seungju and Hwang, Jaehun and Liang, Dongwei and Sap, Maarten and Hajishirzi, Hannaneh and Choi, Yejin},
4 journal={arXiv preprint arXiv:2406.18510},
5 year={2024}
6}