PASS: Probabilistic Agentic Supernet Sampling for Interpretable and Adaptive Chest X-Ray Reasoning

TL;DR

PASS tackles the critical need for interpretable, adaptive, and safe multimodal chest X-ray reasoning by introducing a probabilistic controller over a multimodal agentic supernet. It learns a task-conditioned distribution to sample layer-wise tool sequences, producing probability-annotated decision paths and an evolving memory for in-context reasoning. A principled three-stage training regime—expert warm-up, contrastive path ranking, and cost-aware reinforcement learning—aligns accuracy with computational efficiency, enabling early exits and a Pareto-optimal cost-accuracy frontier. Evaluations on CAB-E and public radiology benchmarks demonstrate superior diagnostic accuracy, stronger language fidelity, and improved safety metrics, validating the feasibility of auditable, tool-driven medical AI. PASS thus represents a paradigm shift toward interpretable, adaptive, and resource-aware agentic systems for high-stakes multimodal medical reasoning.

Abstract

Existing tool-augmented agentic systems are limited in the real world by (i) black-box reasoning steps that undermine trust of decision-making and pose safety risks, (ii) poor multimodal integration, which is inherently critical for healthcare tasks, and (iii) rigid and computationally inefficient agentic pipelines. We introduce PASS (Probabilistic Agentic Supernet Sampling), the first multimodal framework to address these challenges in the context of Chest X-Ray (CXR) reasoning. PASS adaptively samples agentic workflows over a multi-tool graph, yielding decision paths annotated with interpretable probabilities. Given the complex CXR reasoning task with multimodal medical data, PASS leverages its learned task-conditioned distribution over the agentic supernet. Thus, it adaptively selects the most suitable tool at each supernet layer, offering probability-annotated trajectories for post-hoc audits and directly enhancing medical AI safety. PASS also continuously compresses salient findings into an evolving personalized memory, while dynamically deciding whether to deepen its reasoning path or invoke an early exit for efficiency. To optimize a Pareto frontier balancing performance and cost, we design a novel three-stage training procedure, including expert knowledge warm-up, contrastive path-ranking, and cost-aware reinforcement learning. To facilitate rigorous evaluation, we introduce CAB-E, a comprehensive benchmark for multi-step, safety-critical, free-form CXR reasoning. Experiments across various benchmarks validate that PASS significantly outperforms strong baselines in multiple metrics (e.g., accuracy, LLM-Judge, semantic similarity, etc.) while balancing computational costs, pushing a new paradigm shift towards interpretable, adaptive, and multimodal medical agentic systems.

Figures (2)

• Figure 1: An overview of PASS. Given a multimodal complex reasoning task (CXR image, textual comprehensive query, multimodal personalized context), our probabilistic controller learns a continuous task-conditioned distribution over the agentic supernet (i.e. a directed acyclic graph of medical agent containers). At each step, it samples an action, yielding a workflow annotated with interpretable probabilities for post-audits and directly enhances clinical AI safety. Tool outputs, which can be both text and images, are summarized and fed into an evolving personalized memory and shared in-context to inform subsequent steps. The controller is trained via a principled three-stage strategy (expert knowledge warm-up, contrastive path ranking, cost-aware reinforcement learning) to optimize the accuracy-cost trade-off. Eventually, PASS is enabled to answer multimodal medical questions in free-form text via an interpretable, adaptive, and efficient agentic reasoning process.
• Figure 2: Cost-Accuracy Pareto Frontier analysis. Each orange point on the dashed frontier corresponds to a specific penalty weight ($\lambda$) configuration of PASS, enabling flexible cost–accuracy trade-offs at deployment. MedRAX and LLaVA-Med are plotted as additional points for comparison. Lower normalized inference cost and higher accuracy are preferred; the arrow indicates the desired direction toward the top-left preferred region.