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

Authors: Haiyan Zhao, Heng Zhao, Bo Shen, Ali Payani, Fan Yang, Mengnan Du
Primary area: interpretability and explainable AI
Keywords: Large language model, Explainability, Probing, 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.

Decision: accept (poster)
Reviewer scores: 8, 5, 8, 6

Reviews and discussion:
<thread>
creation date:
1730132131615

thread_type:
official_review

<thread_object>
creator: reviewer_cu7j

<content>
summary:
This paper presents a framework, Gaussian Concept Subspace (GCS), for interpreting concept representation within LLMs. Traditional approaches to probing concepts in LLMs rely on single concept vectors derived from linear classifiers, which can be unstable and lack robustness. The GCS method enhances this by modeling each concept as a Gaussian distribution in the representation space, allowing for more nuanced, multidimensional representations. The experiments demonstrate that GCS-sampled concept vectors can effectively describe specific concepts learned by LLMs, similar to single concept vectors. Additionally, these sampled vectors are shown to be effective in downstream interventions, successfully guiding LLMs to generate more desirable outputs, such as positive sentiment steering in text generation.

strengths:
- The concept vectors are inherently variable across different datasets used for training the linear classifier. By modeling concepts through Gaussian distributions, the proposed approach intuitively captures a robust representation of each concept, reducing dependency on specific datasets. This approach is simple and straightforward for enhancing the robustness of concept vectors.

weaknesses:
- The authors mentioned that concept vectors "can vary significantly depending on the dataset used to train the classifier and the training process." However, as shown in Figure 2, the cosine similarity among concept vectors derived from different datasets consistently exceeds 0.75. While GCS increases the cosine similarity, it is unclear how critical this improvement is. Though the concept vector is significantly unstable for lower layers according to Figure 3, the GCS also shows low accuracy and cannot address this issue.

- Table 1 indicates that, in inference-time intervention, GCS achieves a stronger steering effect than single concept vectors, while coherence scores increase as well. However, GCS is not consistently superior to single concept vectors, as the results vary based on the strength parameter. Statistical tests should be conducted across different parameters to substantiate the significance of these improvements. Additionally, the criteria for highlighting the table are unclear, making them misleading. Furthermore, as mentioned in the questions, the authors should clarify why the sampled vectors lead to such improvements.

- The idea of modeling concepts with Gaussian distributions to capture their multifaceted nature is intuitive. However, it is unfortunate that the experiments do not demonstrate that the Gaussian distribution effectively models such multidimensional subspaces. For instance, it is interesting to see if the intermediate vector between "love" and "comedy" movies represents the concept of "love comedy".

questions:
How about using the mean vector rather than sampled vectors? As the sampled vectors are affected by some noises, the mean vector would be more robust for representing the concept and steering model's outputs.

Do you have any intuitive explanation for why GCS improves inference-time intervention compared to a single concept vector? I’m uncertain about the rationale behind this improvement.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730703246135

thread_type:
official_review

<thread_object>
creator: reviewer_qyqr

<content>
summary:
Proposed new methodology for mapping model internals to concepts, by extracting a gaussian distribution of concept vectors, rather than a single vector. 
This method is more robust to variations in the sampled training set, and achieves comparable steering performance and post-steering coherence.

strengths:
Originality: Good. Creative extension of the current concept vector extraction method that solves a big issue with the current method (robustness)

Quality: Good. experiment methodology and results seems solid and justified, though I did not check math and other implementation details in detail. I would prefer a more robust comparison with single-vector method (e.g. more diverse concepts and text types) before I can confidently say this method is empirically superior, but a priori seems likely that GCS will outperform single-vector for most purposes without being significantly more computationally costly 

Clarity: Good. no significant barriers to quick skim reading

Significance: Good. I can see this being the new standard for concept vector extraction/steering, as it's basically a pareto improvement on the existing concept vector methods, without being very costly/complicated to implement. I can imagine future work in the representation engineering literature being facilitated by the authors' library

weaknesses:
I was unable to identify any substantial weaknesses. 

Some minor suggestions:
- It's not clear to me how "reproducing hierarchical concepts" and "similarity between sampled and observed vectors" correspond to measures of "plausibility" and "faithfulness" respectively. Would be nice if you elaborated on why this is the case.
- Some comparison of the coherence/joyfulness scores provided by GPT vs humans would be nice. Just a tiny sample as a sanity check for whether GPT's scores are way off would do alot for your paper's soundness, as your main results hinge on GPT's evaluations being similar enough to human evaluations.
- Similarly, would be nice to include human-generated text (e.g. google search results for [concept], or joyful/angry reviews from IMDB dataset) on top of GPT-generated text for the training set. Just as a sanity check that GPT-text for [concept] is not too far off from human-text. (But I understand that this is costly)
- Would be nice to check how using training texts beyond movie reviews (e.g. joyful/angry tweets) would affect the extracted concept vectors and steering performance.

