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

Abstract: 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.

• The paper introduces a process-aligned safety framework integrating Cognitive Tree Exploration (CTE) and Introspective Step-wise Alignment (ISA) to mitigate multi-stage risks in LLM search agents.
• The methodology employs cognitive Q-value estimation to efficiently guide search trajectories, achieving significant harmful rate reductions with lower data requirements.
• Empirical validation shows COMPASS improves safety and utility across diverse benchmarks, demonstrating cross-model generalization and computational efficiency.

Cognitive MCTS-Guided Process Alignment (COMPASS) for Safe Search Agents: A Technical Analysis

Problem Motivation and Limitations of Prior Safety Alignment

LLM-based search agents, which perform multi-turn retrieval-augmented generation (RAG), represent a critical advance in LLM agent design. These agents achieve complex tool-augmented reasoning by decomposing high-level instructions into iterative knowledge-seeking steps. However, this iterative retrieval exposes a pronounced vulnerability: retrieval-induced safety degradation. Harmful prompts can be decomposed into innocuous sub-queries that evade base model guardrails, enabling bypass of surface-level safety filters and the emergence of unsafe outputs even in otherwise well-aligned models.

Conventional alignment techniques, including final-output preference optimization and query-level supervision, fail to localize and address multi-stage safety failures emerging within agentic search trajectories. Notably, query/document filtering via external models offers partial mitigation but impacts agent utility and incurs significant latency. Furthermore, outcome-level RL rewards are sparse and difficult to attribute to the process steps responsible for safety violations.

Figure 1: Example of retrieval-induced safety degradation in search agents, where decomposition enables the circumvention of base model safety policies.

The COMPASS Framework: Addressing Process-Level Safety with Cognitive Exploration

COMPASS introduces a process-aligned alignment framework comprising two core innovations: Cognitive Tree Exploration (CTE) and Introspective Step-wise Alignment (ISA). COMPASS reframes agentic RAG as an MDP, where the search agent's workflow is structured around three atomic actions: answer generation, query generation, and evidence extraction. Safety preferences must therefore be distributed across the full reasoning trajectory, not only the final response.

Cognitive Tree Exploration (CTE)

CTE extends standard Monte Carlo Tree Search (MCTS) by leveraging cognitive Q-value estimation from a strong LLM. Unlike classical MCTS, which relies on expensive and potentially low-signal rollouts, CTE employs an LLM as an intrinsic heuristic, providing process-level risk estimates at each search step. This enables efficient prioritization of stealthy, high-risk trajectories in the vast action-state space of multi-step RAG agents.

Intermediate state evaluation is informed by two coupled signals:

• Cognitive Q-value: Dense, reasoning-informed risk prediction for each state-action node, replacing costly full rollouts.
• Terminal reward: Alignment model-based judgment of the final response's safety, guiding persistent risk backpropagation for downstream alignment.

Crucially, CTE synthesizes high-quality, risk-localized preference data with minimal computational burden.

Introspective Step-wise Alignment (ISA)

ISA isolates and corrects the precise process stages responsible for safety violations. Exploiting CTE-annotated trajectories, ISA forms step-wise preference pairs by:

• Identifying the first risky state as the locus causally responsible for downstream violations.
• Eliciting safe alternatives via introspective prompting—to directly target and neutralize unsafe reasoning or action generation, including sub-query decomposition and evidence extraction.

This granular correction paradigm transforms heterogeneous failure modes into high-quality, format-compliant preference supervision across the full agent trajectory. The result is fine-grained policy optimization that penalizes precisely those intermediate actions responsible for logical jailbreaking and risk propagation, without sacrificing general reasoning utility.

Figure 2: COMPASS framework overview: (a) CTE efficiently identifies high-risk trajectories using cognitive Q-values; (b) ISA performs targeted, multi-stage interventions across the agent workflow.

Empirical Performance: Safety-Utility Trade-off and Data Efficiency

