A Tool Bottleneck Framework for Clinically-Informed and Interpretable Medical Image Understanding

TL;DR

The paper introduces the Tool Bottleneck Framework to advance medical image understanding by grounding predictions in spatially localized, clinically meaningful features. It uses a vision–language model to select a toolbox of domain-specific tools and a learned Tool Bottleneck Model to fuse their outputs, forming a bottleneck that enhances interpretability and data efficiency. Tool knockout augmentation and perturbation improve robustness to arbitrary tool subsets, enabling effective use across histopathology and dermatology datasets with competitive performance and clearer interpretability than end-to-end deep nets or text based tool composition. The approach demonstrates strong results in data-limited settings and provides actionable insights into tool importance, guiding future expansion of clinically informed toolboxes and tool selection strategies.

Abstract

Recent tool-use frameworks powered by vision-language models (VLMs) improve image understanding by grounding model predictions with specialized tools. Broadly, these frameworks leverage VLMs and a pre-specified toolbox to decompose the prediction task into multiple tool calls (often deep learning models) which are composed to make a prediction. The dominant approach to composing tools is using text, via function calls embedded in VLM-generated code or natural language. However, these methods often perform poorly on medical image understanding, where salient information is encoded as spatially-localized features that are difficult to compose or fuse via text alone. To address this, we propose a tool-use framework for medical image understanding called the Tool Bottleneck Framework (TBF), which composes VLM-selected tools using a learned Tool Bottleneck Model (TBM). For a given image and task, TBF leverages an off-the-shelf medical VLM to select tools from a toolbox that each extract clinically-relevant features. Instead of text-based composition, these tools are composed by the TBM, which computes and fuses the tool outputs using a neural network before outputting the final prediction. We propose a simple and effective strategy for TBMs to make predictions with any arbitrary VLM tool selection. Overall, our framework not only improves tool-use in medical imaging contexts, but also yields more interpretable, clinically-grounded predictors. We evaluate TBF on tasks in histopathology and dermatology and find that these advantages enable our framework to perform on par with or better than deep learning-based classifiers, VLMs, and state-of-the-art tool-use frameworks, with particular gains in data-limited regimes. Our code is available at https://github.com/christinaliu2020/tool-bottleneck-framework.

Figures (6)

• Figure 1: Overview of our proposed Tool Bottleneck Framework. A VLM selects tools from a pre-specified toolbox of clinically-relevant tools. These tools are passed to a Tool Bottleneck Model, which composes/fuses the tool outputs to make a prediction.
• Figure 2: Model performance of TBF vs. EfficientNet baseline over varying training set sizes in log scale. Mean $\pm$ 95% CI over seeds. TBF exhibits improved performance across all training set sizes.
• Figure 3: Tool-wise importance (Eq. \ref{['eq:tool_importance']}) and normalized frequency of VLM tool selections for TBM across Camelyon17 (left) and ISIC-BM/-MN (right). In each plot, the left axis shows the relative importance of each tool measured by the change in Accuracy (Camelyon17) and AUC (ISIC) when tools are individually removed during inference. The right axis shows the normalized frequency of tools selected by MedGemma during training.
• Figure 4: (a)-(e): Visualization of the nuclei–dropout intervention on two example Camelyon17 contour maps. For each example patch, we randomly remove individual nuclei by masking them out with probability $p_{\text{mask}}$, which we sweep across $p_{\text{mask}} \in \{0.0, 0.2, 0.4,0.6,0.8\}$ (from left to right) to randomly mask out nuclei in the tool output maps. (f): As $p_{\text{mask}}$ is increased (dropout increased), the fraction of images with a Normal label prediction increases monotonically with dropout.
• Figure 5: Distribution of MedGemma selected tool combinations for Camelyon17