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:

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)

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.

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

Qwen3-4B-Instruct harmcompliance @ fpr10

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