Membrane — a self-evolving safety guard

An external guard that learns from adversarial attacks at test time, protecting an unguarded target model from jailbreaks through its self-evolving CSM (Contrastive Safety Memory) store.

📄 Paper: arxiv.org/abs/2606.05743

• Model safety on HarmBench and agent safety on AgentHarm from a single CLI — two methods, vanilla (no guard) and membrane (the CSM guard).
• Paired Self-Evolution (PSE) — the CSM evolves by pairing each harmful prompt with its safe counterpart and distilling the contrast into one memory cell. A recall-preserving retrieval critic — off during training, on at eval — keeps APPLIES/MAYBE candidates.
• Agent-safety extras (AgentHarm only): a safe-side veto and a semantic action context derived from the agent's tool list (always on for agent safety; inert for model safety).
• Test-time adaptation (--tta) — keep evolving the CSM during eval using the label-free DecisionReviewer, with no test-time labels.

Compute prerequisite. The backend is a self-hosted vLLM server (Qwen/Qwen3-8B), so you need a CUDA-capable GPU host. The CLI client runs anywhere (CPU only) and points at the server via VLLM_SERVER_URL.

---

1. How it runs: client + model servers

Membrane separates the client (mbr ..., CPU-only) from the model server(s) (vLLM, CUDA GPU on Linux). The client reaches the servers over HTTP via the URLs in .env, so you can drive the CLI from a laptop while a remote GPU box serves the models.

Server	Role	Required?	Launch script	Example port
Chat vLLM (Qwen/Qwen3-8B)	guard · target · eval judge	always	scripts/setup/start_vllm.sh	8077
Embedding model (embeddinggemma-300m)	CSM retrieval vectors	always	scripts/setup/start_vllm_pooling.sh	8078
Decision Reviewer (e.g. gemma-3-27b-it)	label-free judge for --tta	only for TTA	scripts/setup/start_vllm.sh	8079

One GPU can host the chat server plus the low-memory embedding server together; multiple GPUs split with -g. For a quick CPU-only smoke test, set EMBEDDING_BACKEND=local and skip the embedding server.

---

2. Install

This project uses uv for environments and dependencies (Python ≥ 3.13, fetched automatically).

# Install uv (macOS / Linux)
curl -LsSf https://astral.sh/uv/install.sh | sh
# Windows (PowerShell):  irm https://astral.sh/uv/install.ps1 | iex

Then sync the dependencies for what you'll run:

# (A) Client only — drive the CLI, models served on a remote GPU host
uv sync --extra agentharm

# (B) GPU host — also serve vLLM locally; add tests
uv sync --extra serve --extra agentharm --extra test

Extra	Adds	When you need it
(base)	CLI · guard · memory	always
serve	vllm	to start a vLLM server on this machine
agentharm	inspect-ai, …	to run mbr agent (AgentHarm) experiments
test	pytest	to run uv run pytest tests/

source .venv/bin/activate   # once per shell; to skip it, prefix commands with `uv run`

Activating puts the mbr CLI on your PATH. Every mbr ... command below assumes it — without activation, prefix them with uv run (uv run mbr model eval ...).

---

3. Configure (.env)

cp .env.example .env

The keys a first-time runner must check (this is the paper setup):

# Chat vLLM address — must match the port you start the server on (§4.1)
VLLM_SERVER_URL=http://localhost:8077/v1

# Roles: vLLM for target/guard; the eval judge is Gemini (a hosted model).
# Set EVAL_MODEL=vllm:Qwen/Qwen3-8B to stay fully local instead.
TARGET_MODEL=vllm:Qwen/Qwen3-8B
GUARD_MODEL=vllm:Qwen/Qwen3-8B
EVAL_MODEL=gemini-3-flash-preview

# Gateway for the hosted judge above (OpenAI-compatible; required for EVAL_MODEL=gemini-*)
OPENAI_API_KEY=your_api_key_here
OPENAI_BASE_URL=https://your-gateway.example.com/v1

# Embedding — served on its own vLLM pooling server (§4.2)
EMBEDDING_MODEL=google/embeddinggemma-300m
EMBEDDING_BACKEND=vllm
EMBEDDING_VLLM_URL=http://localhost:8078/v1

