💬 NLP

Efficient Knowledge Probing of Large Language Models by Adapting Pre-trained Embeddings

This paper introduces PEEK, a computationally efficient method that adapts pre-trained embedding models with a linear decoder to accurately predict Large Language Models' factual knowledge at scale, thereby avoiding the high costs of traditional forward-pass probing techniques.

Original authors: Kartik Sharma, Yiqiao Jin, Rakshit Trivedi, Srijan Kumar

Published 2026-01-27

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Kartik Sharma, Yiqiao Jin, Rakshit Trivedi, Srijan Kumar

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, incredibly smart library (a Large Language Model, or LLM) that has read almost everything on the internet. You want to know exactly what facts it has memorized and what it has forgotten.

The Problem:
Currently, to check if the library knows a specific fact (like "Garifuna is a population of Guatemala"), you have to walk up to the librarian, ask the question, and wait for an answer. If you want to check millions of facts, this is like trying to read every single book in the library one by one. It's slow, expensive, and if the library is a "black box" (you can't see how it thinks), you can't peek inside its brain to see what it knows without asking it directly.

The Solution (PEEK):
The authors of this paper propose a new tool called PEEK (Proxy Embeddings to Estimate Knowledge). Think of PEEK as a high-tech "knowledge radar" or a shadow librarian.

Instead of asking the giant library every single question, PEEK uses a smaller, faster, pre-trained "scanner" (an embedding model) to guess what the big library knows.

Here is how it works, step-by-step:

The Training Phase (Teaching the Scanner):
First, the researchers take a small sample of facts and ask the big library, "Do you know this?" They record the library's answer (Yes/No, or how confident it is). Then, they take a smaller, faster scanner and show it the text of those facts. They tweak the scanner's settings (like tuning a radio) until the scanner's "guess" matches the big library's actual answer.
- Analogy: Imagine you have a master chef (the LLM) and a sous-chef (the embedding model). You taste the master chef's soup and tell the sous-chef, "This tastes salty." You keep adjusting the sous-chef's seasoning until they can predict exactly how salty the master chef's soup will be, just by looking at the ingredients list.
The Prediction Phase (The Radar):
Once the scanner is tuned, you don't need to ask the big library anymore. You just feed a new fact into the scanner. The scanner instantly predicts, "The big library knows this," or "The big library doesn't know this."
- Analogy: Now, instead of asking the master chef to taste every new ingredient, the sous-chef looks at the ingredient list and instantly says, "The chef knows how to cook this."

What They Found:

Speed and Accuracy: This "scanner" method is incredibly fast and accurate. In their tests, it could predict what the big library knew with up to 90% accuracy without ever asking the big library a single question.
Text vs. Maps: They tried two types of scanners: one that reads sentences (like a sentence-embedding model) and one that reads maps of connections (like a graph neural network). They found that the sentence scanner worked much better. It seems the big library organizes its knowledge more like a collection of sentences than a rigid map of connections.
No Need to Peek Inside: Unlike other methods that require looking at the library's internal brain waves (hidden layers), PEEK only needs the text of the fact. This means it works even if the library is a "black box" that you can't touch.

Why This Matters:
The paper suggests that before we let these giant AI libraries help us with important tasks (like healthcare or law), we should use this "scanner" to quickly check their knowledge gaps. It's a cheap, efficient way to say, "Hey, this AI doesn't know about X, so let's fix that before we let it talk to a patient."

In a Nutshell:
PEEK is a way to build a fast, cheap proxy that learns to mimic a giant AI's knowledge base. Once trained, this proxy can scan millions of facts to tell you what the giant AI knows, saving you the time and money of asking the giant AI directly.

Technical Summary: Efficient Knowledge Probing of Large Language Models by Adapting Pre-trained Embeddings

Problem Statement

Large Language Models (LLMs) acquire vast amounts of factual knowledge across diverse domains during generative pre-training. However, due to stochastic learning objectives and noisy training data, mapping the specific knowledge landscape of an LLM remains challenging. Current methods for probing this knowledge—such as crafting specific prompts, investigating hidden representations, or estimating epistemic uncertainty—require forward passes through the underlying LLM. These approaches are computationally expensive, time-consuming, and often infeasible for extremely large or black-box models, limiting the ability to efficiently identify factual gaps before deployment in knowledge-intensive applications like healthcare and law.

Methodology: PEEK (Proxy Embeddings to Estimate Knowledge)

