💬 NLP

A Monosemantic Attribution Framework for Stable Interpretability in Clinical Neuroscience Transformer-Based Language Models

This paper introduces a unified interpretability framework that leverages monosemantic feature extraction to generate stable, input-level importance scores for Transformer-based language models, thereby addressing the challenges of polysemanticity and inter-method variability to enable trustworthy clinical applications in Alzheimer's disease diagnosis.

Original authors: Michail Mamalakis, Tiago Azevedo, Cristian Cosentino, Chiara D'Ercoli, Subati Abulikemu, Zhongtian Sun, Richard Bethlehem, Pietro Lio

Published 2026-06-02

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Michail Mamalakis, Tiago Azevedo, Cristian Cosentino, Chiara D'Ercoli, Subati Abulikemu, Zhongtian Sun, Richard Bethlehem, Pietro Lio

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Problem: The "Noisy Radio"

Imagine you have a very smart radio (a Large Language Model) that can listen to a patient's medical history and tell you if they might have Alzheimer's disease. The radio works great at making the prediction, but it's terrible at explaining why.

When you ask the radio, "Why did you say this patient has Alzheimer's?" it gives you a confusing answer. It's like the radio is tuned to a station where all the voices are talking over each other at once. In the paper's language, the radio's internal parts (neurons) are polysemantic. This means one single "voice" in the radio is trying to talk about "memory loss," "age," and "blood pressure" all at the same time. Because everything is mixed up, the explanations the radio gives are unstable, confusing, and sometimes wrong. Doctors can't trust a radio that changes its story every time you ask it a question.

The Solution: The "Translator" (SAE)

The authors built a special translator called a Sparse Autoencoder (SAE). Think of this translator as a "de-mixer" for the radio.

Before the Translator: The radio's internal voices are a chaotic jam session.
After the Translator: The translator takes that chaotic noise and separates it into distinct, single-note instruments. Now, one voice only talks about "memory," another only talks about "walking speed," and another only talks about "family history."

The paper calls this monosemanticity. By forcing the radio to speak in these clear, single-concept notes, the explanations become much more stable and easier to understand.

The "Group Chat" Problem

Even with the translator, the authors noticed a new problem. They tried six different ways to ask the radio for an explanation (like asking six different people to summarize a movie). Sometimes Person A said the movie was about love, and Person B said it was about war. They didn't agree. In the paper, this is called inter-method variability. If the explanation changes depending on which "person" you ask, doctors can't trust it.

The "Editor" (TEO)

To fix the disagreement, the authors created an Explanation Optimizer (TEO). Think of this as a smart editor who sits in the middle of the group chat.

The editor listens to all six different explanations.
It figures out which parts of the story are consistent and which parts are just noise.
It writes a single, perfect summary that combines the best parts of all six, ensuring the story is stable and makes sense.

The "Team Photo" (UMAP Constraint)

Finally, the authors wanted to make sure that when they looked at the explanations for many patients at once, the patterns looked organized.

Imagine taking a photo of a crowd. Without guidance, everyone might be standing in a messy, random pile. The authors added a linear constraint (a rule) that acts like a photographer shouting, "Everyone, line up in a straight row!"

This ensures that when they look at the data for a whole group of patients, the patterns line up neatly. This helps them spot the "big picture" trends that apply to most people, rather than getting lost in individual quirks.

What They Found

The paper tested this system on real medical data from two different groups of patients (one from the US and one from Latin America).

The Result: When they used the "Translator" (SAE) and the "Editor" (TEO), the explanations became stable. The radio stopped changing its story.
The Trade-off: Sometimes, making the explanation super stable made it a little less "sparse" (meaning it highlighted a few more words than strictly necessary), but the authors found a way to balance this so the explanations were both stable and focused.
The Clinical Win: The system successfully identified specific medical tests (like memory tests or questionnaires about daily activities) that were most important for diagnosing Alzheimer's. It did this consistently, even when looking at different groups of patients.

The Bottom Line

The paper doesn't claim to cure Alzheimer's. Instead, it claims to fix the flashlight doctors use to see how AI makes decisions. By untangling the messy internal signals of the AI and organizing the explanations, they created a tool that gives doctors a clear, stable, and trustworthy view of why an AI thinks a patient might have the disease. This makes it safer to use these powerful AI tools in real hospitals.

Technical Summary: A Monosemantic Attribution Framework for Stable Interpretability in Clinical Neuroscience Transformer-Based Language Models

Problem Statement

The deployment of Large Language Models (LLMs) and Transformer-based Language Models (LMs) in high-stakes clinical settings, such as the diagnosis and progression monitoring of Alzheimer's Disease (AD), is hindered by a lack of reliable interpretability. Existing attribution methods (e.g., gradients, perturbations, integrated paths) suffer from two primary limitations:

Polysemanticity: Internal neural representations in Transformers are often polysemantic, meaning individual neurons or features encode multiple, semantically unrelated concepts. This entanglement undermines standard attribution techniques that assume a one-to-one correspondence between features and meanings, leading to ambiguous or misleading explanations.
Instability and Variability: Attribution methods exhibit high inter-method variability and unstable explanations, particularly when facing distribution shifts (e.g., moving from one clinical cohort to another). This instability reduces trust in AI-driven clinical decision-making.

