Beyond Single Concept Vector: Modeling Concept Subspace in LLMs with Gaussian Distribution

Abstract

Probing learned concepts in large language models (LLMs) is crucial for understanding how semantic knowledge is encoded internally. Training linear classifiers on probing tasks is a principle approach to denote the vector of a certain concept in the representation space. However, the single vector identified for a concept varies with both data and training, making it less robust and weakening its effectiveness in real-world applications. To address this challenge, we propose an approach to approximate the subspace representing a specific concept. Built on linear probing classifiers, we extend the concept vectors into Gaussian Concept Subspace (GCS). We demonstrate GCS’s effectiveness through measuring its faithfulness and plausibility across multiple LLMs with different sizes and architectures. Additionally, we use representation intervention tasks to showcase its efficacy in real-world applications such as emotion steering. Experimental results indicate that GCS concept vectors have the potential to balance steering performance and maintaining the fluency in natural language generation tasks.

GCS Framework

< >

Evaluation

RQ1: How faithfully do GCS-sampled concept vectors represent the original concepts?(Faithfulness)

• Experiment 1

Llama-2-7B, layer 30	Gemma-2-7B, layer 26	Llama-2-13B, layer 38

Figure: Histogram of cosine similarity within observed concept vectors, sampled concept vectors, and between both sets for concept "Bird".

Llama-2-7B	Gemma-2-7B	Llama-2-13B

Figure: Layer-wise average cosine similarity from the second layer to the penultimate layer of observed concept vectors, sampled concept vectors, and between both sets for concept "Bird".

• Experiment 2

Llama-2-7B	Gemma-2-7B	Llama-2-13B

Figure: Accuracy of observed and sampled concept vectors aross varying models.

• Conclusion
  • Observed concept vectors are similar to sampled concept vectors in representation space.
  • GCS-sampled concept vectors are more general and robust on classifying concept-related data.

RQ2: To what extent do explanations derived from GCS-sampled concept vectors align with human expectations about the hierarchies among model’s learned knowledge?(Plausibility)

• Experiment 1

Llama-2-7B	Gemma-2-7B	Llama-2-13B

Figure: Heatmap of concept average cosine similarity of 16 concepts across Llama-2-7B, Gemma-7B, and Llama-2-13B. The 16 low-level concepts are grouped into four high-level categories: the first 4 rows/columns represent sports events, the next 4 represent populated places, followed by 4 for animals, and the last 4 for movie genres.

• Experiment 2

Llama-2-7B	Gemma-2-7B	Llama-2-13B

Figure: PCA visualization of 16 concepts across Llama-2-7B, Gemma-7B, and Llama-2-13B. Low-level concepts belonging to the same high-level concept category share the same color.

• Conclusion
  • Low-level concepts within the same high-level category typically exhibit higher similarity scores.
  • Some high-level categories, such as “Sports Event”, shows very high internal similarity across all models, suggesting these concepts are closely related in the models’ representations.
  • Some high-level categories, such as “Populated Place” and “Animal”, show stronger correlation, align with human intuition about real-world relationships between these concepts.

RQ3: Can the proposed GCS method effectively mitigate unwanted behaviors in LLMs?(Effectiveness on emotion steering)

Figure: An illustration of inference-time intervention with a LLM.

Figure: Comparison of generated texts with and without intervention.

• Conclusion
  • GCS-sampled concept vectors can better balance steering performance and maintaining the fluency in natural language generation tasks.