# TTA only: the label-free reviewer model + its server (§4.3)
DECISION_REVIEWER_MODEL=vllm:google/gemma-3-27b-it
DECISION_REVIEWER_VLLM_URL=http://localhost:8079/v1

OPENAI_API_KEY / OPENAI_BASE_URL are only consulted for hosted models — leave them unset if every role is vLLM-served. Each *_VLLM_URL lets a component sit on its own port and falls back to VLLM_SERVER_URL if left empty. Full reference: docs/configuration.md.

---

4. Start the vLLM server(s)

Run these on a GPU Linux host with --extra serve installed. Skip this section if you point VLLM_SERVER_URL at a server someone else is running.

4.1 Chat server (required) — Qwen3-8B

One server handles all three chat roles (guard · target · eval judge).

./scripts/setup/start_vllm.sh                          # Qwen3-8B on :8077
./scripts/setup/start_vllm.sh -g 0,1 -t 2              # multi-GPU, tensor parallel

The launch port must match VLLM_SERVER_URL. Sampling is applied server-side by this script (Qwen3 model-card presets); the client never passes sampling params.

4.2 Embedding server (required)

The embedder runs on a vLLM pooling server so parallel workers share one GPU-resident copy:

./scripts/setup/start_vllm_pooling.sh google/embeddinggemma-300m -p 8078

It only needs ~20% of a GPU, so it can co-locate with the chat server (start_vllm.sh --gpu-mem 0.7). For a quick CPU-only smoke you can instead set EMBEDDING_BACKEND=local and skip this server.

4.3 Decision Reviewer server (required for TTA only — §6.4)

TTA judges without labels using a separate reviewer model — kept off the guard's model to avoid correlated blind spots.

./scripts/setup/start_vllm.sh google/gemma-3-27b-it -p 8079

Setting the reviewer to the guard model (vllm:Qwen/Qwen3-8B) on the chat server is allowed (then leave DECISION_REVIEWER_VLLM_URL empty). There is no silent fallback: if the configured reviewer isn't served, a --tta run aborts.

---

5. Build the dataset splits

The tracked data/<bench>/{warmup,eval}_*.jsonl splits are bundled, but regenerate them once to be safe (idempotent):

uv run python scripts/data/harmbench_splits.py     # HarmBench (model safety)
uv run python scripts/data/agentharm_splits.py     # AgentHarm (agent safety)

Each split interleaves harmful/benign on disk (h0, b0, h1, b1, …), so --sample-size N selects N per label (N harmful + N benign) from the prefix. Layout and rationale: docs/splits.md.

---

6. Reproduce the results

Each experiment is shown as (1) a quick single run to confirm the pipeline works, then (2) the full sweep that produces the paper-scale numbers. Every sweep accepts --dry-run (or DRY_RUN=1) to print the planned commands without executing.

• train evolves memory only (no metrics); eval runs inference on a frozen store and reports metrics.
• membrane = guard + CSM; vanilla = the no-defense reference.
• Each eval prints a headline table and writes …/<dataset>/summary.json (aggregate metrics + run metadata) plus per-sample sample_*.json. Results land under outputs/{model_safety|agent_safety}/{namespace}/... — main for headline sweeps, tta/reviewer for --tta runs.

6.1 Model safety — HarmBench

(1) Quick single run

mbr model train --method membrane --memory-dir shared/memory --sample-size 7
mbr model eval  --method membrane --memory-dir shared/memory --sample-size 10

--method membrane eval reads the store built by train via --memory-dir. The no-defense reference is --method vanilla (no --memory-dir).

(2) Full sweep — membrane + vanilla × 7 attacks × warmup sizes {7, 14, 28, 80}:

./scripts/sweeps/harmbench.sh --port 8077

# Subset: only membrane, specific attacks/sizes
TARGETS=membrane METHODS="original pap" SIZES="28 80" ./scripts/sweeps/harmbench.sh

6.2 Agent safety — AgentHarm

Unlike HarmBench, the tool-list-driven safe-side veto and semantic action context turn on automatically. Needs --extra agentharm.

