PaCE: Parsimonious Concept Engineering for Large Language Models

Summary

This paper introduces Parsimonious Concept Engineering (PaCE), a novel framework for aligning LLMs by modifying their activation space. This framework constructs a large-scale concept dictionary in the activation space and partitions concepts into benign or undesirable categories. During inference, PaCE decomposes activations into these concept directions, removing undesirable components to reorient the LLMs' behavior towards alignment goals while preserving linguistic capabilities. The authors demonstrate the effectiveness of PaCE in tasks such as response detoxification, faithfulness enhancement, and sentiment revision.

Strengths

- This paper proposes a novel activation manipulation framework, different from the existing methods of vector addition and orthogonal projection. - The authors comprehensively and clearly articulate the deficiencies of existing methods and the ways in which the new framework addresses these deficiencies. - The authors validate the superiority of the proposed method in pursuing alignment goals and preserving linguistic capabilities compared to existing methods across multiple tasks.

Weaknesses

Overall, I did not find significant flaws that outweigh the contributions of this paper. The main weakness lies in the efficiency of the framework. During inference, using this method results in each token's prediction taking approximately 2-3 times longer than when not using it. In contrast, simply using vector addition only increases the time by about 30-40%. Additionally, although the concept vocabulary only needs to be obtained once, the concept direction dictionary needs to be constructed and stored for different LLMs.

Questions

* Could the authors provide more discussion on the computation time required and storage space (including memory usage and disk space) needed for constructing the direction dictionary? * Why does OrthoProj require significantly more computation time in Table 2 compared to other methods? This seems to conflict with the statement in Remark 3 that OrthoProj is a special case of the method proposed in this paper.

Rating

7

Confidence

4

Soundness

4

Presentation

4

Contribution

3

Limitations

Yes

Summary

This paper introduces an activation engineering framework for LLM alignment. PaCE addresses the challenge of reducing undesirable outputs while retaining benign concepts through a two-stage process: Concept Construction and Partition, and Activation Decomposition and Intervention. It constructs a large-scale concept dictionary and uses sparse coding to decompose LLM activations and selectively remove undesirable components. The authors demonstrate PaCE's effectiveness on tasks like response detoxification, faithfulness enhancement, and sentiment revising, showing improved alignment performance without compromising linguistic capabilities.

Strengths

1. This paper presents a refined activation engineering method for alignment, achieving outstanding results especially with the 13B model. 2. The method is provided with a clear visualization and detailed interpretation to illustrate its underlying principles.

Weaknesses

1. Quality of Concept Dictionary: the method's performance could be overly dependent on the quality and coverage of the concept dictionary. What measures are taken to ensure its comprehensiveness and that the dictionary is free from biases? 2. Static Activation: Static concept dictionary provides static activation while being decomposed, making PaCE unable to handle context-dependent concepts. 3. Without consideration of polysemy: Many common concepts have different meanings in different scenarios. For example, the concept ‘kill’ is harmful when referring to the act of terminating a life, while it is benign in the usage of ‘killing time’. Both two usages are common. However, this paper ignores the polysemy and treats all 30 pieces of contextual stimuli same in constructing concept dictionary construction. Thus, I am a bit concerned about the accuracy of the extracted concept direction. Besides, the design of retrieval promotes the diversity of the concept stimuli, while may undoubtedly increases the potential of the situation that the same word in different or even opposite meanings. 4. Misuse of representation reading algorithm (arXiv:2310.01405) in extracting concept directions: In Appendix B.2 Eq.6, authors take a difference between the activations of concept t_i and all other concepts t_j (j≠i). However, Zou et al. (in Sec. 3.1) state that stimuli in the pair should have different target concept or function. Take detoxification task as an example, the concept ’bias’ and ‘narrowed’ in Fig. 7 are both harmful content which needed to be removed. Since ‘bias’ and ‘narrowed’ have the same target function, representation reading algorithm will not pair them to take a difference, while this paper also pairs them. This is quite weird. 5. At the risk of compromising helpfulness: Following W3. and W4., the proposed concept directions may not be accurate and removing them may remove some benign concepts, thus limiting the model helpfulness. 6. Slower inference speed and time-efficiency: Since authors have claimed that ‘PaCE is computationally slower than VecAdd’, PaCE will undoubtedly slow down the original model’s inference speed, which is very important during model deployment.

Questions

1. Can PaCE handle concepts that have multiple semantic meanings or are context-dependent? 2. Are the VecAdd/OrthoProj baselines constructed based on the same concept dictionary with PaCE during the experiments? 3. How to choose the intervention decoder layers: Prior works (arXiv:2312.06681) find the optimal layer for steering by sweeping over all layers and performing intervention to evaluate the effect. Out of curiosity, how do authors choose the layers? 4. Can authors provide an evaluation on XSTest dataset (arXiv:2308.01263) when applying PaCE in detoxification task? I am a bit concerned about that PaCE will increase the false refusal rate on benign queries. 5. The MMLU results about VecAdd in Tab. 1 are questionable according to (arXiv:2402.19465). I doubt if the choice of steering layers for VecAdd is also the same as the setting of PaCE (last 19 layers) which results in the significant performance drop in MMLU since for VecAdd choosing the middle layers is better.

Rating

3

Confidence

5

Soundness

3

Presentation

3

Contribution

2

Limitations

In terms of experimental backbone models, it is recommended that the authors include more models (e.g., Mistral) not merely Llama2 Family for a comprehensive validation on the effectiveness of PaCE. In terms of linguistic capability evaluation, only fluency, perplexity and one multi-choice QA benchmark is limited. Generation tasks like open domain QA (e.g., MT-Bench), summarization (e.g., XSum) are supposed to be evaluated.