The authors propose PEEK, a framework that leverages pre-trained embedding models as proxies to estimate the knowledge of LLMs without querying the LLMs directly. The core hypothesis is that embeddings can effectively encode factual knowledge and, when adapted, can predict LLM outputs with high accuracy.

1. Knowledge Probing Functions

To establish ground truth for the proxy models, the authors define four distinct probing strategies to determine if an LLM ( $M$ ) knows a specific fact ( $f$ ):

Binary Generation: Prompting the LLM with yes/no questions regarding the truthfulness of a statement.
Binary Logits Generation: Extracting the logits of the expected tokens ('yes'/'no') from the LLM to measure "external confidence."
Binary Activation Prediction: Training a linear head on the hidden representations (last token activations) of the LLM to predict truthfulness, measuring "internal confidence."
Fact Generation: Using open-ended generation (e.g., biographies) and labeling atomic facts as supported (known) or unsupported (unknown) based on external knowledge bases.

2. Proxy Embedding Models

The framework utilizes two types of pre-trained embedding models to represent facts:

Sentence Embeddings: Models (e.g., MPNET, NVE2, Linq, GTE) that encode sentences into low-dimensional spaces, trained to represent semantic similarity.
Graph Neural Networks (GNNs): Specifically ULTRA, which encodes knowledge graph triples (head, relation, tail) using shortest-path information.

3. Proxy Tuning

Since pre-trained embeddings are not natively aligned with specific LLM knowledge, the authors adapt them using a small training set of facts identified via the probing strategies. The adaptation involves minimizing a loss function to match the embedding-based prediction ( $\hat{P}_\theta$ ) with the LLM-probed output ( $P_M$ ). Two tuning strategies are employed:

Linear Tuning: Updating only the final linear decoder layer ( $W$ ) while keeping the embedding parameters ( $E$ ) frozen.
Low-Rank Tuning (LoRA): Updating parameters via a low-rank matrix addition.

The loss function is binary cross-entropy for discrete outputs (0/1) and knowledge distillation loss (with temperature $T$ ) for continuous outputs (logits/activations).

Key Contributions

Novel Probing Approach: Introduction of PEEK, which uses trainable linear heads on pre-trained embedding models to probe LLM knowledge, avoiding costly forward passes through the target LLM.
Multi-Strategy Validation: Uncovering factual knowledge through diverse strategies, including binary queries, free-form generation, and internal activation analysis.
Comprehensive Evaluation: Extensive experiments across 4 LLMs (Llama-3.1-8B, GPT-4o, GPT-4o-mini), 3 datasets (DBP100k, YAGO310, FactScore), and 7 embedding models.
Efficiency and Accuracy: Demonstration that embeddings can predict LLM knowledge on held-out sets with up to 90% accuracy and high AUC scores, significantly outperforming baselines like majority classifiers and random guessing.

Experimental Results

Performance: Sentence embedding models (particularly NVE2 and Linq) consistently outperformed graph embeddings (ULTRA) and baselines. In the binary generation task on DBP100k, NVE2 achieved up to 91.52% accuracy and 86.89% AUC for GPT-4o-mini.
Model Comparison: While Linq performed best on YAGO310, NVE2 excelled on DBP100k. Notably, the graph-based ULTRA model performed poorly, suggesting LLMs do not induce facts via the same shortest-path mechanisms as GNNs on knowledge graphs.
Tuning Efficiency: Linear tuning was found to be as effective as LoRA tuning, suggesting that a simple linear layer is sufficient to adapt embeddings to an LLM's knowledge space.
Robustness: Results remained stable across varying sampling percentages (0.1% to 10%) and negative sample ratios, though performance generally improved with the inclusion of negative samples due to class balance.
Limitations Observed: The method struggled with Llama-3.3-70B, where embedding models failed to estimate knowledge better than random chance, indicating that estimating knowledge in larger, more complex models remains challenging.

Significance and Claims

The paper claims that PEEK provides a scalable, cost-effective alternative to traditional knowledge probing. By using proxy embeddings, researchers and practitioners can:

Identify factual gaps in LLMs at scale without querying the models.
Gain insights into the internal inductive bias of LLMs regarding factual landscapes.
Optimize Retrieval-Augmented Generation (RAG) by disregarding facts the model already knows, thereby improving context utilization.

The authors emphasize that while their benchmark relies on Wikipedia-derived knowledge, the methodology is not inherently restricted to closed domains. However, they modestly note that specialized domain knowledge may require training domain-specific embeddings from scratch. The work concludes that embedding models should be utilized in the pre-deployment stage of LLMs to preemptively find and fix factual gaps.