On this page: Repository layout · Setup · Data · Running AEGIS · Protocol · Ablations · Tests · Citation

Repository layout

Setup

Each model needs a Transformers version that can load it. The upstreams disagree, so we used one environment per model family:

• Models load from the Hub by default. To run offline, point the override variable at a local snapshot.
• Llama-Guard-3-8B, Gemma 4, and the Llama-3.1 backbone that Ultravox downloads are gated on the Hub. Accept their licenses and run huggingface-cli login, or set GUARD and AEGIS_LLAMA31_PATH to local copies.
• VITA-1.5 needs its GitHub code: clone https://github.com/VITA-MLLM/VITA and set AEGIS_VITA_REPO to the checkout (default: external/VITA).
• Hardware: one 48 GB GPU is enough for most models. Gemma 4 training needs a single card with at least ~40 GB, because the 262k-vocabulary loss does not shard well. Voxtral Small needs two 48 GB cards (DEVICE=auto shards the model; the gate and adapters follow the layers they hook). Qwen2-Audio and Phi-4 also benefit from DEVICE=auto on the long AJail clips.

Data

Download the five datasets as described in data/README.md, then build the manifests:

The script checks the row counts used in the paper: AJail 1,490, JALM 946, SACRED 1,364, XSTest 250, Benign_train 215. It also writes the in-domain test splits (718 / 499 / 683 rows) from the id lists in data/splits/.

Running AEGIS

For each model (PY is the Python of that model's environment; JUDGE_PY can point to a different environment for Llama Guard):

Every stage is resumable; rerun the same command after an interruption. Outputs go to runs/<model>/:

Protocol

• Evaluation settings. In-domain: AEGIS is trained on a benchmark's train split and evaluated on its held-out test split, one adapter per benchmark. The splits are 50/50, seed 42, grouped by prompt so variants of the same request stay on one side; the ids are in data/splits/. LOBO (leave one benchmark out): AEGIS is trained on the other two benchmarks and evaluated on the whole held-out benchmark. In both settings the evaluated data never takes part in threshold selection, and XSTest is never used for training.
• Training data (scripts/build_trainsets.py). Benign audio is Benign_train with the model's own undefended answers as targets. Benign prompts that look harmful and that the model refused at baseline are dropped. Harmful targets are generic refusal templates. The number of harmful clips equals the number of benign clips. Seed 42.
• Optimisation. AdamW, learning rate 2e-4, 4 epochs, gradient accumulation 8, LoRA rank 16, λ_BCE = 1.0, λ_L1 = 0.01.
• Prompt. Qwen2-Audio receives the audio only (PROMPT_MODE=native); the other models receive the audio plus one fixed text prompt. The gate reads a prompt-dependent hidden state, so an adapter is valid only under the prompt mode it was trained with.
• Threshold (scripts/calibrate_threshold.py). XSTest ids are sorted and split alternately into a validation half and a test half. The threshold is the smallest gate score whose added over-refusal on the validation half is at most 10%, counted on prompts the undefended model answered.
• Closed loop. The refusal probability sums the next-token probability of refusal openers at the last layer. In-domain runs use the multilingual opener list (REFUSAL_PHRASE_SET=multi, set by run_aegis.sh), so non-English refusals count.
• Metrics (scripts/summarize.py): unsafe rate is the share of evaluated rows that Llama-Guard-3-8B labels unsafe; over-refusal is the share of XSTest responses matching aegis/refusal.py, reported over all 250 prompts and over the test half; gate AUROC uses raw gate scores with attacks as positives and XSTest as negatives.

Ablations

Tests

Citation

License

The code is released under the MIT License. The base models (e.g. Gemma 4, Llama 3.1 in Ultravox, Llama-Guard-3-8B) and the benchmarks keep their own licenses and terms of use.