💬 NLP

LMs as Task-Specific Knowledge Bases: An Interpretability Analysis

This paper challenges the view of language models as unified knowledge bases by demonstrating through behavioral and mechanistic analyses that factual knowledge is encoded in a task-specific manner, where distinct parameter subsets are activated depending on the query context, thereby undermining the consistency and reliability of factual retrieval.

Original authors: Amit Elhelo, Amir Globerson, Mor Geva

Published 2026-06-26

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Amit Elhelo, Amir Globerson, Mor Geva

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

Imagine you have a giant, super-smart encyclopedia that has read almost everything on the internet. You might think of this encyclopedia as a single, perfect library where every fact is stored in one specific spot. If you ask, "What is the capital of France?" or "Is Paris the capital of France?", you expect the library to pull the exact same file from the same shelf and give you the same answer every time. This is how traditional databases work: one truth, one source.

This paper investigates whether Large Language Models (LLMs)—the AI brains behind tools like chatbots—work like that perfect library. The researchers, Amit Elhelo, Amir Globerson, and Mor Geva, found that they do not.

Instead of a single, unified library, the AI's "knowledge" is more like a collection of specialized, task-specific toolkits. Here is a breakdown of their findings using simple analogies:

1. The "Different Toolkits" Analogy

Imagine you have a Swiss Army knife. It has a screwdriver, a blade, and a corkscrew.

Traditional View: You might think the "knowledge" of how to open a bottle is stored in one central part of the knife.
Paper's Finding: The AI is actually like a toolbox where the "knowledge" of how to open a bottle is stored in the corkscrew only when you are asked to open a bottle. But if you ask, "Is this a corkscrew?" (a different type of question), the AI uses a completely different set of tools to answer, even though the answer is about the same object.

The researchers tested this by asking the AI the same fact (e.g., "Paris is the capital of France") in different ways:

Generation: "What is the capital of France?" (The AI has to write the answer).
Discrimination: "Is Paris the capital of France? Yes or No." (The AI has to pick a choice).

They found that the AI often learns the fact for the "Write it" task but fails to "know" it for the "Pick it" task, and vice versa. It's as if the AI has a "Writing Brain" and a "Picking Brain" that don't always talk to each other.

2. The "Training Camp" Observation

The researchers watched the AI learn over time, like a coach watching a student study for different exams.

The Expectation: If the student learns that "Paris is the capital" for a writing test, they should immediately know it for a multiple-choice test too, because it's the same fact.
The Reality: The student often passed the writing test but still failed the multiple-choice test days later. The knowledge didn't "co-emerge" (appear at the same time) across different formats. This suggests the AI isn't storing the fact in one central place; it's learning a specific way to handle that fact for that specific type of question.

3. The "Surgery" Experiment (Mechanistic Analysis)

To prove this, the researchers performed "surgery" on the AI's brain (its internal parameters). They tried to find the tiny, specific group of neurons responsible for knowing "Paris is the capital" when asked in a specific way.

The Result: They found distinct, separate groups of neurons for each task.
- If they "turned off" the neurons used for the "Write it" task, the AI forgot the answer to the writing question but could still answer the "Yes/No" question perfectly.
- If they "turned off" the "Yes/No" neurons, the AI failed that test but could still write the answer.
The Metaphor: It's like having two different keys for the same door. One key opens the door when you push it from the left; the other opens it when you pull it from the right. If you lose the "push" key, the door is locked for that action, even though the "pull" key still works.

4. The "Chain of Thought" Surprise

The paper also looked at "Chain of Thought" (CoT) reasoning, where the AI is asked to "think step-by-step" before answering.

The Finding: When the AI is asked to think out loud, it seems to use a mix of these different toolkits. If you remove the specific neurons for the final answer, the AI can still often get the right answer by "thinking" through the problem, because it's borrowing knowledge from other task-specific toolkits that it wouldn't normally use for a direct answer.
The Analogy: It's like a detective who usually solves a case by looking at a single clue (direct answer). But if you force them to write a full report (Chain of Thought), they start checking other files and cross-referencing different toolkits, which helps them solve the case even if the main clue is missing.

The Big Takeaway

The paper concludes that the idea of an AI as a "Knowledge Base" (a single source of truth) is a myth. Instead, what the AI knows and how you ask it are tangled together.

Reliability Risk: If you edit the AI to "forget" a fact for one type of question (like a multiple-choice quiz), it might still remember that fact for a writing prompt.
Evaluation Risk: If you only test an AI on one type of question, you might think it knows a fact, but it might be completely clueless when asked in a different format.

In short, the AI doesn't have a single, unified memory. It has a patchwork of specialized memories, each tuned to a specific way of asking questions.

Technical Summary: LMs as Task-Specific Knowledge Bases: An Interpretability Analysis

Problem Statement

Language models (LMs) are frequently conceptualized as knowledge bases, where parameters store factual information accessible via various queries. A fundamental property of well-designed knowledge bases is task-invariance: different queries for the same fact should retrieve consistent results from a single source of truth. For instance, a traditional database should return "Paris" for both "What is the capital of France?" and "The capital of France is ___."

This paper investigates whether LMs satisfy this task-invariance property. Specifically, it asks if factual knowledge is stored in a task-independent manner or if the model's internal representation of a fact is intertwined with the specific task format (e.g., multiple-choice vs. open-ended generation) used to query it. If LMs lack task-invariance, the "knowledge base" analogy is undermined, with implications for the reliability, consistency, and controllability of factual knowledge in these systems.

Methodology