questions:
No further questions beyond the issues raised in Weaknesses section
</content>
</thread_object>

</thread>
<thread>
creation date:
1730796528341

thread_type:
official_review

<thread_object>
creator: reviewer_upan

<content>
summary:
This paper introduces the Gaussian Concept Subspace (GCS) framework, which aims to estimate the subspace representing specific concepts within large language models (LLMs). The authors extend the traditional single concept vector approach by modeling the concept subspace using a Gaussian distribution. The effectiveness of GCS is demonstrated through its faithfulness and plausibility across multiple LLMs with different sizes and architectures. The paper also showcases the practical application of GCS in real-world inference-time intervention tasks, such as emotion steering, where it balances steering performance and maintaining fluency in natural language generation tasks.

strengths:
1. The authors conduct extensive experiments to validate the faithfulness and plausibility of GCS across multiple LLMs, including different sizes and architectures. 
2. GCS reveals hierarchical concept structures that align with human understanding, providing insights into how concepts are represented within LLMs.
3. The paper is well-organized and easy to follow.

weaknesses:
1. The use of Gaussian distributions in representation learning is not entirely new. The paper could benefit from a clearer distinction between GCS and other probabilistic models used in similar contexts.
2. The paper primarily focuses on a specific set of tasks and datasets. To fully establish the significance of GCS, the authors should explore its applicability to a wider range of tasks and domains.

questions:
See above
</content>
</thread_object>

</thread>
<thread>
creation date:
1731286076792

thread_type:
official_review

<thread_object>
creator: reviewer_jclm

<content>
summary:
The paper proposes a robust alternative to the standard practice of learning linear probes in LLMs to find concept representations (for steering).
Specifically, the proposed "Gaussian Concept Subspace" (GCS) approach models the concept representation as a multivariate Gaussian (with diagonal covariance), thereby capturing the variance in the representations of a concept.
The overall procedure is to first train several linear probes using different probing datasets (generated by a LLM); estimating their mean and variances; and then sampling several concept vectors according to the learned Gaussian distribution (within a $1\sigma$ range).
In a set of experiments, it is shown that the resulting GCS vectors are faithful to the concept, aligns with known hierarchical semantic relations (in a topic hierarchy), and can be used to improve the robustness of steering tasks (in sentiment steering).