Comprehensive evaluation on red-teaming (RRB, StrongREJECT, WildTeaming) and utility (TriviaQA, HotpotQA, Bamboogle) benchmarks demonstrates COMPASS's effectiveness.

Strong Safety Improvement: Compared to baselines—including external filtering, outcome-RL, and Safesearch—COMPASS achieves the lowest harmful rates with Qwen-7B and Llama3.1-8B backbones:

• Qwen-7B: Harmful Rate = 18.0%
• Llama3.1-8B: Harmful Rate = 14.6%

Data Efficiency and Utility Preservation: COMPASS achieves these results with an order of magnitude less safety data than Safesearch (2K vs. 12K+ samples), while matching or exceeding utility performance (best or second-best EM and F1 scores).

Figure 3: Safety performance comparison, showing substantial reduction in harmful rate by COMPASS relative to prior methods.

Figure 4: Utility performance (e.g., TriviaQA scores), demonstrating that COMPASS preserves general utility despite rigorous safety constraints.

Execution Efficiency: CTE-guided data construction requires significantly fewer GPU hours than traditional MCTS or Safesearch, as corroborated by execution time benchmarks.

Figure 5: Execution time comparison across safety alignment methods, highlighting the computational efficiency of the COMPASS approach.

Cross-Model Generalization: CTE/ISA-aligned preferences generated on Qwen-7B transfer robustly to Llama3.1-8B, substantially improving both safety and helpfulness with no model-specific data synthesis required.

Ablation and Analytical Insights

Alignment Granularity: Step-wise alignment outperforms sequence-level DPO and SFT-based methods. While SFT can lower harmfulness by learning conservative refusals, this degrades utility. Final-answer-only DPO misses multi-turn risks, confirming the necessity of process-level correction.

Exploration Guidance: CTE with cognitive Q-value supervision outperforms both random trajectory generation and unguided MCTS in safety discovery and preference data coverage.

Data Scaling: Safety metrics continue to improve as the scale of red-teaming data increases, though utility remains robust to this variation, indicating that critical risk states are well-captured for practical data budgets.

Search-Induced Harm: Harmful rate analysis under "no search," "safe search," and "harmful search" scenarios confirms that agentic search exposes the model to a pronounced increase in unsafe outcomes—a risk COMPASS mitigates more consistently than external filtering or conventional RL.

Figure 6: Harmful Rate (HarmR) across retrieval modes on StrongReject, showing COMPASS's robustness to retrieval-induced risks.

Theoretical and Practical Implications

The integration of CTE and ISA represents a paradigm shift toward process-level safety alignment for agentic LLMs. By aligning agents not only on outcomes but also on reasoning trajectories and sub-action decomposition, COMPASS constrains the model’s internal planning and execution logic, directly mitigating logic-based jailbreaks.

Practical Implications:

• Enables safety-compliant LLM agents capable of leveraging search and external tools without incurring linear increases in risk exposure.
• Data- and computation-efficient, thereby scalable to broader classes of LLM architectures and operational settings.
• Compatible with introspective preference optimization protocols, supporting unified utility-safety optimization.

Theoretical Implications and Future Directions:

• Suggests applicability of cognitive MCTS heuristics and step-wise RLHF for broader process-level alignment, including toolformer-style agents and decision-support pipelines.
• Opens investigation into automated risk localization and introspection models in the presence of even more complex action/state spaces or external APIs.
• Future research may extend COMPASS to align agents in open-world and real-time environments, further scaling red-teaming coverage and introspective alignment.

Conclusion

COMPASS introduces cognitive MCTS-guided exploration and introspective step-wise alignment to achieve robust, data-efficient safety specification for LLM-based search agents. Experimental evidence demonstrates substantial reductions in harmful outputs and strong preservation of task utility, with favorable computational cost and cross-architecture generalization. The process-level alignment paradigm instantiated by COMPASS addresses critical deficiencies of prior safety approaches and sets a concrete direction for safe deployment of advanced agentic LLMs in adversarial and open-ended environments.