The authors employ a dual approach combining behavioral analysis (tracking knowledge emergence during training) and mechanistic analysis (localizing specific parameter subsets).

1. Behavioral Analysis: Co-emergence Hypothesis

The authors track the acquisition of factual knowledge across training checkpoints of the OLMo-3-7B IT model.

Data: Five relational datasets (e.g., country-capital, landmark-country) formulated as (subject, relation, object) triplets.
Tasks: Six distinct task formats: Next-token completion (COMPLETION), Fill-in-the-blank (FITB), Open-ended QA (OPENQA), Four-way Multiple Choice QA (MCQA), Negative MCQA, and Binary Verification.
Hypothesis: If knowledge is task-invariant, a fact acquired for one task should co-emerge (appear simultaneously) on other tasks the model is already competent at. Formally, if a fact $f$ emerges on task $t$ at checkpoint $E(f, t)$ , and the model is competent on task $t'$ at $E(\cdot, t')$ , then $E(f, t')$ should be $\le \max(E(f, t), E(\cdot, t'))$ .
Statistical Test: The authors formalize task-invariance as conditional independence between facts and tasks ( $P(\text{correct} | f, t) = P(\text{correct} | f) \cdot P(\text{correct} | t)$ ) and test for interaction terms using a two-way ANOVA on log-probabilities.

2. Mechanistic Analysis: Parameter Localization

To determine if different tasks rely on different parameters for the same fact, the authors adapt a localization framework (Bayazit et al., 2024) to identify sparse subsets of parameters (attention heads and MLP neurons) specific to individual (fact, task) pairs.

Objective: Learn binary masks $m$ $m$ over model parameters to minimize a loss function comprising four terms:
- Necessity: Removing the masked parameters must degrade performance on the target (fact, task) pair.
- Sufficiency: The masked parameters must be sufficient to recover performance on a corrupted prompt (where the subject is replaced by placeholders) when their activations are patched in.
- Specificity: Removing or patching these parameters should have minimal effect on the same fact under other tasks or other facts under the same task.
- Sparsity: The mask should be as small as possible.
Models: Experiments were conducted on OLMo-2-7B IT, OLMo-2-13B IT, and Gemma-2-9B IT.

3. Cross-Task Entanglement and CoT Analysis

Entanglement Metrics: The authors define a score to quantify how "cleanly" a (fact, task) encoding can be ablated without affecting other pairs.
Chain-of-Thought (CoT): The authors test whether CoT reasoning recovers knowledge inaccessible to direct answering by engaging task-specific parameters beyond those tied to the evaluation task. They compare direct answering vs. CoT under zero-ablation of localized parameters.

Key Results

1. Violation of Task-Invariance (Behavioral)

Limited Co-emergence: The co-emergence hypothesis is frequently violated. In 47.9% of (fact, task) pairs, a fact does not emerge on a target task by the expected step, even when the model is competent on that task and the fact has emerged elsewhere.
Statistical Rejection: The null hypothesis of conditional independence (no interaction between fact and task) is rejected at every training checkpoint ( $p \approx 0$ ). The interaction term explains up to 23% of the variance in the final model, indicating that knowledge retrieval is highly dependent on the task format.

2. Distinct Parameter Subsets (Mechanistic)

Task-Specific Encodings: The localization experiments successfully identified distinct, sparse subsets of parameters for each (fact, task) pair.
Necessity and Specificity: Zero-ablating the localized parameters for a specific (fact, task) pair caused significant performance drops (29%–89% relative decrease) on that pair. Crucially, performance on the same fact evaluated under other tasks, and other facts under the same task, remained largely unaffected (drops $\le 8\%$ ).
Sufficiency: Patching the localized activations back into a corrupted prompt successfully recovered performance (69%–102% reconstruction rate) for the target pair, confirming these parameters carry the necessary information.

3. Entanglement Patterns

Generation vs. Discrimination: Task format is a dominant predictor of entanglement. Discrimination tasks (MCQA, Verification) are consistently more entangled with other tasks than generation tasks (OPENQA, FITB).
Asymmetry: Facts acquired via generation tasks tend to co-emerge on other tasks, but facts acquired via discrimination tasks rarely co-emerge on generation tasks.

4. Role of Chain-of-Thought

Recovery: When a task's own localized parameters are ablated, CoT reasoning recovers a significant portion of the lost performance (dropping only 12%–30% vs. 20%–72% for direct answering).
Cross-Task Reliance: CoT suffers more collateral damage when other tasks' encodings are ablated compared to direct answering. This suggests CoT routes through multiple task-specific encodings to reconstruct the answer, rather than relying solely on the evaluation task's specific parameters.

Significance and Claims

The paper concludes that the "knowledge base" analogy for LMs is flawed because what the model knows and how it is asked are intertwined in parameter space.

Reliability and Controllability: The lack of task-invariance implies that knowledge editing or unlearning interventions targeting a single task format may leave the fact intact in other formats. Conversely, single-task evaluations provide only a partial view of a model's encoded knowledge.
Mechanism of Reasoning: The findings offer a mechanistic explanation for why CoT reasoning unlocks knowledge inaccessible to direct answering: CoT engages parametric encodings from multiple task formats, effectively bypassing the limitations of a single, task-specific encoding.
Future Directions: The authors suggest that training schemes encouraging task-invariant factual encodings (e.g., incorporating intermediate reasoning during training) could improve the reliability and controllability of LMs.

The authors remain modest, noting that their findings primarily describe well-known relational knowledge and that the interaction between these patterns and model scale remains an open question. They do not claim that LMs possess no shared knowledge, but rather that the storage is organized, at least in part, by task format.