# DRAFT: Task Decoupled Latent Reasoning for Agent Safety

arXiv: https://arxiv.org/abs/2604.03242 · HTML: https://mlsanity.com/paper/2604.03242/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Placed in Reasoning for agents, planning and code, one of 17 directions in LLM reasoning.
- Builds on 36 of its 41 references that are in the map.
- First posted on arXiv 11 Feb 2026.

## Placed in

- LLM reasoning / [Reasoning for agents, planning and code](https://mlsanity.com/lr/reasoning-for-agents-planning-and-code/)
- LLM agents and tool use / [Agent threats, guardrails and governance](https://mlsanity.com/agents/agent-threats-guardrails-and-governance/)

## Builds on (most cited first, 30 of 36)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [LoRA: Low-Rank Adaptation of Large Language Models](https://mlsanity.com/paper/2106.09685/) (2021)
- [Chain-of-Thought Prompting Elicits Reasoning in Large Language Models](https://mlsanity.com/paper/2201.11903/) (2022)
- [ReAct: Synergizing Reasoning and Acting in Language Models](https://mlsanity.com/paper/2210.03629/) (2022)
- [Self-Consistency Improves Chain of Thought Reasoning in Language Models](https://mlsanity.com/paper/2203.11171/) (2022)
- [Toolformer: Language Models Can Teach Themselves to Use Tools](https://mlsanity.com/paper/2302.04761/) (2023)
- [Tree of Thoughts: Deliberate Problem Solving with Large Language Models](https://mlsanity.com/paper/2305.10601/) (2023)
- [Sequence to Sequence Learning with Neural Networks](https://mlsanity.com/paper/1409.3215/) (2014)
- [TruthfulQA: Measuring How Models Mimic Human Falsehoods](https://mlsanity.com/paper/2109.07958/) (2021)
- [Voyager: An Open-Ended Embodied Agent with Large Language Models](https://mlsanity.com/paper/2305.16291/) (2023)
- [WebGPT: Browser-assisted question-answering with human feedback](https://mlsanity.com/paper/2112.09332/) (2021)
- [The Rise and Potential of Large Language Model Based Agents: A Survey](https://mlsanity.com/paper/2309.07864/) (2023)
- [Least-to-Most Prompting Enables Complex Reasoning in Large Language Models](https://mlsanity.com/paper/2205.10625/) (2022)
- [OpenHands: An Open Platform for AI Software Developers as Generalist Agents](https://mlsanity.com/paper/2407.16741/) (2024)
- [Deep Variational Information Bottleneck](https://mlsanity.com/paper/1612.00410/) (2016)
- [Ignore Previous Prompt: Attack Techniques For Language Models](https://mlsanity.com/paper/2211.09527/) (2022)
- [Training Large Language Models to Reason in a Continuous Latent Space](https://mlsanity.com/paper/2412.06769/) (2024)
- [Identifying the Risks of LM Agents with an LM-Emulated Sandbox](https://mlsanity.com/paper/2309.15817/) (2023)
- [WildGuard: Open One-Stop Moderation Tools for Safety Risks, Jailbreaks, and Refusals of LLMs](https://mlsanity.com/paper/2406.18495/) (2024)
- [ToxiGen: A Large-Scale Machine-Generated Dataset for Adversarial and Implicit Hate Speech Detection](https://mlsanity.com/paper/2203.09509/) (2022)
- [Agent Security Bench (ASB): Formalizing and Benchmarking Attacks and Defenses in LLM-based Agents](https://mlsanity.com/paper/2410.02644/) (2024)
- [ReST-MCTS*: LLM Self-Training via Process Reward Guided Tree Search](https://mlsanity.com/paper/2406.03816/) (2024)
- [Agent-SafetyBench: Evaluating the Safety of LLM Agents](https://mlsanity.com/paper/2412.14470/) (2024)
- [Quiet-STaR: Language Models Can Teach Themselves to Think Before Speaking](https://mlsanity.com/paper/2403.09629/) (2024)
- [R-Judge: Benchmarking Safety Risk Awareness for LLM Agents](https://mlsanity.com/paper/2401.10019/) (2024)
- [ShieldGemma: Generative AI Content Moderation Based on Gemma](https://mlsanity.com/paper/2407.21772/) (2024)
- [SafetyBench: Evaluating the Safety of Large Language Models](https://mlsanity.com/paper/2309.07045/) (2023)
- [AEGIS: Online Adaptive AI Content Safety Moderation with Ensemble of LLM Experts](https://mlsanity.com/paper/2404.05993/) (2024)
- [Adaptive Attacks Break Defenses Against Indirect Prompt Injection Attacks on LLM Agents](https://mlsanity.com/paper/2503.00061/) (2025)

## Built on by (most cited first, 1 of 1)

- [LPG: Balancing Efficiency and Policy Reasoning in Latent Policy Guardrails](https://mlsanity.com/paper/2605.17329/) (2026)

## Abstract (arXiv)

The advent of tool-using LLM agents shifts safety monitoring from output moderation to auditing long, noisy interaction trajectories, where risk-critical evidence is sparse-making standard binary supervision poorly suited for credit assignment. To address this, we propose DRAFT (Task Decoupled Latent Reasoning for Agent Safety), a latent reasoning framework that decouples safety judgment into two trainable stages: an Extractor that distills the full trajectory into a compact continuous latent draft, and a Reasoner that jointly attends to the draft and the original trajectory to predict safety. DRAFT avoids lossy explicit summarize-then-judge pipelines by performing evidence aggregation in latent space, enabling end-to-end differentiable training.Across benchmarks including ASSEBench and R-Judge, DRAFT consistently outperforms strong baselines, improving accuracy from 63.27% (LoRA) to 91.18% averaged over benchmarks, and learns more separable representations. Ablations demonstrate a clear synergy between the Extractor and the Reasoner.Overall, DRAFT suggests that continuous latent reasoning prior to readout is a practical path to robust agent safety under long-context supervision with sparse evidence.