(1) Quick single run (--sample-size is per label)

mbr agent train --method membrane --memory-dir shared/memory --sample-size 8
mbr agent eval  --method membrane --memory-dir shared/memory --sample-size 10

(2) Full sweep

./scripts/sweeps/agentharm.sh --port 8077

# Subset
TARGETS=membrane SIZES="8 16 24 32" ./scripts/sweeps/agentharm.sh

AgentHarm selects the jailbreak variant with --attack-method (not --dataset): default original, plus autodan-turbo tap pair pap renellm flipattack. Eval concurrency is set by VLLM_MAX_CONNECTIONS (default 32).

6.3 Warmup-size learning curve

The learning curve is the SIZES axis of the sweeps above — one trained store + eval per warmup size (the defaults do 7 14 28 80 for HarmBench, 8 16 24 32 for AgentHarm). Plot defense / over-refusal vs warmup size from each size's summary.json.

6.4 Test-time adaptation (TTA)

Keep adapting the CSM on the eval stream itself, with no test-time labels — the label-free DecisionReviewer supplies the verdict. Prereqs: chat + embedding servers and the Decision Reviewer server (§4.3).

(1) Single run — add --tta to a membrane eval (auto-routes to the tta/reviewer namespace):

mbr model eval --method membrane --tta --memory-dir shared/memory
mbr agent eval --method membrane --tta --memory-dir shared/memory

(2) Sweep — add the reviewer leg to the headline sweep with TTA=1:

TTA=1 ./scripts/sweeps/harmbench.sh --port 8077
TTA=1 ./scripts/sweeps/agentharm.sh --port 8077

The frozen baseline stays in main/; the adaptive leg lands in outputs/<regime>/tta/reviewer/..., so the two sit side-by-side. TTA eval is sequential (memory writes are racy). Design details: docs/memory-system.md.

---

How it works

A user query first reaches the guard, which retrieves matching cells from the CSM store, optionally filters them through the retrieval critic, and emits SAFE or UNSAFE. SAFE queries fall through to the vanilla target model; UNSAFE queries are refused with a generated explanation. During training the loop continues — this is Paired Self-Evolution: each harmful prompt is paired with its safe counterpart, a judge labels each verdict (JAILBROKEN / OVER_REFUSAL / CORRECT), a memory planner picks an action (CREATE / UPDATE / DELETE / SKIP), and the writer distills the pair into one memory cell. Eval reuses the same pipeline with read-only memory (or keeps evolving under --tta).

Pipeline diagrams and the memory design rationale live in docs/architecture.md and docs/memory-system.md.

Repository layout

src/                  CLI, config, schema, LLM client, top-level utilities
  frameworks/         inference (guard) · evolution (learning) · evaluation (judge)
  memory/             ChromaDB-backed safety memory store
  benchmarks/         AgentHarm benchmark adapter
  evaluation/         orchestrator, data loader, metrics, results I/O
prompts/              YAML / Jinja2 templates for every LLM role
policies/             Safety taxonomies and attack-mechanism definitions
data/                 Jailbreak datasets and generated splits
scripts/{setup,data,sweeps}/   vLLM launchers · split generation · sweep runners
tests/                pytest suite
docs/                 Long-form architecture, memory, configuration, splits

Documentation

File	Contents
docs/architecture.md	Full pipeline diagrams for vanilla and membrane, training and eval.
docs/memory-system.md	CSM design — memory cells, retrieval, retrieval critic, evolution, TTA.
docs/configuration.md	All .env and CLI knobs with defaults and rationale.
docs/splits.md	Dataset partitioning rationale and reproducibility notes.

Citation

If you use Membrane, please cite the paper:

@article{choi2026membrane,
  title={Membrane: A Self-Evolving Contrastive Safety Memory for LLM Agent Defense},
  author={Choi, Minseok and Yang, Seungbin and Kim, Dongjin and Kim, Subin and Son, Jungmin and Lee, Yunseung and Choo, Jaegul and Kwak, Youngjun},
  journal={arXiv preprint arXiv:2606.05743},
  year={2026}
}