PaCE: Parsimonious Concept Engineering for Large Language Models

Dear Reviewer rWAy,

Thank you for your insightful reviews and kind support for acceptance. It is our pleasure to reply to your comments.

Time efficiency

Thank you for discussing PaCE's time efficiency. We would like to emphasize that the main contributions of the paper are (1) a new dataset of concept dictionary and (2) a sparse coding process using this dictionary to decompose latent LLM representations. VecAdd and OrthoProj do not sufficiently model the representation space, and PaCE adopts a dictionary to have enhanced alignment performance.

• PaCE is faster than OrthoProj ([18, 79], Table 2) and we provide an option of a time-efficient solver (OMP, see discussion below).
• PaCE has much higher averaged safety (Table 1), faithfulness and sentiment revising scores (Table 4) than VecAdd, OrthoProj, Prompting, and the vanilla model.

Hence, the goal of this paper is to offer a competitively novel solution to representation manipulation for alignments of LLMs. We gladly take your comment and explore alternative optimization techniques for Equation 1 in the paper from the sparse coding literature that significantly reduces inference times. Below we briefly describe various solvers.

• The elastic-net solver (Line 213) optimizes Equation 1 to obtain coefficients simultaneously.
• We also consider Orthogonal Matching Pursuit (OMP), a fast greedy solver [D1, D2]. OMP iteratively adds concepts to the support based on maximum coherence with the residual and updates the residual until the support size k is reached. In other words, k is the number of non-zero coefficients.

Comparing solvers on response detoxification (\S 4.1) shows certain setups of OMP are faster than Elastic Net (e.g., OMP with $k=50$ is an order of magnitude faster with 12.3% lower safety score). This demonstrates that a greedy solver improves speed at a performance cost. We will include the experiment in $\S4.1$ of the paper based on your suggestions.

Computing concept dictionary for different LLMs

Thank you for bringing up the point. We would like to explain our motivation for collecting the dataset, and how the cost may not be a concern in terms of our efforts in open-sourcing and low amortized cost.

• (Motivation) We noted that the scale of existing datasets on the concept stimulus (e.g., RepE [79], ActAdd [64] in the paper) was limited. This incentivized us to (1) collect a large stimulus dataset to better model the concepts and (2) pre-compute the concept dictionary in an LLM representation space.

• (Open-sourcing) Future researchers need not rerun dataset collection as we will open-source our dataset of human-readable concepts, stimuli, and concept dictionaries for LLaMA2 7B/13B and the Mistral 7B. Pre-computed concept partitions for downstream alignment tasks will also be shared."

• (Low amortized cost) Each target model's concept dictionary needs to be extracted only once, with costs amortized across tasks. For example, extracting our dictionary for LlaMA2-7B took 25 minutes on an NVIDIA A40 machine and the dictionary was used for all evaluations on that model.

Specific time and space requirements for constructing the dictionary

We provide the statistics of computation and storage space needed for constructing the dictionary. We will include them as a new subsection (Appendix C.4) in our paper.

• GPU memory usage for hosting the dictionary during each decomposition:
  • LLaMA2 7B: 0.31 GB
  • LLaMA2 13B: 0.45 GB
• Local disk storage usage for storing dictionaries (for all involved 19 layers):
  • LLaMA2 7B: 5.80 GB
  • LLaMA2 13B: 8.26 GB
• Time taken to compute and extract the dictionaries (NVIDIA A40 GPU machine):
  • LLaMA2 7B: 25.2 minutes
  • LLaMA2 13B: 40.8 minutes

Computation time of OrthoProj

• We appreciate your excellent comments on this point. The difference lies in the fact that the OrthoProj method computes a projection at multiple layers of the LLM to maintain a fair estimation of the coefficient. While PaCE uses a concept dictionary to model the representation space, OrthoProj does not sufficiently model it and could remove the contribution of desirable concepts. To maintain a stable steering process, the coefficient needs to be re-estimated at every decoder layer. Instead, VecAdd pre-sets a coefficient as a hyper-parameter and reuses it across the rest of the layers. Similarly, in experiments ($\S 4$ of the paper), PaCE performs decomposition at the first encountered intermediate layer and reuses it across the rest of the layers. In practice, one can freely tune the VecAdd coefficient through grid search and choose the layers where PaCE decomposition happens.
• We find the choice of PaCE is sufficient for steering the LLM. We hypothesize this is because sparse coding can find a more suitable steering direction that directly reflects the desired concept change as compared to simply running orthogonal projection. Lastly, thank you for carefully reading the remarks. Your time and effort are greatly appreciated.

