Llama3.1-8B-TRACE

TRACE is a trajectory-aware defense model for multi-turn jailbreaks. Instead of judging each user message in isolation, it commits to an explicit safety assessment of the whole conversation so far — in a <STATE> block — and only then writes its reply in an <ANSWER> block. The assessment is part of the generation, so the reply is conditioned on it.

This checkpoint is the paper's TRACE-GRPO model: a LoRA policy trained with group-decoupled GRPO on top of the stage-1 SFT adapter, merged into the base weights for release. Across seven multi-turn attack benchmarks it averages 14.5% ASR, against 31.4% for the strongest baseline and 74.9% for the undefended target, while keeping 93.3% average compliance on over-refusal benchmarks.

• Base model: meta-llama/Llama-3.1-8B-Instruct
• Release form: full merged weights, bf16, single model.safetensors
• Paper / code: TRACE
• Companion model: Dipto084/Qwen3-8B-TRACE (same recipe on Qwen3-8B)

Paper

TRACE: Trajectory Aware Reasoning for Multi-Turn Adversarial Conversation Evaluation — arXiv:2608.15594 · Code: github.com/Dipto084/TRACE

Output format

<STATE>
{
  "cues": {"<cue_name>": "<turn reference and reasoning>"},
  "benign_interpretation":     {"score": 1-10, "explanation": "..."},
  "adversarial_interpretation":{"score": 1-10, "explanation": "..."},
  "jailbreak_score": 1-5,
  "action": "ALLOW | CAUTION | DECLINE",
  "justification": "..."
}
</STATE>
<ANSWER>
<the reply to the current user turn>
</ANSWER>

Only the <ANSWER> block should be shown to an end user. The <STATE> block is an internal reasoning commitment; it is useful for logging, auditing, and routing (e.g. escalate on action == "DECLINE"), but it is not user-facing text.

Usage

The model requires the TRACE system prompt — it defines the cue taxonomy, the scoring rubric, and the output contract above. Without it the model will not emit well-formed <STATE> blocks. The prompt ships with the model as system_prompt.txt; it is the file the attack evaluations load (agents/state_answer_action_prompt.txt in the code repo).

A second variant, system_prompt_or.txt, adds an explicit "harmful vs. harmful-looking" distinction and requires the first sentence of the <ANSWER> to be substantive. It is the prompt used for the PHTest over-refusal measurement; use it when over-refusal on sensitive-but-benign requests matters more than anything else.

The conversation is passed as a single user message, not as a list of chat turns. The model was trained and evaluated on the whole trajectory collapsed into one message with numbered turns; the last turn holds only the user message being answered:

[Turn 1]
USER: ...
ASSISTANT: ...
[Turn 2]
USER: ...

The model replies with <STATE>…</STATE><ANSWER>…</ANSWER>. Show the user only the ANSWER, and append only the ANSWER to the history for the next turn — the STATE block never re-enters the context.

The easiest way to get all of this right is the reference package at github.com/Dipto084/TRACE, which provides the formatting and parsing helpers plus an OpenAI-compatible proxy that lets any client (or attack framework) talk to the model with ordinary chat messages. Doing it by hand:

import re
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "Dipto084/Llama3.1-8B-TRACE"
tok = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, dtype="bfloat16", device_map="auto")
system_prompt = open("system_prompt.txt").read()

def format_trajectory(history, user_message):
    # history: list of (user, answer) pairs already exchanged; answers are ANSWER text only
    lines = []
    for i, (u, a) in enumerate(history, start=1):
        lines += [f"[Turn {i}]", f"USER: {u}", f"ASSISTANT: {a}"]
    lines += [f"[Turn {len(history) + 1}]", f"USER: {user_message}"]
    return "\n".join(lines)

def parse(raw):
    m = re.search(r"<ANSWER>(.*?)</ANSWER>", raw, re.S)
    return m.group(1).strip() if m else re.sub(r"<STATE>.*?</STATE>", "", raw, flags=re.S).strip()

history = []
for user_message in ["first user turn", "second user turn"]:
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": format_trajectory(history, user_message)},
    ]
    ids = tok.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt").to(model.device)
    out = model.generate(ids, max_new_tokens=4096, do_sample=False)
    raw = tok.decode(out[0][ids.shape[-1]:], skip_special_tokens=True)
    answer = parse(raw)
    history.append((user_message, answer))
    print(answer)

Serving with vLLM (the evaluations used greedy decoding, temperature 0):

vllm serve Dipto084/Llama3.1-8B-TRACE --max-model-len 65536 --dtype bfloat16

Budget generously for the completion: the <STATE> block is generated before the answer.

Training

Stage 1 — SFT. The training corpus is built by orchestrating five attack frameworks (Crescendo, ActorAttack, Chain-of-Attacks, ICON, X-Teaming) over 300 harmful behaviors from JailbreakBench and HarmBench, paired with benign multi-turn dialogs and harm-adjacent benign dialogs generated from OR-Bench seeds: 4k adversarial + 2.4k benign + 600 sensitive-but-benign conversations, which expand to 18.2k per-turn trajectories. A frontier LLM annotates every trajectory with the full reasoning trace (cues, dual-hypothesis scores, jailbreak score, action). Stage 1 trains a LoRA adapter (rank 32, alpha 64) on 12.5k of those trajectories, minimizing cross-entropy over the concatenated STATE and ANSWER blocks.