Summary

This work proposes an activation engineering framework for alignment, i.e., PaCE, consisting of a task-driven concept dictionary and linear decomposition algorithm. Thus, the toxic or harmful input can be decomposed into a new activation vector combination, that removes the harmful activation. This is a generalization and adaptive method compared to VecAdd and OrthoProj. However, the presentation is somewhat confusing. It is expected that the authors can clarify it and improve the paper's readability.

Strengths

This paper presents a novel safety alignment method based on activation engineering. The evaluation of PaCE is comprehensive. The results are insightful.

Weaknesses

There are some confusing parts (see detailed comments below). The presentation needs significant improvements. There is a large time and computation cost in the preparatory phase. The efficiency is questionable. It is necessary to analyze more serious jailbreaking examples, e.g., toxic，biased.

Questions

It is not clear how the dictionaries and scores are provided to the LLMs along with the input? Why does this linear degradation performed at each layer of the LLM not cause a decrease in fluency and efficiency? Is it possible to classify activation alignment as a scheme for the inhibition of malicious neurons or prompt purification? Compared to knowledge editing (e.g., "Detoxifying Large Language Models via Knowledge Editing"), what advantage of PaCE? Is the decomposition of semantic concepts the key to PaCE not affecting the normal response? According to Remark 3 of 3.3, I understand that the limitation that exist in VecAdd and OrthoProj are directional deletion and some limitations claimed in Figure 2. So, is dictionary construction the key to helping PaCE address limitations such as ambiguous deletions? How is it that independence between concepts can be solved by sparsity? Compared to RLHF, SFT, and other secure alignment schemes, I understand that the dataset of RLHF is replaced with a knowledge-driven concept dictionary, and the dynamic scoring during training is replaced with a task-driven concept splitter of GPT-4. If the understanding is correct, can the authors further explain the advantages? As far as I know, the Chat model of the LLaMA series has strong security alignment capabilities, but the jailbreak examples given by the author are relatively weak. If possible, can the authors give a comparison of the detoxification of the more offensive or malicious jailbreak examples on AdvBench, for example?

Rating

6

Confidence

4

Soundness

3

Presentation

2

Contribution

3

Limitations

The authors discussed limitations in Appendix B.6.

Summary

This paper presents a framework for aligning LLMs by using sparse coding techniques on a comprehensive concept dictionary. PaCE effectively controls and modifies neural activations of LLMs to achieve alignment goals such as response detoxification, faithfulness enhancement, and sentiment revising. The proposed method addresses limitations of existing alignment methods, such as costly fine-tuning, inadequate removal of undesirable concepts, and unnecessary removal of benign concepts. The paper demonstrates state-of-the-art performance in alignment tasks while maintaining the linguistic capabilities of LLMs.

Strengths

Originality: The paper introduces a unique approach to LLM alignment by combining sparse coding with a large-scale concept dictionary. This method provides a novel perspective on activation engineering, differentiating itself from traditional parameter fine-tuning and prompt engineering approaches. Quality: The methodology is well-developed, with thorough explanations and clear diagrams illustrating the concept construction, partitioning, and activation intervention processes. The use of sparse coding for accurate activation decomposition is well-motivated and justified. Clarity: The paper is well-written, with a clear structure and logical flow. Each section builds on the previous one, providing a comprehensive understanding of the proposed framework. The experimental setup and results are presented in a detailed and easy-to-follow manner. Significance: The proposed PaCE framework addresses critical issues in LLM alignment, offering a scalable and efficient solution. The state-of-the-art performance demonstrated in various alignment tasks highlights the practical significance and potential impact of this work on the broader AI community.

Weaknesses

Computational Efficiency: While the paper demonstrates that PaCE outperforms existing methods in alignment tasks, the computational efficiency could be further optimized. The comparison with OrthoProj shows a significant improvement, but PaCE is still slower than VecAdd. Future work could focus on enhancing the speed of the proposed method. Generalization to Other Models: The paper primarily evaluates PaCE on llama 2 7B, 13B models. It would be beneficial to demonstrate the generalizability of the framework across a wider range of LLMs to strengthen the claims of broad applicability. Interpretability of Concept Dictionary: While the paper provides examples of concept vectors and their partitioning, a deeper analysis of the interpretability and semantic consistency of the concept dictionary would be valuable. Understanding the intrinsic properties of the concept vectors could offer further insights into the alignment process. Reproducibility: The paper mentions the availability of the PaCE-1M dataset and plans to release the source code. Ensuring that all experimental details, including hyperparameters and specific configurations, are thoroughly documented will be crucial for reproducibility and wider adoption of the proposed framework.

Questions

Computational Efficiency: Are there any ongoing efforts or planned approaches to optimize the computational efficiency of PaCE, particularly in comparison to VecAdd? Model Generalization: Have you tested PaCE on other LLMs beyond llama 2 7B, 13B models? If so, could you provide insights or results from those experiments? Concept Dictionary Analysis: Can you provide more detailed analyses or case studies on the interpretability and semantic consistency of the concept dictionary used in PaCE? Ablation Studies: The ablation studies provide valuable insights into the contributions of various components of PaCE. Are there any additional ablation results, particularly focusing on the impact of different sparsity-promoting regularizers or dictionary sizes?

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors adequately address the limitations and potential societal impacts of their work. They discuss the computational overhead and the need for extensive concept dictionaries, as well as the importance of ensuring ethical considerations in the alignment of LLMs. Constructive suggestions for improvement include exploring more efficient algorithms for sparse coding and investigating the scalability of PaCE to even larger concept dictionaries. The paper could also benefit from a deeper discussion on mitigating potential biases in the concept partitioning process.