Also, our paper’s Appendix B.4 includes the implementation details of frameworks and hyper-parameters for inferring on open-source LLMs. Based on your suggestion, we will include the justification in a new subsection Appendix B.5 and update the existing description to better clarify the difference between OrthoProj and other methods.

Hopefully this alleviates your concerns; if not, please engage with us during the discussion phase.

Thanks and regards,

Authors of Submission #8804

[D1] Orthogonal matching pursuit: Recursive function approximation with applications to wavelet decomposition. Pati et al., 1993.

[D2] Orthogonal matching pursuit for sparse signal recovery with noise. Cai et al., 2011.

Thank you for your time on the paper. We are happy to address your comments below.

Quality of the concept dictionary

Excellent question! The efforts are made in two-fold: the concept words (indexes) and the context stimuli (contents) that define the semantics of each word.

• For the concept indexes, we took the most frequent 40,000 words from the Brown Corpus ([14] in our paper), which models standard American English [C2] to ensure that our dictionary covers comprehensive concepts.
• We defined the meaning of each concept via retrieving external knowledge from Wikipedia and instructing GPT4 to generate context stimuli. The semantics of a concept may develop over time: e.g., “mouse” was not regarded as a digital input device before the era of PC [C1]. We consider that GPT’s parametric knowledge is fixed at a timestamp. We propose to retrieve external knowledge for stimulus synthesis because the external knowledge is more updatable and explicit, which serves as enriched descriptions for concepts and helps GPT to generate more grounded and diverse scenarios. To further add to the transparency, we will open-source the dataset so that the community can reuse and improve on it.

Words have polysemy

We agree with you that concepts may have multiple meanings. We tried to be frank on this in our paper on Lines 954-956: “Current practice usually finalizes a single vector per concept by SVD, but theories on polysemanticity and recent studies on the causal models of language suggest that a concept might be better represented by a union of subspaces, each corresponding to different semantic meanings”.

However, we argue that, while VecAdd and OrthoProj may be affected by polysemy, PaCE's overcomplete dictionary allows accurate analysis of the target representation through sparse decomposition. Tables 1 and 2 in the paper show PaCE outperforming OrthoProj and VecAdd on linguistic metrics. We further provide empirical evidence on MT-Bench (QA) and XSum (summarization) with explanation for polysemy not harming PaCE’s model helpfulness.

Empirical evidence. Table 1 and Table 4 in $\S 4$ show that PaCE is better than OrthoProj and VecAdd on linguistic capability metrics (fluency, perplexity, MMLU). Based on your suggestions, we conduct further experiments on MT-Bench and XSum, whose implementation is detailed in the one-page PDF. PaCE maintains superior linguistic capability compared to VecAdd and OrthoProj.

Explanations. We attribute this to the large-scale dictionary with sparse coding, explained as follows.

• Since the dictionary is large, concepts with single and clear semantics are involved. E.g., if stimuli of “kill” may have different meanings, there exist other more polarized concept vectors such as “murder” (more harmful) and “spend” (more benign).
• Sparse coding aims to choose the fewest concepts to reconstruct the latent representation (i.e., parsimony). For the sake of argument, assuming the sentence is about “killing time” and the vanilla LLM has the correct semantic understanding of its benignness, the latent representation of the whole sentence will be closer to concepts such as “spend” and “time” rather than string-matching to “kill” (which in your setup could have mixed harmful and benign senses). As the sparse coding of the target representation promotes the parsimonious selection of concepts with monosemantics, it helps to represent benign contexts correctly without assigning significant weights to ambiguous terms like "kill".

Response for pairing in representation reading; time efficiency; context-dependent concepts

Thank you for your insightful discussion. Due to OpenReview's limitations on figures, please refer to the one-page PDF under Global Response for our detailed responses.

Compliance to benign requests in the XSTest dataset

