Technical Summary: Med-SegLens

Problem Statement

Modern medical image segmentation models, particularly those based on convolutional and transformer architectures, achieve strong predictive performance on benchmarks but remain largely opaque. This opacity limits the ability to diagnose systematic failures, understand dataset shifts (e.g., when models are deployed across heterogeneous populations), or intervene in a principled manner. Current mitigation strategies often rely on costly retraining or post-hoc interpretability methods (e.g., saliency maps, attention mechanisms) that provide limited causal insight into how internal representations encode population-specific priors or drive errors under distribution shift. Consequently, failures are difficult to attribute, compare across datasets, or correct without retraining.

Methodology

The authors introduce Med-SegLens, a mechanistic model-diffing framework designed to decompose segmentation model activations into interpretable latent features. The framework operates through the following pipeline:

1. Latent Extraction via Sparse Autoencoders (SAEs):
   Instead of analyzing raw activations, the framework extracts intermediate activations from trained segmentation models (specifically SegFormer and U-Net) and trains Sparse Autoencoders (using a BatchTopK operator) to decompose these activations into a sparse, disentangled latent space. This enforces an average sparsity (e.g., $k=32$ active latents per sample) to encourage monosemantic features.

2. Automated Latent Interpretation:
   To ground these latents in the imaging domain, the authors propose a geometry- and spatially grounded automated interpretation pipeline. This assigns anatomical and pathological semantics to latents by analyzing their Top-K highest-activating samples. Metrics include:

   • Spatial Concentration: Spatial entropy of activation maps.
   • Geometric Localization: Brain/organ edge preference, depth within tissue, and centroid localization.
   • Class Association: Fraction of activation mass assigned to specific segmentation classes (e.g., edema, necrotic core, white matter).

3. Cross-Dataset Model Diffing:
   The core contribution is a latent alignment procedure that explicitly aligns SAE features across models trained on different datasets (e.g., Healthy, Adult Glioma, Pediatric Glioma, Sub-Saharan African Glioma).

   • Alignment: Cosine similarity matrices are computed between encoder and decoder weights of SAEs trained on different datasets.
   • Matching: The Hungarian algorithm is applied to find strict one-to-one correspondences between latent vectors.
   • Classification: Latents satisfying high similarity thresholds in both encoder and decoder spaces are classified as shared (population-invariant). Those failing this criterion are classified as dataset-specific.

4. Latent-Level Intervention:
   The framework enables targeted interventions at inference time without retraining. By identifying "causal bottlenecks" (specific dataset-specific latents), the authors apply Latent Steering (additive or multiplicative scaling) to correct errors. This involves amplifying or suppressing specific latent vectors to recover missing structures or suppress spurious predictions.

Key Contributions

• Med-SegLens Framework: The first mechanistic model-diffing framework for medical image segmentation, enabling principled comparison of internal representations across datasets via interpretable latent features.
• Automated Interpretation: A geometry- and spatially grounded method for assigning semantic meaning to SAE latents without manual labeling, applicable across different organs and modalities (demonstrated on Brain MRI, CT, and non-medical face parsing).
• Causal Discovery: Demonstration that dataset shift is driven by differential reliance on population-specific latents rather than a lack of shared representations. Shared latents act as a stable backbone, while dataset-specific latents act as causal bottlenecks for performance.
• Retraining-Free Adaptation: A method to diagnose and partially mitigate segmentation failures through targeted latent-level interventions, recovering performance in a significant portion of failure cases without accessing target-domain data or retraining.

Results

The framework was evaluated across four cohorts (Healthy, Adult Glioma, Pediatric Glioma, SSA Glioma) using SegFormer and U-Net architectures.

• Representation Analysis:

  • Models trained on different populations learn both shared and dataset-specific representations. Adult and Pediatric models shared the most latents (58.9%), while overlap with SSA cohorts was lower (43.2%), likely due to scanner and acquisition differences.
  • Feature swapping experiments confirmed that shared latents preserve spatial localization and semantic structure, whereas non-shared latents produce disrupted responses.

• Failure Diagnosis and Recovery:

  • Instance-Level: Targeted steering of the most active SAE latents recovered performance in 70.6% of Adult, 60.0% of Pediatric, and 63.3% of SSA failure cases.
  • Class-Level: For the Adult model, edema segmentation (a known weakness) improved from a Dice score of 65.8% to 69.1% on failure cases, with overall mean Dice increasing from 81.8% to 82.6%.
  • Cross-Dataset Adaptation: Without retraining, the framework improved cross-dataset adaptation. Specifically, for the Adult $\to$ Pediatric transfer, suppressing edema-related latents improved the edema Dice score from 39.4% to 74.2%. Conversely, amplifying these latents improved Pediatric $\to$ Adult performance.

• Baseline Comparison:

  • Med-SegLens outperformed standard inference-level refinements like Connected-Component Filtering (CCF) and Uncertainty Filtering (UF), which often failed to recover missing structures (false negatives) or provided only marginal gains.
  • The proposed latent alignment method (using encoder-decoder consistency) achieved perfect agreement and stability compared to Greedy, Sinkhorn, or encoder-only matching baselines.

Significance and Claims

The paper claims that Med-SegLens provides a practical and mechanistic tool for diagnosing failures and mitigating dataset shift in medical segmentation models.

• Mechanistic Insight: The work suggests that population-specific priors are encoded in identifiable and manipulable latent features. This moves beyond "black box" performance metrics to a causal understanding of why models fail under distribution shift.
• Intervention without Retraining: A primary impact is the ability to correct errors and adapt models to out-of-distribution data at inference time. This is particularly significant in high-stakes clinical settings where collecting large, representative datasets or retraining models is infeasible.
• Equity and Reliability: By identifying and correcting population-specific failure modes, the framework has the potential to reduce disparities in model performance across diverse populations (e.g., pediatric vs. adult, or different geographic cohorts).
• Generalizability: While focused on brain MRI, the authors demonstrate that the pipeline generalizes to multi-organ CT and non-medical semantic segmentation, suggesting the methodology is applicable to other high-dimensional perception models.

The authors emphasize that these techniques are intended to support model auditing, robustness analysis, and scientific understanding, rather than to replace clinical judgment. Deployment in clinical environments would still require rigorous validation and expert oversight.