Stage 2 — GRPO with group-decoupled advantages (GDPO). The SFT adapter is trained further with reinforcement learning against a co-located Qwen3-8B-AWQ judge. Standard GRPO normalizes the summed reward across the group, which collapses distinct reward combinations into identical advantages; GDPO instead normalizes each component independently within the group before aggregating into a per-token advantage. Harm-adjacent benign trajectories sit in the RL mixture, so safety and over-refusal are coupled into every gradient step rather than balanced post-hoc.

Reward components:

Component	Grounded in	Signal
R_jb	<STATE>	jailbreak-score accuracy
R_con	<ANSWER>	behavioral consistency: the generated response checked against the ground-truth action
R_cue	<STATE>	cue-set agreement

A structural gate (R_struct) floors the reward at -2.0 when the output does not parse as a well-formed STATE/ANSWER pair.

Hyperparameter	Value
Advantage estimator	GDPO, norm_adv_by_std_in_grpo=False
Reward weights (R_jb, R_con, R_cue)	0.6, 0.6, 0.4
LoRA rank / alpha	32 / 64 (q,k,v,o,gate,up,down)
Learning rate	1e-5
KL loss coefficient	0.03
Entropy coefficient	1e-3
Train batch / PPO mini / micro per GPU	64 / 32 / 2
Rollouts per prompt, temperature	8, 0.9
Max prompt / response length	16384 / 3072
Judge	Qwen3-8B-AWQ (4-bit), vLLM, temperature 0
Hardware	4x H100 80GB

Evaluation

All numbers are from the TRACE paper (Table 2 and Table 3). Every model in the comparison, TRACE included, is built on Llama-3.1-8B-Instruct. TRACE-SFT is the stage-1 ablation — this model without the RL stage; TRACE-GRPO is this checkpoint.

Behavior-level attack success rate (ASR, %) across seven multi-turn attack frameworks; lower is better. FITD and AMA are held-out attacks, not represented in the training corpus.

Defense	X-Teaming	Crescendo	ActorAttack	CoA	ICON	FITD	AMA	Avg
Llama-3.1-8B-Instruct (undefended)	90.8	74.2	45.0	98.3	86.7	80.8	48.3	74.9
Self-Reminder-MT	71.7	25.8	11.7	70.8	61.7	40.8	25.2	44.0
LLaMA-Guard-3-MT	89.2	49.1	20.8	94.2	70.0	25.8	36.5	55.1
X-Guard	58.3	28.3	19.2	79.2	65.0	37.5	23.3	44.4
Red-Queen-Guard	30.8	20.8	9.2	45.8	80.0	47.5	28.3	37.5
NBF-LLM	85.8	60.0	7.5	87.4	7.5	74.2	31.7	50.6
STAIR	44.2	10.0	10.0	42.5	69.2	24.2	20.0	31.4
TRACE-SFT	38.3	40.8	18.3	39.2	2.5	35.0	29.2	29.0
TRACE-GRPO (this model)	20.8	14.2	4.2	21.7	1.7	20.0	19.2	14.5

Lowest ASR on six of the seven attacks, and less than half the average ASR of the strongest baseline (14.5 vs 31.4).

Over-refusal — full-compliance rate (%) on benign prompts; higher is better.

Defense	PHTest	XSTest	Avg
Llama-3.1-8B-Instruct (undefended)	93.2	92.8	93.0
Self-Reminder-MT	59.4	62.8	61.1
LLaMA-Guard-3-MT	88.6	92.0	90.3
X-Guard	83.3	91.6	87.5
Red-Queen-Guard	71.1	86.4	78.8
NBF-LLM	82.5	92.4	87.5
STAIR	44.1	62.0	53.1
TRACE-SFT	85.8	94.4	90.1
TRACE-GRPO (this model)	93.0	93.6	93.3

This is the axis where multi-turn defenses usually pay: the next-best defense by ASR (STAIR, 31.4) complies with only 53.1% of benign prompts, while TRACE-GRPO matches the undefended model's compliance (93.3 vs 93.0) at less than half that ASR.

General capability (%), base vs. the two TRACE stages:

Benchmark	Llama-3.1-8B-Instruct	TRACE-SFT	TRACE-GRPO
ARC-Challenge (25-shot)	81.1	80.7	80.6
BBH (3-shot CoT)	61.5	69.5	68.2
GSM-8K (0-shot)	81.1	79.2	79.6
HellaSwag (10-shot)	80.0	77.8	78.1
MMLU-Pro (5-shot CoT)	43.8	44.8	43.5

Single-turn robustness under AutoDAN-Turbo, a strong single-turn attacker that discovers jailbreak strategies autonomously (a length-1 trajectory under the TRACE formulation). ASR@k is over k independent attempts — lower is better; Avg. Attempts to jailbreak per behavior — higher is better.

Target	ASR@3	ASR@5	ASR@10	Avg. attempts
Llama-3.1-8B-Instruct	67.5	81.7	90.8	3.5
+ TRACE-GRPO (this model)	11.7	21.7	38.3	8.0

Trajectory-aware reasoning does not cost single-turn robustness — it more than doubles the attacker effort per successful jailbreak.

Citation

@article{miah2026trace,
  title   = {TRACE: Trajectory Aware Reasoning for Multi-Turn Adversarial Conversation Evaluation},
  author  = {Miah, Md Messal Monem and Anika, Adrita and Yu, Zhiyuan and Huang, Ruihong},
  journal = {arXiv preprint arXiv:2608.15594},
  year    = {2026}
}

Use of this model is subject to the Llama 3.1 Community License.