Thanks for inquiring on this very insightful dataset. PaCE decomposes input based on underlying semantics, not string-matching. It removes only malicious parts of the representation, preserving compliance with benign requests. We evaluate benign requests in XSTest, following the XSTest classification: I: full compliance, II: full refusal, III: partial refusal. Compliance is minimally affected. More detailed evaluations are in the one-page PDF.

Middle layers for representation manipulation

Thank you for this comment that allows us to correct a typo: In Line 909 (Appendix B.4), we wrote “activation vectors are extracted from the last 19 layers of the target LLM’s decoder layer.” We meant to say “activation vectors are extracted from the last-29th to the last-11th layer (totaling 19 layers) of the target LLM’s decoder layers”. Thus, the VecAdd in our experiments did use the middle layers.

PaCE on Mistral-7B

Please refer to our response and evaluation in Generalization to Other LLMs (Reviewer Cdr5).

Thanks and regards,

Authors of Submission #8804

[C1] Biases in Large Language Models: Origins, Inventory, and Discussion. Navigli et al., 2023.

[C2] Role of the Brown Corpus in the History of Corpus Linguistics. Kholkovskaia et al., 2017.

Response to Reviewer s54k (1/2)

Dear Reviewer s54k,

Thank you for the thoughtful feedback and constructive questions. We address the concerns and provide additional evaluations to validate our work. We add the suggested improvements to enhance the quality of our work.

Preparatory Phase Concerns

Thank you for recognizing our efforts in dataset collection. Early in developing PaCE, we noted that the scale of existing datasets on the concept stimulus (e.g., RepE [79], ActAdd [64]) was limited. This incentivized us to (1) collect a large stimulus dataset to better model the concepts (2) pre-compute the concept dictionary in a larger LLM's representation space. Below we try to address your comments on computational costs in this phase.

Future researchers need not rerun dataset collection as we will open-source our large dataset of human-readable concepts and stimuli. Further, each target model needs a concept dictionary to be extracted only once and this cost is amortized across different tasks on this model. For example, extracting our dictionary for LlaMA2-7B model took 25 minutes on our eight-card GPU machine and was used for all evaluations on that model. Lastly, despite this one-time preparation cost, the decomposition and intervention stages of PaCE are flexible in time efficiency and task performance. In Table 2 of the paper, we show that PaCE is faster than OrthoProj with a higher performance.

Providing dictionary and scores to LLM

We appreciate your question which helps improve the clarity of the paper! In short, dictionaries are a collection of concept vectors and are frozen for representation decomposition.

• First, the LLM takes an input prompt (e.g., malicious requests)
• Compared to a regular LLM, the representation engineering framework [B1] extracts the activations at each decoder block of the transformer. Such extraction results in a vector corresponding to the input prompt, which can then be modified for steering in different ways.
• For PaCE, this steering has two main stages:
• Stage 1 pre-computes the large-scale concept dictionary offline and the partition (i.e., scores) of which concepts represent benign/harmful concepts.
• In Stage 2, at inference time, our approach extracts the representation of an input prompt and uses sparse coding to decompose this as the linear combination of atoms in our frozen dictionary. We then modify this linear combination by removing undesirable components and proceeding with inference in the LLM with the detoxified representation. We hope this clears up any misunderstandings.

We will include the justifications in Appendix B.3 of our paper and revise other parts of the manuscript correspondingly.

Effects on fluency and efficiency, and the role of decomposition to preserve the normal response

Thank you for acknowledging the advantages of PaCE. Below we explain in two parts: fluency (Part 1) and efficiency (part 2).

Part 1: The key to keeping fluency is to maintain the benign components of the activation of user input prompts. PaCE handles it by decomposing the user input’s activation along benign and harmful concepts, while VecAdd and OrthoProj do not sufficiently model the concepts in the representation space. The decomposed solution of PaCE accurately estimates the benign components and preserves them during the intervention process.

Part 2: The decomposition is efficient because the optimization problem for Equation (1) in the paper is convex with known fast solvers such as active-set (used in this paper) and OMP (see the discussion with Reviewer Cdr5/rWAy,).

Connection between activation alignment and the inhibition of malicious neurons or prompt purification

