latent-mas-attn-aligner

LatentMAS Attention Aligner

A latent-space safety aligner for multi-agent systems (MAS) whose agents communicate in latent space rather than natural language.

In a latent MAS, inter-agent messages are hidden states, so text-based moderation cannot read them without first decoding every latent back to tokens. This aligner reads the pre-Judger latents directly and maps them into the representation space of a frozen Llama-Guard classifier tail, producing p_unsafe for the communication. A deployment thresholds p_unsafe at a chosen tau and, when exceeded, injects a safety notice before the final aggregating agent.

Architecture

Because the pooling attends over a variable number of latent tokens, one checkpoint accepts any agent count N without retraining or reshaping — the same weights run 3-agent chains and 9-agent branch-and-fuse DAGs.

Files

• aligner.pt — {d_a, d_g, hidden_dim, pool, state_dict}

The frozen Llama-Guard tail is not included; supply your own bundle.

Usage

import torch
from scripts.aligner.aligner_loader import build_aligner_from_bundle

bundle = torch.load("aligner.pt", weights_only=False, map_location="cpu")
aligner, arch, meta = build_aligner_from_bundle(bundle)

# h_a: [B, K, 2560] pre-Judger latents; head: frozen Llama-Guard tail
p_unsafe = torch.softmax(head(aligner(h_a)), dim=-1)[:, 1]
inject_notice = p_unsafe > tau

Choosing tau

tau sets the operating point on the harm-blocking / over-refusal tradeoff; it is not a fixed property of the model. Pick it from the quantiles of p_unsafe on your own benign traffic — e.g. the 95th percentile targets a 5% flag rate on benign prompts. Thresholds calibrated on a validation split do not transfer reliably to deployment, since the score distributions differ.

Intended use and limitations

Research artifact for studying safety in latent multi-agent systems. It is trained on Qwen3-4B latents at d_a=2560 and will not transfer to models with a different hidden size or latent convention without retraining. It is not a general-purpose content classifier and has not been evaluated as a standalone moderation system.

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

Required companion file: guard_head_bundle.pt

The checkpoint in this repo is only the aligner head (~0.11–0.15 GB). It projects a LatentMAS run's pre-Judger latents h_a into a frozen slice of Llama-Guard-3-8B, and that slice is what turns the projection into a p_unsafe score. Without it the aligner cannot be evaluated.

The slice is not published here because it is ~7.6 GB and contains no trained parameters — every tensor in it comes straight from meta-llama/Llama-Guard-3-8B:

Rebuild it

Requires a GPU and access to the gated Llama-Guard-3-8B repo:

python -m scripts.aligner.extract_cli \
  --mode dump-head-bundle \
  --layer 16 \
  --out data/aligner/phase0_full/guard_head_bundle.pt \
  --device cuda:0

This is a one-time bootstrap; the resulting bundle is reused by every aligner (all models, all seeds, all taus). Point aligner_guard_head at it in the model config, alongside aligner_path for the checkpoint from this repo.

Why the split

The aligner is deliberately small: it learns only the projection from latent space into the guard's representation, while the guard's own harm-classification weights stay frozen. That keeps the trained artifact ~0.1 GB and means a single guard bundle serves every model in the collection.