TLDR: A systematic study investigated how Large Language Models (LLMs) represent and process medical knowledge using four interpretability techniques: UMAP projections, gradient-based saliency, layer lesioning, and activation patching. The research created ‘knowledge maps’ for five LLMs, revealing that for Llama3.3-70B, most medical knowledge is processed in the first half of its layers. Key findings include non-linear age representation with a discontinuity at age 18, circular disease progression, and drugs clustering by medical specialty. The study also noted activation collapse in Gemma/MedGemma models at intermediate layers. These results offer guidance for fine-tuning, unlearning biases, and applying causal interventions in medical LLMs.
Large Language Models (LLMs) have demonstrated remarkable capabilities across various tasks, from coding to complex reasoning. However, understanding precisely how these models store and process information, especially in critical domains like medicine, remains a significant challenge. This lack of transparency is particularly concerning for medical applications, where insights into how LLMs represent patient demographics, diseases, and drug treatments are crucial for identifying biases and building safe, trustworthy AI systems.
A recent study delves into this complex area, presenting a systematic investigation into the medical-domain interpretability of LLMs. The research explores how these models both represent and process medical knowledge, aiming to create ‘knowledge maps’ that reveal where specific medical information is stored within the model’s layers. This is vital for guiding future efforts in fine-tuning, unlearning biases, or de-biasing LLMs for medical tasks.
Unveiling LLM Internal Workings
The researchers employed four distinct interpretability techniques to probe the internal mechanisms of five open-source LLMs: Llama3.3-70B, Gemma3-27B, MedGemma-27B, Qwen-32B, and GPT-OSS-120B. These techniques included:
-
UMAP Projections of Intermediate Activations: This method visualizes how the model’s internal representations (activations) cluster together, providing insights into how similar concepts are grouped.
-
Gradient-Based Saliency: By analyzing the gradients with respect to model weights, this technique identifies which parts of the model are most sensitive or important for specific medical concepts.
-
Layer Lesioning/Removal: Similar to neuroscience studies, this involves temporarily disabling specific layers of the LLM to observe the degradation in its medical responses, thereby pinpointing layers crucial for certain knowledge.
-
Activation Patching: This technique involves replacing the activations of a single layer with those from a different prompt to see if a specific piece of information can be ‘patched’ into the model’s processing flow.
By integrating these diverse methods, the study aimed to build confidence in identifying the specific layers where medical knowledge is stored, leveraging the unique strengths of each technique.
Key Discoveries from the Knowledge Maps
The study generated detailed knowledge maps, particularly for Llama3.3-70B, revealing fascinating insights into how medical information is organized within the model:
-
Age Representation: Knowledge about a patient’s age appears to be processed primarily in the initial layers (0-5) of Llama3.3-70B. Interestingly, age is often encoded in a non-linear and sometimes discontinuous manner. A notable discontinuity was observed around age 18, suggesting the model distinguishes between teenagers and adults, which could imply potential biases.
-
Medical Symptoms: Symptoms are processed in layers 0-9 and also 15-40.
-
Diseases: Knowledge related to diseases is found in layers 0-5 or potentially 27-37.
-
Drug Knowledge: Information about drugs is most likely learned in layers 15-45. Furthermore, the model tends to cluster drugs more effectively by their medical specialty (e.g., cardiology, neurology) rather than their mechanism of action (how they work at a molecular level).
-
Drug Dosage: While less conclusive, drug dosage knowledge seems to be processed in the first half of the layers (0-40).
Beyond Llama3.3-70B, the research also uncovered other intriguing phenomena. For instance, Gemma3-27B and MedGemma-27B showed instances where their internal activations ‘collapsed’ at intermediate layers, although they managed to recover by the final layers. This suggests a potential inefficiency or unique processing strategy within these models.
Also Read:
- Unlocking Entity Understanding in Large Language Models
- Unveiling How Language Models Navigate Logic Through Geometric Paths
Implications for Future Medical AI
These findings have significant implications for the development and application of LLMs in medicine. By identifying the specific layers where different types of medical knowledge reside, researchers can more effectively:
-
Fine-tune LLMs: Target specific layers for medical tasks, potentially improving performance and efficiency.
-
Unlearn Biases: Address and mitigate hidden biases related to age, gender, or disease representation by focusing interventions on the relevant layers.
-
Causal Interventions: Apply targeted interventions to modify or enhance medical concepts within the model.
The study acknowledges limitations, such as the absence of ground-truth data for validating internal representations. However, the use of four distinct interpretability methods provides a robust framework, as agreement across these diverse techniques increases confidence in the results. This systematic approach marks a crucial step towards making medical LLMs more transparent, reliable, and ultimately, safer for real-world applications. You can read the full research paper here: Medical Interpretability and Knowledge Maps of Large Language Models.
