asatheesh/latent-mas-aligner

LatentMAS Aligner — Qwen3-4B ↔ Llama-Guard-3-8B (attention-pool MLP4k, 60 epochs)

A 27M-parameter alignment module that projects Qwen3-4B's multi-agent latent reasoning (h_a) into Llama-Guard-3-8B's residual space so that Llama-Guard's frozen tail can decode a safety verdict directly from the audited model's mid-reasoning state — no text bridge required at deploy time.

Trained on 195,589 (prompt, response) records from a LatentMAS run with Qwen3-4B agents (Planner → Critic → Refiner → Judger), with WildGuard-generated safe/unsafe labels.

• Validation AUC: 0.9538 (18,915 held-out records)
• Peak val AUC during training: 0.9566 at epoch 53
• Full deploy-time results across a three-threshold sweep are in Evaluation below.

Files

file	purpose
aligner_qwen3-4b_llamaguard-3-8b_ep60.pt	model checkpoint (state dict + config)

Architecture

h_a  [B, 24, 2560]   (Qwen3-4B pre-Judger latents; 3 stages × 8 K-vectors)
  │
  ▼
attention pool with learnable query [2560]
  │
  ▼
2-layer MLP:  Linear(2560, 4096) → GELU → Linear(4096, 4096)
  │
  ▼
w_h  [B, 4096]   (aligned into Llama-Guard block-16 residual space)

w_h is then fed to Llama-Guard-3-8B's blocks 17-31 + norm + lm_head as if it were a length-1 sequence residual at position 0. Softmax over {safe_token_id, unsafe_token_id} gives p_unsafe.

Load recipe

import torch
from scripts.aligner.aligner import Aligner  # see github.com/Asatheesh6561/LatentMASHarmBench

ckpt = torch.load("aligner_qwen3-4b_llamaguard-3-8b_ep60.pt",
                  weights_only=False, map_location="cpu")

aligner = Aligner(
    d_a=ckpt["d_a"],           # 2560 (Qwen3-4B hidden size)
    d_g=ckpt["d_g"],           # 4096 (Llama-Guard hidden size)
    hidden_dim=ckpt["hidden_dim"],   # 4096
    pool=ckpt["pool"],         # "attention"
)
aligner.load_state_dict(ckpt["state_dict"])
aligner.eval()

At deploy time, combine with:

• Llama-Guard-3-8B tail (blocks 17-31 + norm + lm_head) — download from meta-llama/Llama-Guard-3-8B and extract via the repo's guard_head.py.
• Live h_a from LatentMAS's per-stage tap (stage_last_layer_K in LatentMAS/methods/latent_mas.py).

Full deploy code is in Asatheesh6561/LatentMASHarmBench under LatentMAS/methods/latent_mas.py::_run_aligner_defense and scripts/aligner/deploy.py.

Training config

hyperparameter	value
optimizer	AdamW
learning rate	3e-4
batch size	512
epochs	60
loss	MSE(W(h_a), h_g) + CE(head(W(h_a)), v_g) (λ = 1.0)
frozen decoder	Llama-Guard-3-8B blocks 17-31 + norm + lm_head
tap layer	ℓ_G = 16 (Llama-Guard's block-16 residual, position -1)
features	195,589 records, pool=none (24 K-vectors per record)

Training set

The (h_a, h_g, v_g) training tensors are packed together in this file:

• Extraction shards on HF: asatheesh/latent-mas-safety-dataset (shards 0..764, WildGuard labels, prompt/response bundles).
• The Llama-Guard-side h_g residuals are packaged as h_g_ellG=16.pt on the dataset repo.
• Combine with your local Llama-Guard-3-8B download to reproduce.

Evaluation

Val AUC 0.9538 on the 18,915 held-out records (WildGuard labels).

Deploy-time results (three-threshold sweep)

Wire-in to LatentMAS-Qwen3-4B, evaluated on 4 benchmarks. Harmful ASR and refusal rate — lower is better on both.

threshold	HB ASR ↓	StrongREJECT ↓	OR-Bench-hard-1k refuse ↓	PHTest refuse ↓
τ = 0.40 (strict)	0.159	0.171	0.023	0.043
τ = 0.25 (balanced)	0.141	0.053	0.055	0.125
τ = 0.10 (aggressive)	0.128	0.079	0.120	0.182

The three thresholds correspond to val-ROC operating points of FPR ≈ {5%, 10%, 20%} respectively.

Comparison with baselines

architecture	HB ASR	StrongREJECT	OR-Bench	PHTest
latent-MAS (no defense)	0.344	0.283	0.003	0.004
latent-MAS + Llama-Guard (weak decode)	0.372	0.293	0.007	0.000
latent-MAS + Llama-Guard (full decode)	0.303	0.210	0.011	0.007
aligner τ=0.40	0.159	0.171	0.023	0.043
aligner τ=0.25	0.141	0.053	0.055	0.125
aligner τ=0.10	0.128	0.079	0.120	0.182
text-MAS + Llama-Guard	0.081	0.127	0.043	0.020

Compared to the strongest existing latent-space defense (latent_mas_cotguard_full):

• HB ASR: −58% relative at τ=0.10 (0.128 vs 0.303).
• StrongREJECT: −75% relative at τ=0.25 (0.053 vs 0.210).

Compared to text-space text_mas_cotguard (which runs full Llama-Guard on decoded outputs — much more expensive at deploy):

• Beats it on StrongREJECT at every τ.
• Slightly behind on HarmBench (0.128–0.159 vs 0.081).
• Roughly comparable over-refusal on OR-Bench.

ROC curves

Full validation ROC across three aligner variants (linear-mean-bypass, MLP-attention-bypass, MLP-attention-fulltail — this checkpoint) at results/paper/aligner/val_roc.png in the code repo.

Intended use

• Research on latent-space safety monitoring for multi-agent systems.
• Reproducing the paper's aligner-defense results on Qwen3-4B LatentMAS.
• Ablations that vary the aligner architecture while keeping the audited-model and safety-monitor sides fixed.

Not intended for

• Production deployment as a standalone content classifier — this aligner is meaningful only in the LatentMAS context (it consumes Qwen3-4B's pre-Judger latents, not text).
• Any use that requires distributing Llama-Guard's weights — those remain under Meta's Llama-Guard-3-8B license and must be obtained separately.

License

Apache 2.0 for the aligner weights (this file). Llama-Guard-3-8B weights are NOT included and remain under Meta's original license.

Citation

If you use this checkpoint, please cite the paper describing LatentMAS alignment (link forthcoming).