rt-antigravity-wildguard-v1

LoRA fine-tune of allenai/wildguard hardened against jailbreak attacks that the base wildguard model historically let through.

Trained as part of the rt_antigravity red-teaming pipeline.

What it does

Same interface as the base model — classifies (prompt) or (prompt, response) pairs as harmful / safe with a refusal flag. Outputs follow the wildguard template:

Harmful request: yes/no
Response refusal: yes/no
Harmful response: yes/no

Use the chat template shipped in tokenizer_config.json to format inputs.

Training data

• Unsafe pool: guard_jailbreaks_wildguard.jsonl — attack prompts that bypassed the base wildguard guard in earlier pipeline runs (across multiple attack papers and a Reddit-jailbreak corpus). Mined with aggregate_guard_jailbreaks.py --guard wildguard --score 4.
• Safe pool: allenai/wildguardmix (wildguardtrain split) — to preserve the original classifier's behaviour on benign prompts and avoid over-refusal.
• Mix proportion: 40% unsafe / 60% safe.
• Augmentations: homoglyphs and character-level typos applied to each prompt to harden against simple obfuscation attacks.
• Final dataset size: ~25k examples (see datasets/wildguard.json).

Training setup

• Method: SFT via TRL SFTTrainer with PEFT LoRA, then merged back into the base weights for single-checkpoint deployment.
• LoRA: r=16, alpha=32, dropout=0.05, target_modules=[q_proj, k_proj, v_proj, o_proj].
• Training args: batch_size=4, epochs=1, lr=1e-4, max_length=2048, cosine LR schedule, warmup_ratio=0.05.
• Hardware: single H200 GPU, bf16.

Build/train pipeline (reproducible from the repo):

python3 src/finetune/aggregate_guard_jailbreaks.py --guard wildguard --score 4 \
        --output guard_jailbreaks_wildguard.jsonl
python3 src/finetune/build_all_guards.py --guards wildguard \
        --unsafe guard_jailbreaks_wildguard.jsonl --wg-unsafe-n 2000 \
        --proportion 0.4 --copies 3
python3 src/finetune/train_all.py --guard wildguard --gpu 0
python3 src/finetune/merge_lora.py --guard wildguard \
        --lora-path ./models/wildguard-wildguardmix-only-alpha/final \
        --out-path ./models/wildguard-ft-merged

Intended use

• Safety classifier inside multi-guard ensembles for red-teaming pipelines.
• Drop-in replacement for allenai/wildguard where the original model was too lenient on dialogue-style jailbreaks.

Limitations

• Specialised on attack distributions present in the training pool — generalisation to attack styles not represented there (e.g. low-resource language jailbreaks, image-injection attacks) is untested.
• Inherits all upstream limitations and licensing of allenai/wildguard.
• Single-judge model; for high-stakes deployments use it as part of an ensemble with diverse base architectures.

Citation

If you use this model, please cite the upstream wildguard paper and link to rt_antigravity.