# COMPASS: Cognitive MCTS-Guided Process Alignment for Safe Search Agents

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Agent threats, guardrails and governance, one of 16 directions in LLM agents and tool use.
- Builds on 35 of its 55 references that are in the map.
- First posted on arXiv 29 May 2026.

## Placed in

- 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 35)

- [Training language models to follow instructions with human feedback](https://mlsanity.com/paper/2203.02155/) (2022)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks](https://mlsanity.com/paper/2005.11401/) (2020)
- [Direct Preference Optimization: Your Language Model is Secretly a Reward Model](https://mlsanity.com/paper/2305.18290/) (2023)
- [Let's Verify Step by Step](https://mlsanity.com/paper/2305.20050/) (2023)
- [Finetuned Language Models Are Zero-Shot Learners](https://mlsanity.com/paper/2109.01652/) (2021)
- [HotpotQA: A Dataset for Diverse, Explainable Multi-hop Question Answering](https://mlsanity.com/paper/1809.09600/) (2018)
- [TriviaQA: A Large Scale Distantly Supervised Challenge Dataset for Reading Comprehension](https://mlsanity.com/paper/1705.03551/) (2017)
- [Self-RAG: Learning to Retrieve, Generate, and Critique through Self-Reflection](https://mlsanity.com/paper/2310.11511/) (2023)
- [Search-R1: Training LLMs to Reason and Leverage Search Engines with Reinforcement Learning](https://mlsanity.com/paper/2503.09516/) (2025)
- [Improving language models by retrieving from trillions of tokens](https://mlsanity.com/paper/2112.04426/) (2021)
- [Measuring and Narrowing the Compositionality Gap in Language Models](https://mlsanity.com/paper/2210.03350/) (2022)
- [Math-Shepherd: Verify and Reinforce LLMs Step-by-step without Human Annotations](https://mlsanity.com/paper/2312.08935/) (2023)
- [Interleaving Retrieval with Chain-of-Thought Reasoning for Knowledge-Intensive Multi-Step Questions](https://mlsanity.com/paper/2212.10509/) (2022)
- [DeepSeek-V3.2: Pushing the Frontier of Open Large Language Models](https://mlsanity.com/paper/2512.02556/) (2025)
- [Tree of Attacks: Jailbreaking Black-Box LLMs Automatically](https://mlsanity.com/paper/2312.02119/) (2023)
- [Safe RLHF: Safe Reinforcement Learning from Human Feedback](https://mlsanity.com/paper/2310.12773/) (2023)
- [Safety Alignment Should Be Made More Than Just a Few Tokens Deep](https://mlsanity.com/paper/2406.05946/) (2024)
- [LongLLMLingua: Accelerating and Enhancing LLMs in Long Context Scenarios via Prompt Compression](https://mlsanity.com/paper/2310.06839/) (2023)
- [Adaptive-RAG: Learning to Adapt Retrieval-Augmented Large Language Models through Question Complexity](https://mlsanity.com/paper/2403.14403/) (2024)
- [Step-DPO: Step-wise Preference Optimization for Long-chain Reasoning of LLMs](https://mlsanity.com/paper/2406.18629/) (2024)
- [RECOMP: Improving Retrieval-Augmented LMs with Compression and Selective Augmentation](https://mlsanity.com/paper/2310.04408/) (2023)
- [FlashRAG: A Modular Toolkit for Efficient Retrieval-Augmented Generation Research](https://mlsanity.com/paper/2405.13576/) (2024)
- [Safety Tax: Safety Alignment Makes Your Large Reasoning Models Less Reasonable](https://mlsanity.com/paper/2503.00555/) (2025)
- [Process Reward Models That Think](https://mlsanity.com/paper/2504.16828/) (2025)
- [STAIR: Improving Safety Alignment with Introspective Reasoning](https://mlsanity.com/paper/2502.02384/) (2025)
- [Process vs. Outcome Reward: Which is Better for Agentic RAG Reinforcement Learning](https://mlsanity.com/paper/2505.14069/) (2025)
- [RedAgent: Red Teaming Large Language Models with Context-aware Autonomous Language Agent](https://mlsanity.com/paper/2407.16667/) (2024)
- [s3: You Don't Need That Much Data to Train a Search Agent via RL](https://mlsanity.com/paper/2505.14146/) (2025)
- [Equilibrate RLHF: Towards Balancing Helpfulness-Safety Trade-off in Large Language Models](https://mlsanity.com/paper/2502.11555/) (2025)

## Abstract (arXiv)

LLM-powered search agents enable multi-step reasoning and tool use. However, these capabilities introduce retrieval-induced safety degradation, as harmful intents may decompose into seemingly innocuous sub-queries that lead to unsafe outcomes. Existing alignment methods struggle to capture sparse safety signals and fail to supervise diverse violations across multi-step interactions. We propose COMPASS, a Cognitive MCTS-Guided Process Alignment framework designed to achieve robust safety alignment throughout the agent workflow while preserving general utility. COMPASS integrates cognitive tree exploration (CTE) to efficiently synthesize stealthy attack trajectories, and introspective step-wise alignment (ISA) to isolate risky intermediate actions for fine-grained process supervision. Empirical results show that COMPASS achieves a favorable safety-utility trade-off while requiring substantially less training data.