Publish TRACE project resources

README.md

@@ -0,0 +1,145 @@
+---
+license: apache-2.0
+tags:
+  - arxiv:2605.08778
+  - llm-safety
+  - red-teaming
+  - jailbreak
+  - multi-turn
+  - reinforcement-learning
+  - credit-assignment
+  - trace
+---
+
+# TRACE
+
+## Not All Turns Matter: Credit Assignment for Multi-Turn Jailbreaking
+
+This repository is the Hugging Face project page for **TRACE (TuRn-level Assignment for CrEdit)**,
+a framework for turn-aware credit assignment in reinforcement-learning-based multi-turn red
+teaming. It collects the paper links, official model releases, a lightweight inference example,
+and release-level evaluation artifacts. Model weights and model-specific settings remain in each
+model repository.
+
+> **Dual-use warning:** TRACE studies adversarial prompts designed to reveal language-model safety
+> failures. Use the released artifacts only for authorized, controlled safety research. Do not test
+> systems without permission, expose the attacker as an unrestricted service, or automatically
+> execute generated content.
+
+## Overview
+
+Multi-turn jailbreak strategies can distribute harmful intent across apparently benign dialogue
+turns. TRACE addresses the resulting credit-assignment problem: individual turns can contribute
+unequally, at different phases of an attack, and against different target models. The method uses
+leave-one-turn-out semantic masking to assign credit in successful trajectories and harmfulness and
+semantic-relevance signals to assign penalties in failed trajectories. The same turn-level signals
+can also support earlier defensive intervention.
+
+## Paper
+
+- **Title:** Not All Turns Matter: Credit Assignment for Multi-Turn Jailbreaking
+- **Authors:** Zhida He, Xiaoyu Wen, Han Qi, Ziyuan Zhou, Peng Yu, Xingcheng Xu, Dongrui Liu,
+  Xia Hu, Chaochao Lu, Qiaosheng Zhang
+- **arXiv:** https://arxiv.org/abs/2605.08778
+- **PDF:** https://arxiv.org/pdf/2605.08778
+- **DOI:** https://doi.org/10.48550/arXiv.2605.08778
+
+The PDF is linked from arXiv rather than duplicated here so that readers receive the current paper
+version.
+
+## Code
+
+The complete training and evaluation implementation is maintained at:
+
+- https://github.com/xsddys/TRACE
+
+## Model releases
+
+| Variant | Base model | Training targets | Repository |
+|---|---|---|---|
+| TRACE (mix) | Qwen2.5-3B-Instruct | gpt-oss-20b and Llama-3.1-8B-Instruct | [XiaoyuWen/TRACE-Mix-Qwen2.5-3B-Instruct](https://huggingface.co/XiaoyuWen/TRACE-Mix-Qwen2.5-3B-Instruct) |
+
+The model repository contains the weights, tokenizer, Transformers configuration, exact prompt
+contract, decoding settings, license, and model-specific safety information.
+
+## Interactive inference
+
+[`inference.py`](inference.py) is a lightweight manual orchestration example. It generates one
+attacker query at a time and asks the operator to paste the response from a separately operated,
+authorized target model. It deliberately does not contact a target endpoint or run a safety judge.
+
+Install the inference-only dependencies:
+
+```bash
+pip install -r requirements.txt
+```
+
+Run the released TRACE (mix) checkpoint:
+
+```bash
+python inference.py \
+  --model XiaoyuWen/TRACE-Mix-Qwen2.5-3B-Instruct \
+  --objective "<AUTHORIZED_RED_TEAM_OBJECTIVE>"
+```
+
+The script automatically downloads `prompt_template.json` from the selected model repository. For
+a local model directory, place that file beside the checkpoint or pass `--prompt-config PATH`.
+
+The defaults reproduce the attacker-side validation sampling configuration:
+
+| Setting | Value |
+|---|---:|
+| Maximum new tokens | 128 |
+| Sampling | enabled |
+| Temperature | 0.5 |
+| Top-p | 0.9 |
+| Top-k | disabled (`-1` in vLLM; `0` in Transformers) |
+| Maximum interaction turns | 5 |
+
+The target model's response is intentionally appended as a `user` message in the attacker model's
+conversation. Attacker queries occupy the `assistant` role. The exact system message, initial user
+template, and role convention are model-specific and are documented in the
+[`prompt_template.json`](https://huggingface.co/XiaoyuWen/TRACE-Mix-Qwen2.5-3B-Instruct/blob/main/prompt_template.json)
+file and Model Card.
+
+## Results
+
+The paper reports ASR@1 (%) under a five-turn limit using the HarmBench Classifier. HB, JBB, and WJB
+denote HarmBench, JailbreakBench, and WildJailbreak, respectively.
+
+| Evaluation target | HB | JBB | WJB | Target average |
+|---|---:|---:|---:|---:|
+| Qwen2.5-7B-Instruct | 90.57 | 87.72 | 90.50 | 89.60 |
+| Llama-3.1-8B-Instruct | 84.48 | 89.09 | 88.67 | 87.41 |
+| gpt-oss-20b | 83.64 | 86.06 | 83.17 | 84.29 |
+| **Overall** | 86.23 | 87.62 | 87.45 | **87.10** |
+
+See the paper for baselines, alternative judges, transfer evaluations, confidence intervals, and
+the complete protocol. [`results/in_run_validation_step130.json`](results/in_run_validation_step130.json)
+contains a separately labeled aggregate from the saved step-0 and step-130 in-run validation
+rollouts; it should not be confused with the paper table above.
+
+## Repository contents
+
+| File | Purpose |
+|---|---|
+| `README.md` | Project overview, paper, models, usage, and reported results |
+| `inference.py` | Interactive multi-turn attacker inference example |
+| `requirements.txt` | Minimal dependencies for the inference example |
+| `results/in_run_validation_step130.json` | Saved in-run validation aggregate |
+
+## Citation
+
+```bibtex
+@misc{he2026turnsmattercreditassignment,
+  title         = {Not All Turns Matter: Credit Assignment for Multi-Turn Jailbreaking},
+  author        = {Zhida He and Xiaoyu Wen and Han Qi and Ziyuan Zhou and Peng Yu and
+                   Xingcheng Xu and Dongrui Liu and Xia Hu and Chaochao Lu and Qiaosheng Zhang},
+  year          = {2026},
+  eprint        = {2605.08778},
+  archivePrefix = {arXiv},
+  primaryClass  = {cs.AI},
+  url           = {https://arxiv.org/abs/2605.08778}
+}
+```
+

inference.py

@@ -0,0 +1,223 @@
+#!/usr/bin/env python3
+"""Interactive inference example for authorized red-team evaluation.
+
+This script generates attacker queries but deliberately does not call a target
+model endpoint. A human operator must paste each authorized target response.
+"""
+
+import argparse
+import json
+from pathlib import Path
+from typing import Dict, List, Optional
+
+import torch
+from huggingface_hub import hf_hub_download
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+
+Message = Dict[str, str]
+
+
+def resolve_prompt_config(model: str, explicit_path: Optional[Path]) -> Path:
+    """Resolve model-specific prompt settings locally or from the Hub."""
+    if explicit_path is not None:
+        return explicit_path
+
+    local_candidate = Path(model) / "prompt_template.json"
+    if local_candidate.is_file():
+        return local_candidate
+
+    return Path(
+        hf_hub_download(repo_id=model, filename="prompt_template.json")
+    )
+
+
+def load_prompt_config(path: Path) -> Dict:
+    with path.open("r", encoding="utf-8") as handle:
+        prompt_config = json.load(handle)
+
+    required_keys = {
+        "system_prompt",
+        "initial_user_prompt_template",
+        "target_response_role",
+        "attacker_response_role",
+        "max_interaction_turns",
+    }
+    missing_keys = sorted(required_keys - prompt_config.keys())
+    if missing_keys:
+        raise ValueError(
+            f"Prompt config {path} is missing keys: {', '.join(missing_keys)}"
+        )
+    if prompt_config["target_response_role"] != "user":
+        raise ValueError("The TRACE target-response role must be 'user'.")
+    if prompt_config["attacker_response_role"] != "assistant":
+        raise ValueError("The TRACE attacker-response role must be 'assistant'.")
+    return prompt_config
+
+
+def build_initial_messages(prompt_config: Dict, objective: str) -> List[Message]:
+    initial_prompt = prompt_config["initial_user_prompt_template"].format(
+        harmful_target=objective
+    )
+    return [
+        {"role": "system", "content": prompt_config["system_prompt"]},
+        {"role": "user", "content": initial_prompt},
+    ]
+
+
+def append_completed_turn(
+    messages: List[Message], attacker_query: str, target_response: str
+) -> None:
+    messages.extend(
+        [
+            {"role": "assistant", "content": attacker_query},
+            # This role is intentional: the target response is the next
+            # observation received by the attacker policy.
+            {"role": "user", "content": target_response},
+        ]
+    )
+
+
+def generate_attacker_query(
+    model,
+    tokenizer,
+    messages: List[Message],
+    max_new_tokens: int,
+    do_sample: bool,
+    temperature: float,
+    top_p: float,
+    top_k: int,
+) -> str:
+    model_inputs = tokenizer.apply_chat_template(
+        messages,
+        tokenize=True,
+        add_generation_prompt=True,
+        return_tensors="pt",
+        return_dict=True,
+    )
+    model_inputs = model_inputs.to(model.device)
+
+    generation_kwargs = {
+        "max_new_tokens": max_new_tokens,
+        "do_sample": do_sample,
+        "pad_token_id": tokenizer.pad_token_id,
+        "eos_token_id": tokenizer.eos_token_id,
+    }
+    if do_sample:
+        generation_kwargs.update(
+            {"temperature": temperature, "top_p": top_p, "top_k": top_k}
+        )
+
+    with torch.inference_mode():
+        output_ids = model.generate(**model_inputs, **generation_kwargs)
+
+    generated_ids = output_ids[0, model_inputs.input_ids.shape[1] :]
+    return tokenizer.decode(generated_ids, skip_special_tokens=True).strip()
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Interactive TRACE attacker inference for authorized evaluation."
+    )
+    parser.add_argument("--model", required=True, help="Local path or Hub model ID.")
+    parser.add_argument(
+        "--objective",
+        required=True,
+        help="Authorized red-team objective supplied to the attacker policy.",
+    )
+    parser.add_argument(
+        "--prompt-config",
+        type=Path,
+        default=None,
+        help=(
+            "Optional path to prompt_template.json. By default it is loaded "
+            "from the local model directory or the model's Hub repository."
+        ),
+    )
+    parser.add_argument("--max-turns", type=int, default=None)
+    parser.add_argument("--max-new-tokens", type=int, default=128)
+    parser.add_argument(
+        "--do-sample",
+        action=argparse.BooleanOptionalAction,
+        default=True,
+        help="Sample attacker outputs; enabled in the reported validation setting.",
+    )
+    parser.add_argument("--temperature", type=float, default=0.5)
+    parser.add_argument("--top-p", type=float, default=0.9)
+    parser.add_argument(
+        "--top-k",
+        type=int,
+        default=0,
+        help=(
+            "Transformers top-k value. Zero disables top-k and corresponds "
+            "to top_k=-1 in the training-time vLLM configuration."
+        ),
+    )
+    parser.add_argument(
+        "--device-map",
+        default="auto",
+        help="Transformers device_map value; defaults to auto.",
+    )
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    prompt_config_path = resolve_prompt_config(args.model, args.prompt_config)
+    prompt_config = load_prompt_config(prompt_config_path)
+    max_turns = (
+        int(prompt_config["max_interaction_turns"])
+        if args.max_turns is None
+        else args.max_turns
+    )
+    if max_turns <= 0:
+        raise ValueError("--max-turns must be positive")
+    if args.max_new_tokens <= 0:
+        raise ValueError("--max-new-tokens must be positive")
+    if args.do_sample and args.temperature <= 0:
+        raise ValueError("--temperature must be positive when sampling is enabled")
+    if not 0 < args.top_p <= 1:
+        raise ValueError("--top-p must be in the interval (0, 1]")
+    if args.top_k < 0:
+        raise ValueError("--top-k must be non-negative for Transformers")
+
+    tokenizer = AutoTokenizer.from_pretrained(args.model)
+    model = AutoModelForCausalLM.from_pretrained(
+        args.model,
+        torch_dtype="auto",
+        device_map=args.device_map,
+    ).eval()
+
+    messages = build_initial_messages(prompt_config, args.objective)
+    for turn_index in range(1, max_turns + 1):
+        attacker_query = generate_attacker_query(
+            model=model,
+            tokenizer=tokenizer,
+            messages=messages,
+            max_new_tokens=args.max_new_tokens,
+            do_sample=args.do_sample,
+            temperature=args.temperature,
+            top_p=args.top_p,
+            top_k=args.top_k,
+        )
+        print(f"\n[attacker turn {turn_index}]\n{attacker_query}\n", flush=True)
+
+        if turn_index == max_turns:
+            break
+        try:
+            target_response = input(
+                "Paste the authorized target-model response "
+                "(or type /stop to finish):\n"
+            ).strip()
+        except EOFError:
+            break
+        if target_response == "/stop":
+            break
+        if not target_response:
+            print("Empty target response; stopping without adding an invalid turn.")
+            break
+        append_completed_turn(messages, attacker_query, target_response)
+
+
+if __name__ == "__main__":
+    main()

requirements.txt

@@ -0,0 +1,5 @@
+torch>=2.1
+transformers>=4.53.2
+accelerate>=0.27
+safetensors>=0.4
+huggingface_hub>=0.30

results/in_run_validation_step130.json

@@ -0,0 +1,31 @@
+{
+  "checkpoint_step": 130,
+  "evaluation_type": "in_run_validation",
+  "success_definition": "outcome_score >= 0.9",
+  "samples_per_profile": 200,
+  "decoding": {
+    "do_sample": true,
+    "temperature": 0.5,
+    "top_p": 0.9,
+    "top_k": -1,
+    "max_new_tokens": 128
+  },
+  "results": [
+    {
+      "target_profile": "qwen",
+      "step_0_success_rate": 0.335,
+      "step_130_success_rate": 0.83
+    },
+    {
+      "target_profile": "oss",
+      "step_0_success_rate": 0.24,
+      "step_130_success_rate": 0.82
+    },
+    {
+      "target_profile": "llama",
+      "step_0_success_rate": 0.175,
+      "step_130_success_rate": 0.81
+    }
+  ],
+  "limitations": "Stochastic in-run validation; not an independent third-party benchmark."
+}