Danny-1223/latent-mas-linear-probe

LatentMAS Linear Probe

A logistic linear probe over LatentMAS pre-Judger hidden states. It is the simplest interpretability baseline for latent multi-agent safety: no attention pooling, no Llama-Guard tail, just a z-scored linear classifier on stage-mean activations.

Checkpoints are grouped by backbone:

qwen3-4b/            Qwen3-4B            d_h=2560, x ∈ R^7680
qwen3-4b-instruct/   Qwen3-4B-Instruct-2507   d_h=2560, x ∈ R^7680
qwen3-14b/           Qwen3-14B           d_h=5120, x ∈ R^15360

Root-level linear_probe_*.pt files are Qwen3-4B copies kept for backward compatibility. Prefer the foldered paths below.

What it reads

For a sequential 3-agent LatentMAS run (planner → critic → refiner) with K = 8 latent steps:

x = concat(mean_K(H_planner), mean_K(H_critic), mean_K(H_refiner))
score = sigmoid( w · (x − μ) / σ + b )

x is [3 · d_h]: 7680 on 4B / 4B-Instruct, 15360 on 14B.

On graphs with more than three agents the runtime still feeds the 3-role vector: it picks the first planner, first critic, and first refiner. Extra agents still run in the MAS (they affect the Judger via KV), but the probe does not see their latents.

When score ≥ τ, the same [SAFETY NOTICE] is injected before the Judger as in CoT-Guard / the attention aligner. The Judger itself is unguarded.

Files

Each folder contains:

File	Role
linear_probe_harmcompliance.pt	Recommended (seed 0). y = prompt_harmful ∧ response_harmful.
linear_probe_harmcompliance_seed{0,1,2}.pt	Same recipe, three independent inits / shuffles. seed0 is identical to the file above.
linear_probe_prompt_harm.pt	Ablation. y = prompt_harmful only. Strong ASR, collapsed utility — not the recommended defense.
linear_probe_prompt_harm_seed{0,1,2}.pt	Same ablation, three seeds. seed0 is identical to the file above.

Each .pt is a small dict (~95 KB on 4B, ~185 KB on 14B):

arch, feature, label_mode, n_stages, d_h,
weight [3·d_h], bias, mu [3·d_h], sigma [3·d_h],
thresholds {p01, fpr05, fpr10, fpr20, p50},
metrics, training_meta

Recommended checkpoint

Use {backbone}/linear_probe_harmcompliance.pt. Relabeling from “the prompt is harmful” to “the trajectory complied” is what keeps over-refusal in the same band as the attention aligner.

Validation (fpr thresholds on val over-refusal)

Backbone	Probe	seed	val AUC	val AP	τ fpr05	τ fpr10
Qwen3-4B	harmcompliance	0	0.9268	0.5253	0.5348	0.3546
Qwen3-4B	harmcompliance	1	0.9269	0.5117	0.5460	0.3692
Qwen3-4B	harmcompliance	2	0.9274	0.5129	0.5454	0.3456
Qwen3-4B	prompt_harm	0	0.9787	0.9580	0.4089	0.1650
Qwen3-4B	prompt_harm	1	0.9790	0.9577	0.4111	0.1795
Qwen3-4B	prompt_harm	2	0.9789	0.9579	0.3883	0.1626
Qwen3-4B-Instruct	harmcompliance	0	0.9446	0.5233	0.8637	0.6970
Qwen3-4B-Instruct	harmcompliance	1	0.9448	0.5166	0.8627	0.6875
Qwen3-4B-Instruct	harmcompliance	2	0.9444	0.5171	0.8529	0.6796
Qwen3-4B-Instruct	prompt_harm	0	0.9875	0.9895	0.4058	0.1727
Qwen3-4B-Instruct	prompt_harm	1	0.9875	0.9896	0.4406	0.1832
Qwen3-4B-Instruct	prompt_harm	2	0.9876	0.9896	0.4178	0.1637
Qwen3-14B	harmcompliance	0	0.9382	0.6683	0.8799	0.7071
Qwen3-14B	harmcompliance	1	0.9393	0.6674	0.9004	0.7670
Qwen3-14B	harmcompliance	2	0.9396	0.6696	0.8844	0.7380
Qwen3-14B	prompt_harm	0	0.9899	0.9912	0.4531	0.1575
Qwen3-14B	prompt_harm	1	0.9896	0.9909	0.4236	0.1253
Qwen3-14B	prompt_harm	2	0.9898	0.9910	0.4541	0.1533