Thank you for bridging PaCE with other alignment schemes. If the concept vectors in PaCE are sparse/axis-aligned, i.e. each concept vector only has a few specific non-zero elements, then yes – subtracting such a vector from the activation is inhibiting the support neurons. We observe that concept vectors are dense and not axis-aligned, which means that many neurons are involved.

Detoxification on AdvBench

AdvBench uses GCG [B2] to adversarially optimize a jailbreak suffix for a harmful behavior request. In the table below, we show the LlaMA-7B-Chat safety score (%, ↑) on the effective set of suffix attacks for AdvBench harmful behavior set. We observe that PaCE outperforms other baselines. We also note that the outperformance of PaCE in our paper’s jailbreaks is more significant than that in suffix attacks. This could be because story-telling and roleplay jailbreaks (used in our paper) contain more complex and entangled concepts. Under this scenario, PaCE decomposes the target representation and well estimates the malicious component, while VecAdd and OrthoProj do not model the space sufficiently. In the AdvBench case, instead, the optimized adversarial suffix can be regarded as the text-space inversion of straightforward malicious concepts. PaCE and other defense mechanisms in latent space and prompt space shall effectively defend these suffixes more easily. We will include the evaluation in $\S 4.1$ of our paper.

(continued in the comment below)

(rebuttal continued here)

Response to Reviewer s54k (2/2)

Dictionary construction to address limitations such as ambiguous deletions (Part 1); Sparsity to handle independence between concepts (Part 2).

Part 1: Yes, PaCE has a large-scale dictionary modelling sufficient concept directions in the latent space, which allows PaCE to accurately analyze the compositionality of benign concept directions in the target.

Part 2: When the concept vectors are not independent, infinitely many linear combinations of concept vectors can reconstruct the activation equally well. Therefore, one needs regularization to break ties.

• Sparsity serves as one regularization, which models the belief that the activation can be written as a linear combination of only a few concept vectors. This is favorable since shorter explanations are deemed more interpretable [B3, B4].
• A standard regularizer is the sum-of-squares of the combination coefficients, as in ridge regression. However, this tends to give combinations that use most of the concept vectors in the dictionary,
• and since the dictionary is large this is not efficient.

Advantages of PaCE over RLHF, SFT, and Knowledge Editing (KE)

Your analogy of the correspondence among different alignment paradigms is inspiring. We would like to point that the main advantage of PaCE over the mentioned RLHF, SFT, and KE are two-fold.

• Training-free: RLHF, SFT, KE all need to tune the parameters of LLM, which potentially degrade the well-structured priors of the pre-trained LLM. Taking a step back, even if LoRA is adopted for these paradigms, the training/tuning incur significant computation and memory costs. PaCE does not modify the parameters of LLM and requires no training. It better preserves priors of LLM, provides a low-resource alignment solution, and retains the general linguistic capabilities.
• Interpretable and Adaptive: The solved coefficients are an accurate interpretation of how a user input’s representation is composed in the concept space. Also, when a new alignment goal is set, RLHF, SFT, and KE need to collect sufficient task samples and tune the LLM on the new dataset. In contrast, PaCE just needs to run the concept partitioner through PaCE-1M, which is expected to be much faster and more convenient.

We are glad to respond to further comments or other concerns (if any) during the discussion period. We appreciate your insightful suggestions which have helped to validate and strengthen our framework.

Thanks and regards,

Authors of Submission #8804

[B1] Representation Engineering: A Top-Down Approach to AI Transparency. Zou et al., 2023.

[B2] Universal and Transferable Adversarial Attacks on Aligned Language Models. Zou et al., 2023.

[B3] Interpretable by design: Learning predictors by composing interpretable queries. Chattopadhyay, et al., 2022.

[B4] Leveraging sparse linear layers for debuggable deep networks. Wong et al., 2021.

Thanks for the authors' response. Most of my concerns have been addressed. It would be great if the authors incorporate those insights into the revised version. I will increase the score from 5 to 6.

Dear Reviewer Cdr5,

We greatly appreciate your constructive feedback and kind acknowledgment of the novelty and advantages of our approach. It is our pleasure to address your comments and provide clarification below.

Computational Efficiency

