jailbreaker-llm

An LLM-vs-LLM automated red-teaming research platform. An attacker model, warm-started from an existing open-weight base and then fine-tuned via reinforcement learning, learns to discover prompts and multi-turn conversation strategies that cause a target LLM to violate its own policies -- with an independent judge ensemble scoring outcomes and a mutation module mining which strategies transfer across targets.

This is LLM engineering first, security second: the bulk of the work is RL fine-tuning (custom PPO), judge-model training, and multi-turn agent design. The security framing is what the trained system is pointed at.

This repo intentionally stops at a generic scored vulnerability report -- there is no vendor-specific disclosure-ticket export here. If you need to feed a specific vulnerability-management process, build that as a separate step that reads report.json; don't fork this repo to add it.

responsible use

This repo is for authorized safety evaluation only -- running it against a model you own, control, or have explicit permission to test (your own fine-tunes, local open-weight models, or a target you've been given permission/API access to red-team). It is not configured with any working exploits out of the box; data/ includes placeholder-based integration data, high-level local smoke seeds, and a reproducible importer for the official JailbreakBench behavior benchmark. Do not point the attacker loop at third-party production systems without authorization -- that's the same rule as any other pentesting tool.

Enforcement, not just policy: configs/allowlist.yaml lists the only API hosts the attacker/agent loop is permitted to call, and that check runs INSIDE APITarget.__init__ itself (src/targets/allowlist.py), not in one factory function some call sites might skip. A circuit breaker halts a run after repeated consecutive failures against a target instead of retrying forever. See "architecture" below for the rest of the safety-relevant design.

architecture

 ┌─────────────────────────┐
 │  Base LM (warm-start)      │  src/model/pretrain.py --
 │  existing open-weight base,  │  configs/base_model.yaml:
 │  its own tokenizer             │  pretrain.warm_start_from
 └──────────────┬─────────────────┘
                ▼
 ┌─────────────────────────────┐        ┌────────────────────┐
 │  SimplePPOTrainer              │◄──────│ Reward Judge          │  src/judge/ --
 │  src/attacker/ppo_trainer.py   │       │ (train-only classifier)│  train-only, never
 │  generate -> judge -> GAE ->   │       └────────────────────┘  scores a final report
 │  clipped PPO update            │
 │  + diversity term              │  reward.py: judge score minus
 │  (src/attacker/reward.py)      │  embedding similarity to this
 └──────────────┬────────────────┘  run's recent successful attacks
                ▼
 ┌─────────────────────────────┐
 │   Multi-turn agent loop       │  src/agent/multi_turn_loop.py --
 │   technique menu (persona,      │  best-of-N branch search: try a
 │   hypothetical, obfuscation,      │  few techniques against the
 │   escalate/retreat) + tree-search   │  target, keep whichever the
 │   scores the FULL conversation        │  live judge scores highest
 │   so far, not just the latest turn      │
 └──────────────┬─────────────────────────┘
                ▼
     ┌─────────────────────┐
     │  Scope gate            │  src/targets/allowlist.py --
     │  allowlist + circuit     │  enforced inside
     │  breaker                   │  LocalTarget/APITarget.respond()
     └──────────┬─────────────────┘
                ▼
      ┌───────────────────┐
      │   Target LLM(s)    │  src/targets/ (local or API, retry+backoff+timeout)
      └──────────┬─────────┘
                 ▼
      ┌───────────────────┐
      │  Campaign runner   │  src/eval/run_campaign.py -- resumable,
      │                     │  writes an authorization ledger entry
      └──────────┬─────────┘  per row (target_mode/target_host)
                 ▼
     ┌─────────────────────┐
     │  Redacted transcript   │  src/eval/redact.py -- strips
     │  store                    │  PII/diagnostic-shaped data before
     └────────┬────────┬─────────┘  anything downstream reads a row
              ▼        ▼
 ┌────────────────┐   ┌─────────────────────────┐
 │  Eval Judge        │   │  Mutation/clustering       │
 │  classifier + LLM-    │   │  src/mutation/                │
 │  judge ensemble,        │   │  embedding clusters, not        │
 │  scores full transcript   │   │  TF-IDF; drops rows whose         │
 │  (src/judge/eval_judge.py)  │   │  target authorization expired       │
 └────────┬───────────────────┘   └─────────────────────────┘
          ▼
   ┌─────────────────────────┐
   │  Scored report              │  src/eval/report.py --
   │  (flags disagreement            │  generic, no ticket/disclosure
   │  between judges for human           │  export; build that
   │  review, doesn't auto-decide)         │  separately if you need it
   └─────────────────────────┘

Three design decisions worth calling out:

The Reward Judge and Eval Judge are never the same model. If the attacker trained directly against the judge that later "proves" a finding is real, PPO would happily learn to exploit that judge's blind spots instead of finding real policy violations. The Reward Judge exists only to produce a fast training signal; the Eval Judge that scores a completed campaign is a differently-built ensemble (a separately-trained classifier plus an LLM-as-judge) that never sees a gradient from the attacker it's grading.

The judge scores the conversation, not the last message. A judge that only sees the latest exchange can't recognize a violation that only makes sense in light of what was set up two turns earlier -- which defeats the purpose of a multi-turn agent. Both the live per-turn judge (windowed context, src/agent/multi_turn_loop.py::_score_with_context) and the eval-time ensemble (src/judge/eval_judge.py) score against the conversation so far, not just the current turn.

Warm-start, not pretrain-from-scratch, is the default. A from-scratch run is the biggest single time sink in this pipeline and isn't necessary to get the actual differentiator (an attacker trained on reward from the real target) -- see configs/base_model.yaml. Pretraining from scratch is still supported (pretrain.warm_start_from: null) if you deliberately want your own architecture and tokenizer; src/tokenizer/ is only relevant on that path.

The default warm start is now Qwen/Qwen2.5-0.5B-Instruct. PPO updates LoRA adapters plus the value head instead of every base-model parameter. The KL reference uses the same frozen base with adapters temporarily disabled, so the 6 GB laptop GPU does not hold a second 0.5B model copy. Saved attacker checkpoints contain the compact adapter and value head; the unchanged base is reused directly from the Hugging Face cache. This avoids an extra roughly 1 GB project copy on Windows. Set pretrain.copy_to_out_dir: true only when a standalone base-model directory is explicitly needed.

local performance

Local inference is hardware-aware and configured under target::

• dtype: auto uses BF16 on supported CUDA GPUs, FP16 on older CUDA GPUs, and FP32 on CPU.
• attn_implementation: sdpa and use_cache=True select PyTorch's optimized attention path and reuse the KV cache during generation.
• max_input_tokens bounds the growing multi-turn context; truncation keeps the newest turns.
• batch_size batches best-of-N target trials. The selected trial response is reused on the next turn instead of being generated a second time.
• Hugging Face attacker rollouts are also generated as one left-padded batch instead of looping over every PPO prompt serially.
• PPO policy/reference scoring and loss passes use configurable micro-batches (rl.ppo_micro_batch_size) to trade a small amount of VRAM for much faster training without requiring the full rollout batch to fit at once.
• No candidate batch is generated after the final allowed turn.

With the default six turns and three candidates, the old loop could perform 24 serial target generations (including duplicate winners and a final unused branch). The new local path performs six generation rounds and evaluates 16 sequences. Actual wall-clock speedup depends on the target model and GPU.

Optional quantization: 4bit or 8bit is available for CUDA targets after installing bitsandbytes. Quantization is primarily a memory-saving option; for a small model it may be slower than BF16.

module map

Path	What it does
src/tokenizer/	Trains a custom BPE tokenizer -- only used for a from-scratch run
src/model/	Base transformer (nanoGPT-style) + from-scratch pretraining loop; pretrain.py also handles warm-starting from an existing open-weight base
src/attacker/value_head.py	Adds a value head to either backend and supports a shared-base LoRA reference policy
src/attacker/ppo_trainer.py	Custom PPO (GAE + clipped surrogate + KL-to-reference penalty), advantages normalized across the batch
src/attacker/reward.py	Judge score -> scalar reward, plus an embedding-similarity diversity penalty so PPO doesn't collapse onto one trick
src/attacker/technique_library.py	Named attack techniques (persona/hypothetical framing, obfuscation, escalation) the multi-turn agent and mutation module both draw from
src/attacker/rl_finetune.py	The RL fine-tuning entry point -- resumable, redacts sensitive config before any wandb logging
src/judge/train_judge.py, judge_infer.py	Trains/runs the train-only Reward Judge classifier
src/judge/llm_judge.py, eval_judge.py	The independent eval-time judge: LLM-as-judge + a separately-trained classifier, combined into an ensemble that flags disagreement for human review
src/agent/	Multi-turn conversational loop: technique menu + best-of-N branch search, scores the full conversation so far
src/targets/	TargetModel interface, LocalTarget/APITarget (both scope-gated + circuit-breaker protected), shared build_target factory
src/mutation/	Generates variants of a successful attack via the technique library; clusters transcripts by embedding similarity (not TF-IDF) to find transferable patterns, dropping revoked targets
src/eval/	Runs a campaign (run_campaign.py, resumable), redacts transcripts (redact.py), produces a generic scored report (report.py)
scripts/relabel_for_judge.py	Converts a binary refusal/compliance benchmark into the judge's 3-class scheme
src/analysis/	ML coursework analysis of campaign transcripts (see below), run by main_analysis.py

ML coursework analysis (21CSC305P PBL)

python main_analysis.py generates the synthetic campaign (data/generate_synthetic_campaign.py: 978 turns, 220 conversations, with a known hidden compliance state per turn). It then runs every step below and writes plots and metrics.json to outputs/analysis/.

Review	Module	Models
1	eda.py, severity_regression.py, state_classification.py	EDA, linear regression, logistic regression / SVM classification (with Random Forest baselines)
2	unsupervised.py	PCA, K-means, Gaussian mixture (EM), hierarchical clustering, all written from scratch in NumPy
3	strategy_clustering.py, compliance_hmm.py, comparison.py	DBSCAN baseline, HMM, comparative analysis

Review 2: unsupervised models (src/analysis/unsupervised.py)

The question for Review 2: can the hidden state (refusal, partial, full_violation) be recovered from the three live-judge probabilities alone, with no labels? Each model is written from its equations, with no scikit-learn or SciPy.

• PCA. Build the covariance S = (1/N) Σ (xₙ − x̄)(xₙ − x̄)ᵀ, solve S uᵢ = λᵢuᵢ, and project with z = Uᵀ(x − x̄). The three probabilities sum to 1, so the third eigenvalue is about 0. The clustering models therefore run in the 2-D principal subspace.
• K-means. Minimizes J = Σₙ ‖xₙ − μ_k(n)‖². The E-step assigns each point to its nearest centre and the M-step moves each centre to the mean of its points. Uses k-means++ seeding and keeps the best of 10 restarts; k is chosen with the elbow and the silhouette.
• Gaussian mixture via EM. Models p(x) = Σ π_k N(x | μ_k, Σ_k). The E-step computes responsibilities γ_nk ∝ π_k N(xₙ | μ_k, Σ_k). The M-step re-estimates N_k, μ_k, Σ_k and π_k. The log-likelihood never decreases, and BIC is reported for each K.
• Hierarchical (agglomerative) clustering. Clusters whole conversations (6 aggregate features each). It starts from single points and repeatedly merges the closest pair, updating distances with the Lance–Williams rule for single, complete, average or Ward linkage.
• Evaluation. Silhouette (b − a)/max(a, b), Adjusted Rand Index against the true state, and accuracy after matching cluster labels to states.

Model	ARI vs true state	Accuracy after matching	Silhouette
K-means (k = 3)	0.924	0.969	0.684
GMM-EM (K = 3)	0.865	0.951	0.652
Hierarchical, complete linkage (conversations)	0.692	0.805	0.423

The NumPy K-means reaches the same ARI (0.924) as the scikit-learn K-means in strategy_clustering.py. tests/test_unsupervised.py checks every model on toy data.

build order

1. src/model/pretrain.py -- warm-start by default; only touch src/tokenizer/ if you deliberately choose a from-scratch run
2. src/judge/train_judge.py -- train the Reward Judge; train a SEPARATE checkpoint for the Eval Judge's classifier half from a different split/seed
3. src/attacker/ -- RL fine-tuning, depends on (1) and (2)
4. src/agent/ -- multi-turn behavior on top of the trained attacker
5. src/eval/ -- run a campaign and get a scored report
6. src/mutation/ -- once you have real transcripts, mine them for transferable strategies

See scripts/ for the runnable order (01_...sh through 05_...sh).

setup

python -m venv .venv && source .venv/bin/activate
pip install -r requirements.txt

On Windows, the complete local integration path is:

.\scripts\run_smoke_test.ps1

Import and locally validate JailbreakBench with:

python data\import_jailbreakbench.py
.\scripts\run_jailbreakbench.ps1 -Split harmful -Limit 2

JailbreakBench is kept as evaluation data, not PPO or judge-training data.

See docs/FULL_WORKINGS.md for component-level behavior, individual commands, verified measurements, every material change, and the distinction between the synthetic smoke judges and a real evaluation.

tests

pytest

Coverage focuses on the orchestration code that previously had none: the multi-turn agent loop (including a regression test for a real crash bug -- see tests/test_multi_turn_loop.py's docstring), the allowlist/circuit breaker, and the technique library's branch-selection logic. Anything requiring network access (downloading distilbert-base-uncased, a sentence-transformers model, or a warm-start base) is left out of the default test run by construction -- those paths are exercised by actually running the pipeline, not by unit tests.

known limitations

• Single target per campaign. Every config here points at exactly one target. A fleet/campaign-matrix concept (useful once you have more than one authorized target, e.g. multiple product versions) isn't built -- add it once a single-target campaign is solid, not before.
• The Eval Judge's classifier half inherits DistilBERT's 512-token window. It's given the most recent turns that fit, not the oldest -- but on a long conversation it is not seeing everything the LLM-judge half sees. A longer-context encoder or a summarize-then-score step would close this gap.
• The judge training data is heuristically labeled. scripts/relabel_for_judge.py synthesizes the partial_compliance class from regex/length heuristics. Its docstring says to spot-check ~50 rows before trusting it -- nothing enforces that happening.
• PII redaction (src/eval/redact.py) is pattern-based, not a guarantee. It catches structurally-recognizable categories (emails, phone numbers, card numbers, IPs, VIN-shaped strings), not every way sensitive data can appear in free text.