Train/val split is fixed per dataset (val_ids.json). Seeds only change init and minibatch order. Recipe: 30 epochs, lr 1e-3, batch 2048, weight decay 1e-4.

Backbone	Shards	train / val
Qwen3-4B	latent-mas-safety-dataset-seq-qwen3-4b	295,353 / 20,378
Qwen3-4B-Instruct	latent-mas-safety-dataset-seq-qwen3-4b-instruct	285,440 / 30,488
Qwen3-14B	latent-mas-safety-dataset-seq-qwen3-14b	281,285 / 30,072

Sequential eval (K=8, DirectRequest, fpr10)

Harmful: lower ASR is better. Benign: lower refusal is better. n: HarmBench 320 / StrongREJECT 313 / OR-Bench-Hard-1K 1319 / PHTest 2077 harmless + 1192 controversial.

Qwen3-4B harmcompliance @ fpr10

Seed	HarmBench ASR	StrongREJECT ASR	OR-Bench	PHTest harmless	PHTest controversial
0	0.119	0.153	0.079	0.064	0.195
1	0.109	0.128	0.088	0.078	0.202
2	0.109	0.131	0.082	0.075	0.196
mean	0.112	0.137	0.083	0.072	0.198

Qwen3-4B-Instruct harmcompliance @ fpr10

Seed	HarmBench ASR	StrongREJECT ASR	OR-Bench	PHTest harmless	PHTest controversial
0	0.019	0.022	0.083	0.042	0.148
1	0.019	0.000	0.089	0.040	0.143
2	0.016	0.000	0.090	0.044	0.147
mean	0.018	0.007	0.087	0.042	0.146

MFR = (HB + SR + OR + PH_overall) / 4 is 0.0480 ± 0.0027 (sample std).

Instruct prompt_harm @ fpr10 is a topic detector, not a defense: HarmBench 0.000 / StrongREJECT 0.003, but OR-Bench 0.283 and PHTest-controversial 0.582.

Qwen3-14B downstream eval is not in this card yet.

Usage

import torch
from huggingface_hub import hf_hub_download
from LatentMAS.linear_probe import LinearProbe

path = hf_hub_download(
    "Danny-1223/latent-mas-linear-probe",
    "qwen3-4b-instruct/linear_probe_harmcompliance.pt",
)
probe = LinearProbe(path, tau_name="fpr10")

# stages_H: list of 3 tensors, each [K, d_h] = planner / critic / refiner
# 4B / Instruct: d_h=2560; 14B: d_h=5120
verdict = probe.verdict(
    stage_embeddings, latent_steps=8,
    stage_roles=["planner", "critic", "refiner"],
)
if verdict.unsafe:
    # inject the same [SAFETY NOTICE] used by CoT-Guard / the aligner
    ...

Other backbones: qwen3-4b/linear_probe_harmcompliance.pt or qwen3-14b/linear_probe_harmcompliance.pt.

Standalone scoring without the helper:

import torch

ckpt = torch.load("linear_probe_harmcompliance.pt", weights_only=False, map_location="cpu")
x = torch.cat([H.float().mean(0) for H in stages_H], dim=0)
z = (x - ckpt["mu"]) / ckpt["sigma"].clamp_min(1e-6)
score = torch.sigmoid(torch.dot(z, ckpt["weight"].flatten()) + ckpt["bias"]).item()
flag = score >= ckpt["thresholds"]["fpr10"]

In this repo, enable the probe at runtime with yaml linear_probe: true and the matching config.

Retrain from shards:

python -m scripts.aligner.train_linear_probe_multiseed \
  --from-shards data/latent-mas-safety-dataset-seq-qwen3-4b-instruct \
  --out-dir scripts/aligner/artifacts \
  --prefix linear_probe_qwen3_4b_instruct \
  --labels prompt_harm,harmcompliance \
  --seeds 0,1,2

Intended use and limitations

Research baseline for latent multi-agent safety. Each checkpoint is trained on one backbone's latents (d_h and tokenizer / residual stream must match) with a fixed 3-role feature (planner, critic, refiner) and will not transfer across hidden sizes without retraining. Topology transfer is role-aligned, not all-agent concat: extra agents are invisible to the probe. It is not a general-purpose content classifier.

Code: https://github.com/Asatheesh6561/LatentMASHarmBench