Conversely, mechanistic interpretability approaches, such as Sparse Autoencoders (SAEs), successfully identify semantically coherent, monosemantic components but typically lack the attributional resolution required to explain how specific inputs contribute to model predictions in real-world scenarios.

Methodology

The authors propose a Unified Monosemantic Attribution Framework that integrates attributional and mechanistic interpretability perspectives. The framework operates through the following pipeline:

Monosemantic Feature Extraction (SAE):
- The framework employs Sparse Autoencoders (SAEs) to transform the polysemantic activations of a Transformer-based LM (specifically the 22nd layer of ModernBERT) into a monosemantic latent space.
- By mapping activations to a higher-dimensional, sparse feature space, SAEs disentangle overlapping concepts, ensuring individual features correspond to specific semantic factors.
- Variants tested include TopK, JumpReLU, and Gated-SAEs, with TopK showing superior performance.
Dual-Level Attribution:
- Six established attribution methods (Feature Ablation, Layer Activations, Layer Conductance, Layer Gradient-SHAP, Layer Integrated Gradients, and Layer Gradient×Activation) are applied at two levels:
  - From the SAE feature space to the encoder layer.
  - From the encoder layer to the input token embeddings.
- This allows the attribution of importance scores to input tokens based on the behavior of the disentangled, monosemantic internal mechanisms.
Explanation Optimization (TEO/DEO):
- To address inter-method variability, the framework introduces an Explanation Optimizer. Two architectures are explored: a Diffusion UNet1D (DEO) and a Transformer-based autoencoder (TEO).
- The optimizer is trained to reconstruct a unified attribution vector that captures the consensus across the six methods.
- The loss function combines four components: Sparseness (Gini Index), Relative Input Stability (RIS), Relative Output Stability (ROS), and Similarity to the weighted average of the base methods.
Geometric Constraints (UMAP Linear Constraint):
- To ensure global coherence and robustness, the framework applies a UMAP projection to the attribution scores and imposes a linear constraint ( $u_{i1} = u_{i2}$ ) on the embedding components.
- This constraint enforces a linear structure along the dominant statistical directions of the explanation space, acting as a regularizer to produce population-level patterns that are consistent across the cohort.

Key Contributions

Unification of Paradigms: The paper presents the first framework to successfully integrate attributional techniques (surface-level explanations) with mechanistic interpretability (structural insights via SAEs) specifically for LLMs in clinical inference.
Stability via Monosemanticity: The authors demonstrate that transforming polysemantic activations into a monosemantic space significantly reduces inter-method variability and improves the stability of attribution scores.
Optimization Mechanism: The introduction of the Transformer Explanation Optimizer (TEO) with a linear UMAP constraint provides a tunable trade-off between explanation sparsity and stability, outperforming simple weighted averages of attribution methods.
Clinical Validation: The framework is rigorously tested on Alzheimer's Disease datasets (ADNI and BrainLat), demonstrating robustness under both In-Distribution (IID) and Out-of-Distribution (OOD) conditions.

Experimental Results

The framework was evaluated on binary (Control vs. AD) and three-class (Control, EMCI, LMCI) classification tasks using ADNI (IID) and BrainLat (OOD) datasets.

Stability Metrics: The proposed TEO-SAE variant consistently achieved the lowest Relative Input Stability (RIS) and Relative Output Stability (ROS) scores across all tasks and datasets, indicating superior robustness compared to classical attribution methods and SAE-less optimizers.
Sparsity-Stability Trade-off: While SAEs generally improved stability, they sometimes reduced sparsity. The TEO-UMAP (SAE) variant successfully recovered higher sparsity (>0.40) while maintaining low RIS/ROS, offering a tunable balance.
Statistical Significance: Paired t-tests and Wilcoxon signed-rank tests confirmed that adding an SAE bottleneck significantly reduced complexity and RIS (p < 10⁻⁴) in binary tasks, with the most consistent improvements in ROS for the Mild Cognitive Impairment (MCI) group in three-class tasks.
External Semantic Validation: An auxiliary evaluation using ChatGPT-5.1 to interpret the highlighted tokens revealed that TEO-SAE and TEO-UMAP produced explanations containing coherent neurocognitive indicators and demographic risk factors. In contrast, TEO without SAE often fixated on task artifacts (e.g., instruction counts) rather than clinically meaningful biomarkers.
Biomarker Identification: The framework successfully identified class-specific biomarkers (e.g., Functional Activities Questionnaire for Controls, Auditory Verbal Learning Test for AD) that align with established clinical literature, supporting the extraction of minimal, high-yield assessment sets.

Significance and Claims

The paper claims that this framework advances the safe and trustworthy application of LMs in cognitive health and neurodegenerative disease by:

Enhancing Trust: Providing stable, consistent, and clinically meaningful explanations that are less susceptible to the polysemantic entanglement inherent in standard Transformer models.
Improving Robustness: Defining a tunable sparsity-stability frontier that generalizes effectively across demographic and protocol shifts (OOD settings), a critical requirement for clinical deployment.
Facilitating Clinical Integration: Enabling the identification of specific, interpretable cognitive biomarkers (e.g., AVLT, FAQ, CDT) that can streamline screening pipelines and reduce clinical burden.

The authors maintain a modest stance, acknowledging that the study is currently restricted to the neurodegenerative domain and that future work is required to incorporate clinician-in-the-loop evaluations, uncertainty quantification, and fairness auditing to fully realize the framework's potential in broader medical contexts.