What are They Thinking? Delineation, Probing and Tracking of Concepts in LLMs

Technical Summary: Delineation, Probing, and Tracking of Concepts in LLMs

Problem Statement
As Large Language Models (LLMs) increasingly influence decision-making, there is a critical need to understand their internal "thought processes." While previous research has demonstrated that linear probes can detect specific linguistic properties (e.g., parts-of-speech, tense) and concepts (e.g., time, location, truth) within LLM embeddings, existing methods for concept detection face significant limitations. Sparse Autoencoders (SAEs), for instance, require computationally expensive unsupervised training and lack control over which specific concepts are detected, making them unsuitable for targeted monitoring. Other approaches, such as prompting an LLM to identify a concept, are too costly for continuous monitoring as they require separate model invocations. This paper addresses the need for a low-cost, scalable method to create specific probes that can monitor the presence and absence of implied concepts within LLM embeddings during normal operation.

Methodology
The authors propose a three-stage pipeline for creating and utilizing concept probes:

1. Concept Delineation: The core innovation is a semi-automatic method for creating high-quality binary datasets for specific concepts.

   • Template Generation: The authors created 30,000 sentence templates from Project Gutenberg texts, filtered to ensure they focus on human subjects and contain coherent structures.
   • Example Generation: Using a prompted LLM (gpt-4o), they generated pairs of examples for each concept (present vs. absent) by mimicking the syntactic structure of the templates while altering semantics to reflect the target concept. This approach aims to prevent "label leakage" where the model might learn spurious correlations rather than the concept itself.
   • Refinement and Labeling: To ensure reliability, a prompted LLM classifier re-labeled all generated examples. The prompts were iteratively refined to achieve >90% accuracy on manually labeled validation sets. Datasets were filtered to retain only pairs with opposite labels, ensuring the probe learns the concept rather than structural similarities. Four concepts were selected: ambition, investigation, democracy, and envy.

2. Probe Training and Evaluation:

   • Training: Linear probes (binary classifiers) were trained on the embeddings extracted from various layers of the target LLMs. The authors tested two representative embedding strategies: the final token embedding ($N$-th) and the average of all embeddings in a layer.
   • Control Tasks: To verify that the probes detect concepts encoded in the embeddings rather than learning the task independently, the authors employed control tasks:
     • Randomization: Training on randomized embeddings or labels resulted in accuracy dropping to ~50%, confirming the probes rely on meaningful information in the embeddings.
     • Parameter Reduction: Probes were constrained to fewer than 80 parameters (via PCA). High accuracy was maintained even with significant compression, suggesting the embeddings encode the features robustly.

3. Concept Tracking (Waxing and Waning):

   • Story Datasets: The authors constructed 50 three-paragraph stories for each concept. The target concept was present only in the transition sentences connecting the paragraphs and absent elsewhere.
   • Dynamic Monitoring: The trained probes were applied to the embeddings of expanding input contexts (word-by-word). This allowed the authors to observe how the probability of the concept's presence fluctuates as the context evolves, specifically tracking the "waxing and waning" of the concept signal.

Key Results

• Concept Inference: The study confirms that LLMs (tested on Llama-3-8B, Gemma-2, and Qwen2.5 families) infer the four target concepts. Probes achieved accuracy significantly above 50% across most layers, with the embedding layer (layer 0) showing surprisingly high accuracy (87%) for ambition when using average embeddings, though this did not hold for longer contexts.
• Probe Efficiency: Probes achieved strong performance with very few parameters. For ambition, accuracy remained above 75% with only 40 parameters and plateaued around 80 parameters, well below the model's embedding dimension (4,096 for Llama-3-8B).
• Tracking Dynamics: The probes successfully tracked the presence and absence of concepts in real-time as input context expanded. When using final token embeddings, the probe output rose above the 0.5 threshold when the concept was present in the text and fell below it when absent. Conversely, cumulative mean embeddings diluted the signal, failing to capture these dynamics, indicating that concept inference in LLMs involves more than simple bag-of-words representations.
• Layer Specificity: Different layers in different models were found to be optimal for tracking specific concepts. For example, in Llama-3-8B, layer 13 was best for tracking ambition, while layer 31 was best for investigation.

Significance and Claims
The paper claims to take the "first steps" toward a capability for creating broad sets of probes for monitoring LLMs. Its primary contributions are:

1. A reproducible, semi-automatic methodology for delineating concepts into binary datasets that minimize accidental label leakage.
2. Demonstration that linear probes can detect implied concepts with high accuracy using minimal parameters (<80), making them computationally cheap.
3. Evidence that these probes can serve as low-cost, continuous monitors, capable of tracking the dynamic "waxing and waning" of concepts as context expands.
4. The release of the created datasets (templates and concept-specific examples) to enable further research in LLM explainability.

The authors maintain a modest tone, acknowledging limitations such as the potential for generated data to follow specific distributions and the reliance on LLM-based classifiers for labeling. They posit that while their current work focuses on four concepts, scaling this process could enable the monitoring of new models and the exploration of causal links between inferred internal concepts and generated outputs, thereby advancing LLM safety and explainability.