Efficient Knowledge Probing of Large Language Models by Adapting Pre-trained Embeddings
Kartik Sharma, Yiqiao Jin, Rakshit Trivedi, Srijan Kumar
TL;DR
The paper tackles the challenge of mapping what large language models actually know, given their stochastic training and the cost of probing via forward passes. It introduces PEEK, a framework that adapts pre-trained embedding models with a lightweight linear head to estimate an LLM's knowledge from a knowledge base without querying the model. The authors define four probing functions (binary generation, logits, activation, and fact generation) and two embedding families (sentence embeddings and knowledge-graph encoders), and show that adapted embeddings can predict LLM knowledge on held-out facts with up to about 90% accuracy across multiple LLMs and datasets. The work demonstrates that sentence-embedding proxies outperform graph-based encoders, highlights LoRA's limited advantage for tuning, and suggests practical uses in identifying knowledge gaps and guiding retrieval-augmented generation.
Abstract
Large language models (LLMs) acquire knowledge across diverse domains such as science, history, and geography encountered during generative pre-training. However, due to their stochasticity, it is difficult to predict what LLMs have acquired. Prior work has developed different ways to probe this knowledge by investigating the hidden representations, crafting specific task prompts, curating representative samples, and estimating their uncertainty. However, these methods require making forward passes through the underlying model to probe the LLM's knowledge about a specific fact, making them computationally expensive and time-consuming. To bridge this gap, we propose $\textbf{PEEK}$ or $\textbf{P}$roxy $\textbf{E}$mbeddings to $\textbf{E}$stimate $\textbf{K}$nowledge of LLMs, by leveraging the pre-trained embedding models that effectively encode factual knowledge as text or graphs as proxies for LLMs. First, we identify a training set of facts known by LLMs through various probing strategies and then adapt embedding models to predict the LLM outputs with a linear decoder layer. Comprehensive evaluation on $3$ Wikipedia-derived datasets, $4$ LLMs, and $7$ embedding models shows that embeddings can predict LLM knowledge on a held-out set with up to 90 % accuracy. Furthermore, we find that sentence embedding models are more suitable than graph embeddings to predict LLM knowledge, shedding light on the underlying representation of the factual landscape. Thus, we believe that knowledge-adapted embeddings can be used to identify knowledge gaps in LLMs at scale and can provide deeper insights into LLMs' internal inductive bias. The code and data are made available at https://github.com/claws-lab/peek.