
Research implementation for mitigating adaptive prompt injections via on-policy distillation, with training recipes and evaluators for SEP, PISmith, and AgentDojo benchmarks.
Yibo Peng · Long Lian · David Wagner† · Sizhe Chen†
† Joint supervision.
[Paper] [Project page] [Model]
This release implements the paper's final full-response KL formulation,
also described as the no-parsing variant. The student rolls out under an
attacked context. A clean-context teacher, initialized from the same base
model, scores the student's sampled tokens under the paired clean context.
The reverse-KL signal is applied to every sampled response token, including
reasoning tokens; no </think> boundary is used to select a supervision span.
training/tinker/: the full-response KL training recipe and exact paper
configuration.scripts/: data preparation and Tinker checkpoint export utilities.evaluation/sep/: SEP static and adaptive evaluators.evaluation/pismith/: PISmith training/evaluation overlay and launchers.evaluation/agentdojo/: AgentDojo utility and attack runner.evaluation/lm_eval/: MMLU-Pro, GPQA Diamond, GSM8K, and Minerva MATH runner.docs/REPRODUCIBILITY.md: end-to-end commands and experimental settings.cp .env.example .env
python scripts/prepare_data.py
bash training/tinker/setup.sh
bash training/tinker/train_full_kl_qwen36.sh
The setup script checks out the pinned upstream Tinker Cookbook commit and applies the source overlay in this repository. It does not modify another Tinker checkout on the machine.
Tinker checkpoints must be converted with the Tinker Cookbook weight mapper:
python scripts/export_tinker_adapter.py \
--tinker-path 'tinker://RUN_ID:train:0/weights/final' \
--output-dir adapters/secopd-full-kl
Do not merge a raw Tinker adapter with vanilla PEFT. Qwen3.6 linear-attention
and lm_head key names require the Tinker mapping implemented by
tinker_cookbook.weights.
See docs/REPRODUCIBILITY.md for evaluation commands, hardware notes, and the exact frozen configuration.
No API keys, model weights, generated outputs, private paths, or experiment
logs are tracked. Credentials are read from environment variables and runtime
judge configuration is written only under the ignored runtime/ directory.
The training overlay targets Thinking Machines Lab's Tinker Cookbook. SEP evaluation is derived from Meta-SecAlign, PISmith evaluation targets PISmith, and AgentDojo evaluation targets AgentDojo. Exact upstream revisions and licenses are listed in docs/THIRD_PARTY.md.
@article{peng2026secopd,
title = {{SecOPD}: Mitigating Adaptive Prompt Injections by On-Policy Distillation},
author = {Peng, Yibo and Lian, Long and Wagner, David and Chen, Sizhe},
journal = {arXiv preprint arXiv:2608.21500},
year = {2026}
}