LatentMAS Safety Aligners (Qwen3-4B)
Safety aligners for the LatentMAS multi-agent system (Planner→Critic→Refiner→Judger
on Qwen3-4B). Each aligner reads the pre-Judger latent stack h_a [24, 2560]
(the K-vectors from the 3 pre-Judger agents) and acts to make the Judger refuse
harmful requests. Trained on a wildjailbreak-augmented harm-compliance dataset.
Artifacts
aligner_classifier_harmcompliance_WJ.pt (recommended)
The working aligner. Attention-pool + MLP (27M) that projects h_a into
Llama-Guard-3-8B's residual space; the frozen LG tail classifies safe/unsafe.
If p_unsafe > τ (fpr05 τ=0.749), a SAFETY NOTICE is injected into the Judger's
turn so it refuses in its own voice.
- val AUC 0.9555 (harm-compliance label), over-refusal ~0.01 at fpr05.
- Trained: harm-compliance = (prompt_harmful AND response_harmful), wildjailbreak- augmented (~315k rows, ~4× the harm-compliance positives vs the pre-WJ data).
aligner_sft_steer_WJ.pt (negative result — for completeness)
An SFT-trained aligner that projects h_a → a soft-prefix steering vector
injected into the frozen Qwen3-4B Judger, trained by teacher-forced LM-CE to make
the Judger generate a refusal. Does not work at inference: the vector injects
correctly but fails to change greedy generation (teacher-forced loss ≠ generation
control; scale mismatch). Included as a documented negative result. Use the
classifier above for actual defense.
Usage
Both are torch.load-able state dicts. See the LatentMAS repo's
scripts/aligner/deploy.py (classifier) / train_sft_steer.py (SFT-steer) and
the aligner_defense / aligner_sft_steer inference hooks in
LatentMAS/methods/latent_mas.py.