Prospector Heads: Generalized Feature Attribution for Large Models & Data

TL;DR

Prospector heads provide encoder-equipped feature attribution without relying on post hoc explanations, enabling scalable, data-efficient localization of class-specific regions across text, images, and graphs. The method builds a two-layer architecture that first quantizes token embeddings into a compact set of concepts and then performs a concept-based convolution to produce prospect maps that highlight informative regions. Across text, pathology images, and protein structures, prospectors outperform attribution baselines by up to 26.3 points in mean localization AUPRC, with robustness to coarse supervision and clear interpretability via learned concepts and kernels. This approach enhances trust in large models for scientific and biomedical tasks and offers a modular, plug-in mechanism compatible with domain-specific foundation models.

Abstract

Feature attribution, the ability to localize regions of the input data that are relevant for classification, is an important capability for ML models in scientific and biomedical domains. Current methods for feature attribution, which rely on "explaining" the predictions of end-to-end classifiers, suffer from imprecise feature localization and are inadequate for use with small sample sizes and high-dimensional datasets due to computational challenges. We introduce prospector heads, an efficient and interpretable alternative to explanation-based attribution methods that can be applied to any encoder and any data modality. Prospector heads generalize across modalities through experiments on sequences (text), images (pathology), and graphs (protein structures), outperforming baseline attribution methods by up to 26.3 points in mean localization AUPRC. We also demonstrate how prospector heads enable improved interpretation and discovery of class-specific patterns in input data. Through their high performance, flexibility, and generalizability, prospectors provide a framework for improving trust and transparency for ML models in complex domains.

Figures (16)

• Figure 1: Explanation-based attribution can be conceptualized as a "wrapper function" for trained classifiers using internals, forward or backward passes, or input perturbations. Prospector heads are instead encoder-equippable like classifier heads and adapt token embeddings with data- and time-efficiency. Flame icon indicates trainable parameters.
• Figure 2: Prospectors are modality-generalizable, amenable to sequences (e.g., text), images (e.g., pathology), and graphs (e.g., protein structures). They can also operate on embeddings from either partial- or full-context encoders.
• Figure 3: Prospector-equipped encoders produce attribution maps (called "prospect maps") over two layers. Details for fitting and inference are in Sections \ref{['sec:inference']}, \ref{['sec:fit']}, and \ref{['supp:internals']}.
• Figure 4: Layer (I) fitting and inference ($K=5$). Quantized token embeddings define spatially resolved concepts, which together form data sprites.
• Figure 5: Layer (II) fitting and inference ($K=5$). Concept frequencies are used to build sprite embeddings, which are used to fit a K2conv kernel. Flame icon indicates trainable parameters.

Theorems & Definitions (6)

• Definition 3.1: Map Graph
• Definition 3.2: Connectivity
• Definition 3.3: Partial-context Encoder
• Definition 3.4: Full-context Encoder
• Definition 3.5: Self-complete graph
• Theorem 3.6: Range of implicit $n$-grams