Tracing Pharmacological Knowledge In Large Language Models

This study employs causal intervention and linear probing techniques to demonstrate that pharmacological knowledge in Llama-based models is encoded through distributed representations across early layers and multiple tokens, rather than being localized to specific final tokens.

The Gist

Imagine you have a super-smart robot librarian named LLM (Large Language Model). This robot has read almost every medical textbook, research paper, and drug manual in existence. It's incredibly good at answering questions like, "Which drugs help lower blood pressure?" or "What class does this medication belong to?"

But here's the mystery: How does the robot actually "know" this?

Does it have a tiny filing cabinet inside its brain where it keeps a list of "Blood Pressure Drugs"? Or is the knowledge scattered everywhere, like a puzzle where the pieces are mixed up?

This paper is like a team of detectives (the researchers) who decided to open up the robot's brain to see how it stores this medical knowledge. They used two main tools to solve the mystery: Activation Patching and Linear Probing.

Here is the story of what they found, explained simply:

1. The Detective Tool: "Activation Patching" (The Swap Test)

Imagine you are watching a play. You want to know which actor is responsible for the main character's big speech.

• The Test: The researchers watched the robot answer a question about a drug. Then, they "swapped" the robot's brain activity at specific moments with the brain activity from a different question (where the answer was wrong).
• The Result: If the robot suddenly gave the wrong answer after the swap, that specific moment in the brain was crucial for knowing the answer.

What they found:

• It happens early: The most important "thinking" happens in the first few layers of the robot's brain (the early stages of processing), not at the very end.
• It's not the last word: You might think the robot figures out the answer when it reads the last word of the drug's name. But the researchers found that the middle words of the drug name were actually doing the heavy lifting! It's like the robot understands the whole concept of "Aspirin" while it's reading "Asp..." and "...rin," not just when it finishes the word.

2. The Detective Tool: "Linear Probing" (The Snapshot Test)

Imagine taking a photo of the robot's brain at every single step as it reads a sentence. You want to see if the "Drug Class" information is visible in just one photo, or if you need to combine many photos to see the picture.

• The Test: They tried to guess the drug class by looking at the brain activity of just one word (token) at a time. Then, they tried looking at the average of all the words in the drug name combined.
• The Result:
  • One word? The robot's brain looked like static noise. You couldn't tell the drug class from a single word.
  • All words combined? Suddenly, the pattern became crystal clear! The information was distributed. It wasn't stored in one specific "filing cabinet"; it was spread out across the whole group of words, like a scent that fills an entire room rather than sitting in one corner.

3. The Big Surprise: The Knowledge Was There From the Start

The researchers checked the robot's brain before it even started processing the sentence (in the initial "embedding" stage).

• The Finding: The information about drug groups was already there before the robot even started reading! It's as if the robot's raw materials were already pre-sorted into "Blood Pressure" and "Pain Relief" buckets before the actual thinking began.

The Main Takeaway (The "Aha!" Moment)

Before this study, people thought the robot might store medical facts like a human stores a phone number: in one specific spot, waiting to be recalled.

This paper proves that's wrong.

Instead, the robot stores pharmacological knowledge like a symphony:

• No single instrument (word) holds the whole song.
• The music (the meaning) emerges only when all the instruments play together.
• The conductor (the early layers of the brain) sets the tone immediately, and the melody is built up through the middle of the song, not just at the finale.

Why Does This Matter?

If we want to trust robots to help doctors prescribe medicine, we need to know how they think.

• If the knowledge is scattered and built from the ground up, we can't just "delete" a bad fact by erasing one line of code.
• But now that we know where (early layers) and how (distributed across words) this knowledge lives, scientists can build better, safer, and more transparent medical AI. We can fix the "brain" without breaking the whole system.

In short: The robot doesn't have a cheat sheet. It has a complex, distributed understanding of medicine that starts the moment it sees the words, built by the middle of the sentence, and spread out across the whole group of words.

Technical Summary

Here is a detailed technical summary of the paper "Tracing Pharmacological Knowledge in Large Language Models."

1. Problem Statement

Large Language Models (LLMs) have demonstrated strong empirical performance in pharmacology and drug discovery tasks (e.g., target identification, drug-drug interaction prediction). However, the internal mechanisms by which these models encode, store, and retrieve pharmacological knowledge remain poorly understood.

• The Gap: While general factual knowledge in LLMs has been studied (e.g., where facts are stored in the residual stream), it is unclear if biomedical concepts like drug groups (e.g., "adrenergic alpha-agonists") follow similar patterns.
• The Challenge: Existing benchmarks do not explicitly target drug-class relationships, and tokenization issues make it difficult to evaluate knowledge at the individual token level.
• Goal: To provide a systematic, mechanistic analysis of how drug-group semantics are represented in Llama-based biomedical LLMs using causal and probing-based interpretability methods.

2. Methodology