strengths:
- Good motivation and clear introduction. The paper is generally well-written and easy to follow.
- I think it is a meaningful endeavor to model the variance across different probe vectors for a single concept in LLMs, especially knowing that the probe vectors can be unreliable (see, e.g., [Tan et al., 2024](https://arxiv.org/abs/2407.12404), which I think deserves a mention in the paper by the way). The paper proposes an intuitive and sensible approach for this.
- Most of the experiments are well-designed and the results are convincing. I think the plots for plausibility experiments are particularly clear and informative.
- It's very interesting to see that, in the intervention experiment for sentiment, the $1\sigma$ samples from GCS (collectively) outperform the mean difference or the single-probe vector. (Curious to see if this generalizes to other concepts, but that's probably beyond the scope of this paper.)

weaknesses:
- I think the paper makes a meaningful contribution, but it is relatively light on noting the limitations of its main approach. Here are the main ones in my view:
    - The most obvious drawback of the main approach is its reliance on large samples obtained using a high-quality LLM (10,000 samples from GPT-4o per each topic concept). This appears necessary to obtain a variance estimate on the GCS, so it feels inherent to the approach. Perhaps this needs to be mentioned in the introduction and the discussion, as it could be a significant limitation for certain use cases.
    - Another limitation, which I think is fine as long as it is mentioned in the paper, is the assumption of Gaussianity with diagonal covariance for the concept vectors. That said, this is still far better than having no variance information and is not a knock on the paper's contribution.
    - For the intervention step, it appears that the steering is done by applying each of the 1,000? sampled steering vectors and averaging the results. It's good to know that this makes the intervention robust, but it can also make the approach computationally expensive. It would be good to see some discussion on this.
- I think what the similarity score for faithfulness is somewhat confusing. In Section 3, the authors state that we want the sampled concept vectors to be similar to each other as much as the observed concept vectors are. But in Figure 2, the sampled vectors are a lot more similar to each other than the observed ones, which is expected as the sampling restricts to the "within $1\sigma$" range, but then the paper appears to suggest that this is ideal. So, what do we actually want out of this metric? Doesn't Figure 2 just end up being an illustration of how large the variances $(\sigma_j)$ are in each layer?
- While this is an understandable choice, I do think it should be noted that the evaluations for intervention experiments are entirely model-based (GPT-4o) and may not be accurate.
- Finally, I feel that the related work on hierarchical concepts is light in the paper, despite the fact that the plausibility experiments highlight the clusters of topic concepts found by GCS. Some suggestions on representations of hierarchical concepts include [Nickel and Kiela, 2017](https://arxiv.org/abs/1705.08039); [Chen et al., 2021](https://arxiv.org/abs/2104.03869); [He et al., 2024](https://arxiv.org/abs/2401.11374); and [Park et al., 2024](https://arxiv.org/abs/2406.01506).

questions:
- Intro: For the sake of clarity, I think it should be mentioned somewhere that the concepts being considered here are binary, requiring positive/negative prompt pairs. You can just say "following prior work" and reference, e.g., the ITI paper, and maybe give a few examples of what concepts are being considered here.
- p. 3: have you tried removing the independence assumption and estimating a full covariance matrix for the concept subspace (maybe for a smaller model, to reduce $d$)? If so, how does it compare to the current approach?
- p. 3: "randomly sample vectors ... within $1\sigma$" is slightly ambiguous to me. Am I correct in assuming that you  first sample from the learned Gaussian distribution and then reject the sample if it is outside the 1-sigma boundary? In Algorithm 1 it says "in $1\sigma$", which is even more ambiguous. I think this part can be reworded for clarity.
- p. 4: What is the similarity function? If it is cosine similarity, then wouldn't it be skewed by how far away the mean vector is from the origin? If you use centered cosines or some distance metric, do these results change?
- Eq. 8: What exactly is $\bf C$? ${\bf w}_i$ is already indexed by the dataset number, so ${\bf w}_i \in {\bf C}$ is confusing to me.
- p. 5, "Implementation Details" first paragraph: I was wondering about this very detail since Section 3, and it feels like crucial information that shouldn't be left in the "details" part. My suggestion is to bring this up earlier when you introduce the sampling part. I am also curious to see how some of the results change when you nevertheless resort to distributional comparisons, e.g., for plausibility (what does the KL divergence between the concepts show?).
- p. 5, lines 265--267: do you specifically mean that you subsample 1,000 (with/without replacement?) from the 10,000 samples for each concept? This feels like an important detail to be clarified.
- p. 6, lines 319--322: it first says that "We sample concept vectors ... ranging from $1\sigma$ to $5\sigma$, etc." and then says "we focus on sampling concept vectors within the 1$\sigma$ range". Some rewording here appears necessary.
- Figure 5 (PCA): which concept vectors exactly do you use to learn the PC space? Also, could you just project the Gaussian mean vector rather than the mean of 1,000 sampled vectors?
- Table 1 (Intervention): for $1\sigma$--$5\sigma$, am I correct in thinking that you sampled 1,000 steering vectors, applied them, and then averaged the final ratings? This was confusing to me and could be clarified. I'd also recommend adding error bars if this was the case.
- Intervention (App. H.4): in the rating prompt, why specifically ask "repetitive or chaotic" instead of "coherent" or "fluent"?
- Appendix C, lines 737--738: this feels like a critical detail. Exactly how is the "range" of ${\bf h}^\ell$ found?
- Appendix H: I don't think it's a good idea, for science, to show only best performing examples. At least I'd want to see random samples along with the best ones.
- Appendix I: are these random or hand-picked samples within each category?

- Minor stylistic suggestions:
    - p. 3, lines 115--117: maybe instead of repeating the question from the intro, pose it in a theorem environment ("Question 1") in the intro and then simply refer to it here?
    - Sections 3: maybe call the "dataset" specifically as "probing dataset" and give a sense of how many samples one may need per a probe vector. Otherwise, I think the reader can get worried about needing an excessively large dataset for each concept.
    - p. 4, line 197: I think this meant "generalizable" instead of "generalized"?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732504737067

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response to All Reviewers

comment:
We sincerely thank all reviewers for providing many constructive comments and helpful feedback. We are encouraged that they found our contributions to be meaningful and convincing (JcLm), creative and solid (qyQR), simple and straightforward (cu7j), and well-organized (uPaN).

To address the evaluation concerns raised, we have conducted additional experiments:
- We performed preliminary human sanity check, which yielded results comparable to GPT-4o's evalution performance.
- We conducted an additional experiment demonstrating that a similarity value of 0.75 represents a relatively low value.
- We generated PCA visualizations with guassian mean vectors, which produced results nearly identical to the average of sampled vectors.
- We expanded our intervention experiments to show that the proposed GCS framework can effectively steer LLM behavior beyond movie reviews.
- We performed experiments validating that the proposed Gaussian distribution models multidimensional concept subspaces through "Bird in Village" experiment.

Please find our detailed responses to specific questions and concerns below. We have incorporated all comments and comprehensive experimental evaluations into the revised manuscript, with changes highlighted in blue. We are grateful to the reviewers for their valuable suggestions to improve our work.

Paper4838 Authors
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper moves from modeling concept representations as directions in a representation space to modeling them as normal distributions (with diagonal covariance) in the representation space. Reviewers appreciated the novelty and clarity, and are optimistic that characterizing the variance will improve the usefulness of LLM representation approaches. The reviewers offered extensive (though relatively minor) feedback, which should be incorporated in the camera ready revision.