Thank you for acknowledging the task (e.g., detoxification) performance of PaCE. Yes, it is possible to make the method even faster, at the trade-off of task performance. While our method jointly solves for the entire set of coefficients simultaneously, we can solve sparse decomposition in a greedy manner to speed up the computation. That is, the coefficients are obtained procedurally (i.e., one by one) in the order of energy minimization. The main benefit of using greedy algorithms is that the user can arbitrarily set the desired threshold on the number of non-zero coefficients in pursuit of better time efficiency. Based on such observation, we recap specific decomposition details in our paper and show an alternative speed-up version:

• The elastic-net solver (Line 213) used in the experiments in $\S 4$ is designed to compute an exact solution of coefficients for the optimization problem (Equation (1), Line 186-189). The coefficients are obtained simultaneously.
• Motivated by your inquiry, we additionally evaluate Orthogonal Matching Pursuit (OMP), a fast greedy solver in the compressed sensing literature [A1, A2] for our activation decomposition. On a high level, OMP iteratively 1) adds to the support the concept that has maximum coherence with the unexplained residual, and 2) updates the residual by solving a least square using the new support. It stops when a pre-defined maximum size k of support is reached. Intuitively, the k is the number of non-zero elements in the solved coefficients.

We compare two solvers on response detoxification ($\S 4.1$) in the table below.

• Notably, with a small choice of k=50, OMP is an order of magnitude faster than Elastic Net, while it achieves a safety score 12.3% lower than that of Elastic Net.
• Overall, the observations in this experiment validate that a greedy solver can improve computational speed at the cost of safety performance.

Generalization to Other LLMs

We evaluate PaCE on Mistral-7B-Instruct (version 1) for detoxification and observe PaCE’s superior performance while preserving linguistic capability. Mistral-7B’s strong instruction-following leads to a lower initial safety score and higher MMLU performance. Hence unlike LLaMA2-7B-Chat, the prompting baseline for detoxification does not significantly harm linguistic performance.

Interpretability of Concept Dictionary

Figure 8 and Appendix D.2 show semantic similarity between concept vectors by clustering the first 10,000 concepts in our PaCE-1M dictionary. The observed semantic structures (salient clusters and decision boundaries) indicate that the target LLM has an activation space that understands and organizes semantic information of the concepts, enabling further analysis and manipulations in PaCE. Figure 10 and Appendix D.3 show the activation space's utility for concept retrieval, indicating close coherence between the target and relevant concept vectors. With full respect to your and Reviewer s54k’s insights on this, we will add a new section Appendix D.4 for discussions of analyzing other properties (e.g., sparsity, magnitude, direction).

Ablation Study

Thank you for inquiring about ablation studies on key components of PaCE. Figure 6 shows detoxification performance for LLaMA2-13B-Chat across different dictionary sizes, with safety scores increasing and converging around 9000-10000. For the regularization design, we evaluated different setups of $\tau$ in Equation (2) that implement the sparsity-prompting regularizer. The results are shown in the table below:

The results show that the regularization with $\tau$=0.95 (our choice stated in the paper Appendix B.4) yields the best safety performance among the five choices. Pure ridge regression ($\tau$=0) and pure lasso regression ($\tau$=1) do not perform as well as the mixed regularization strategy.

Reproducibility

We appreciate your emphasis on reproducibility. Appendix B.4 details our frameworks, experiments, and hyper-parameters for inferring on open-source LLMs, concept curation, and knowledge retrieval. We have also included details of our computing resources on GPUs. Upon acceptance, we will open-source the PaCE-1M dataset and PaCE implementation with documentation. We will continue maintaining the project (e.g., computing concept dictionaries and partitioning the concepts with new expert models) for future open-source LLMs and new alignment tasks.

We will include the experiments above in $\S4$ and Appendix B.5 of the paper, and we will revise the writing based on your suggestions. We appreciate your valuable insights which have helped to validate and strengthen our framework.

Thanks and regards,

Authors of Submission #8804

[A1] Orthogonal matching pursuit: Recursive function approximation with applications to wavelet decomposition. Pati et al., 1993.

[A2] Orthogonal matching pursuit for sparse signal recovery with noise. Cai et al., 2011.