A Monosemantic Attribution Framework for Stable Interpretability in Clinical Neuroscience Large Language Models

TL;DR

This work tackles the interpretability gap in clinical large language models by addressing polysemantic representations that hinder stable explanations. It introduces a monosemantic attribution framework that fuses attributional methods with mechanistic insight via Sparse Autoencoders (SAEs), embedding activations into a disentangled latent space and applying a Transformer-based explanation optimizer (TEO) with a diffusion-based variant (DEO) and a UMAP linear constraint. Across IID ADNI and OOD BrainLat data, the approach yields a tunable stability–sparseness trade-off, with TEO-SAE and TEO-UMAP providing the most stable and clinically coherent input-level explanations and revealing meaningful biomarkers such as FAQ and AVLT. An external GPT-based evaluation corroborates the clinical relevance of SAE-enabled explanations, supporting safer deployment of LLMs in cognitive health diagnostics and enabling more efficient screening through high-yield biomarkers.

Abstract

Interpretability remains a key challenge for deploying large language models (LLMs) in clinical settings such as Alzheimer's disease progression diagnosis, where early and trustworthy predictions are essential. Existing attribution methods exhibit high inter-method variability and unstable explanations due to the polysemantic nature of LLM representations, while mechanistic interpretability approaches lack direct alignment with model inputs and outputs and do not provide explicit importance scores. We introduce a unified interpretability framework that integrates attributional and mechanistic perspectives through monosemantic feature extraction. By constructing a monosemantic embedding space at the level of an LLM layer and optimizing the framework to explicitly reduce inter-method variability, our approach produces stable input-level importance scores and highlights salient features via a decompressed representation of the layer of interest, advancing the safe and trustworthy application of LLMs in cognitive health and neurodegenerative disease.

Figures (5)

• Figure 1: Proposed interpretability framework for LLM in Alzheimer's diagnosis. The model integrates k-attributional methods with a SAE to generate a monosemantic feature space. An explanation optimizer refines attribution outputs, enhancing clarity and reducing variability. Global explanation quality is visualized and assessed using UMAP and a linear meta-rule, supporting both individual prediction interpretability and cohort-level pattern discovery.
• Figure 2: PCA of token-level attribution representations (LMCI, ADNI test set). The first two principal components are shown, computed from the top eight PCA components of the attribution matrix.
• Figure 3: Left: Stability--sparsity frontier for explanation optimizers on ADNI (Late MCI) in the testing cohort. Scatter points show TEO, TEO--SAE, and TEO--UMAP (SAE) on ADNI (IID) and BrainLat (OOD). Metrics: Sparseness (higher is better) vs. RIS/ROS stability (lower is better). Middle: Token-level heatmap produced by the proposed framework, with feature-attribution scale (green: positive relevance; red: negative relevance; white: neutral). Right:$1^\text{st}$ PCA and 2D UMAP projections of the full testing cohort. Thresholding uses 60% feature attribution over 512 tokens. The generated input text is split into nine subgroups (different colours) based on input modality, as detailed in Supplementary Material 1.1.5, for pattern analysis and biomarker identification.
• Figure 4: GPT-5.1 generated global explanations characterizing attribution-score performance in biomarker identification for TEO, both with and without the SAE layer, and for the TEO-UMAP model. (a) Binary ADNI: Control vs. Alzheimer's disease. (b) Binary BrainLAT: Control vs. Alzheimer's disease.
• Figure 5: (continued) (c) Three-class ADNI: Control, MCI, and LMCI categories.