The authors employed a two-pronged approach: Activation Patching (causal intervention) and Linear Probing (correlational analysis).

A. Dataset Construction

• Source: Drug names and categories were parsed from the U.S. National Library of Medicine (NLM) and National Center for Biotechnology Information (NCBI).
• Format: A two-choice question-answering dataset was created.
  • Rationale: Drug names are tokenized variably, making single-token evaluation unreliable. Additionally, multiple drugs can belong to the same class, so a unique "correct" token doesn't exist.
  • Structure: Prompts ask which compound belongs to a specific class (e.g., "Which compound is categorized as vasoconstrictor agents?"). Distractors were randomly sampled, and answer positions were varied.

B. Activation Patching (Causal Analysis)

• Objective: To identify where and when drug-group knowledge is stored by intervening on model activations.
• Procedure:
  1. Clean Run: A prompt with a known correct answer (e.g., Drug A is a vasoconstrictor).
  2. Counterfactual Run: A prompt where the drug group is swapped (e.g., Drug A is now a bronchoconstrictor), changing the correct answer.
  3. Patching: Activations from the clean run are injected into the counterfactual run at specific layers and token positions.
• Metric: Normalized Logit Difference (LD) to measure the causal impact of the intervention on the model's output preference.
• Targets:
  • Residual Stream: Patching at the start of transformer blocks across all layers.
  • MLP Layers: Patching outputs of Multi-Layer Perceptrons (MLPs) in contiguous 10-layer windows.

C. Linear Probing (Representation Analysis)

• Objective: To determine if pharmacological semantics are linearly decodable from internal representations and if they are localized to specific tokens or distributed.
• Setup: Paired drug groups with opposing semantics (e.g., agonists vs. antagonists; stimulants vs. depressants).
• Training: Logistic regression classifiers were trained on:
  1. Individual Token Activations: Activations of specific tokens within the drug-group span.
  2. Sum-Pooled Activations: Aggregated activations across the entire drug-group span.
• Scope: Probing was performed across all layers, including the embedding space (pre-layer 0).

3. Key Results

A. Model Performance

• Llama-based models (e.g., Llama-3.1-8B-Instruct, OpenBioLLM-8B) achieved high accuracy (>90%) on the drug class–name relationship task, outperforming general-purpose models and specialized models like BioGPT in some cases. This confirms the models possess substantial pharmacological knowledge.

B. Causal Mechanisms (Activation Patching)

• Early Layer Dominance: The strongest causal effects were observed in the early layers (specifically the first 10 layers) of the model. Information injected here propagates through the network to support task performance.
• Intermediate Token Importance: Contrary to findings on general factual knowledge (where the final subject token is most critical), the strongest causal effects for drug groups occurred at intermediate tokens within the drug-group span, not the final token.
• MLP vs. Residual Stream: Patching early MLP layers (0–10) consistently produced positive causal effects, confirming that semantic refinement happens early in the processing pipeline.

C. Distributed Representations (Linear Probing)

• Token-Level vs. Pooled: Linear probes trained on individual tokens performed near chance levels. In contrast, probes trained on sum-pooled representations (aggregating the whole span) achieved near-perfect accuracy.
• Embedding Space Presence: Sum-pooled probes achieved maximal accuracy even on activations from the embedding space (pre-layer 0).
• Conclusion: Pharmacological semantics are not localized to a single token but are distributed across the token span and are already present in the initial embedding space.

4. Key Contributions

1. First Systematic Mechanistic Analysis: This is the first study to apply causal activation patching and linear probing specifically to pharmacological knowledge in LLMs.
2. Discovery of Distributed Semantics: The paper challenges the "single-token" hypothesis for biomedical concepts, demonstrating that drug-group knowledge is a distributed representation emerging from the aggregation of multiple tokens.
3. Identification of Early-Layer Circuits: It identifies that early layers (0–10) and intermediate tokens are the critical loci for encoding drug-group semantics, differing from the mechanisms of general factual knowledge.
4. Methodological Framework: Establishes a robust pipeline for benchmarking and interpreting biomedical LLMs, addressing tokenization challenges through two-choice QA and sum-pooled probing.

5. Significance and Implications

• Scientific Trustworthiness: Understanding how models encode medical knowledge is crucial for deploying LLMs in high-stakes biomedical domains. Knowing that knowledge is distributed and early-layer dependent helps in debugging and improving model reliability.
• Model Architecture Insights: The findings suggest that biomedical LLMs may rely on different circuit-level mechanisms for domain-specific concepts compared to general facts.
• Future Directions: The work provides a foundation for "mechanistic editing" of biomedical models (e.g., correcting hallucinations by intervening in early layers) and highlights the need for future research to analyze individual drug representations and specific attention circuits.

6. Limitations

• The study is currently limited to drug groups and has not yet been extended to individual drug compounds or other biomedical categories.
• The specific attention heads or circuits responsible for forming these representations were not explicitly identified.