JoPA: Explaining Large Language Model's Generation via Joint Prompt Attribution

We thank the reviewer for the valuable suggestions.

W1. Insufficient contribution to the ML and DL field.

We believe our work considers an important problem of explaining LLMs' entire generation considering the jointly effect of prompt tokens, offering a unique perspective in LLM interpretability. This initial attempt offers a general and efficient interpretation framework to locate important tokens for the whole generation sequence. And we have demonstrated several use cases for future works to use our framework on potential applications.

Method applications: Our method has potential usages on improving the model safety and helping in the model diagnosis as discussed in Appendix A.3. By applying our method to detect malicious tokens in the input prompt for jailbreak attacks, the model can refuse to respond to malicious queries, thereby making the attacks ineffective. Moreover, by identifying the important input attributions, the relevance between the input prompt and output generations will be assessed, contributing to the model diagnosis.

Without considering unsafe input: If there are unsafe contents in the input, our method can serve as a defense mechanism, enabling the model to refuse responses to malicious prompts, as discussed earlier. Additionally, our work focuses on explaining the joint effects of input prompts on output generations. While it touches upon safety-related applications, addressing safety issues can be a potential use case but is not the primary focus or main topic of our paper.

Not limited to prompt attribution task: The primary goal of our work is to provide explanations for the effects of joint tokens in input prompts on output generations, with a particular focus on capturing their semantic interactions. We uniquely formulate this unresolved challenge as an optimization problem and design a specialized algorithm to address it. While we acknowledge the reviewer's observation that our algorithm may have applications in other research areas, this does not diminish its contribution to the field of XAI. Furthermore, it is important to note that prompt compression and prompt optimization are distinct research topics that fall outside the scope of this work.

W2. Theoretical proofs lack persuasiveness.

Local optima: The local optima proof guarantees that our algorithm could reach a local optimum after sufficient iterations, which is useful in showing that the algorithm is well-behaved without the issue of divergence. Therefore though intuitive, the proof ensures a nice property of the algorithm. We agree that having a theoretical guarantee on the convergence rate or global convergence would be nicer, however, we acknowledge that it is very challenging and still an open problem, considering the complicated optimization landscape. To the best of our knowledge, existing works for explaining LLMs have not yet established such theoretical analysis.

Use of gradient: The proof is designed to demonstrate that the algorithm converges to the local optimum, regardless of the mechanism used to propose swaps. The specific strategy for generating swaps under the guidance of gradient is orthogonal to the convergence proofs. Moreover, gradients are commonly used in attribution and feature importance tasks, as they capture the sensitivity of the model's predictions with respect to the input perturbations. The use of gradients has demonstrated success in numerous works and aligns with the common pratices.

The effectiveness of our algorithm is demonstrated through its strong performance across all metrics in Table 1 and the convergence plot in Figure 4. In Table 1, JoPA achieves the lowest probability ratio compared to other methods, indicating its ability to significantly reduce the generation probability of the same output by the masked input. Furthermore, as the number of training iterations increases, the generation probability of the output from the modified prompt stabilizes, as shown in Figure 4. These experimental results collectively highlight the effectiveness of our algorithm.

Dear Reviewer ZoiX,

We greatly appreciate your initial comments and we totally understand that you may be extremely busy at this time. However, we kindly hope you can take a moment to review our responses to your concerns. We greatly value any feedback you can provide and would appreciate it if you could consider updating your rating if your questions have been adequately addressed. Additionally, we are more than happy to answer any further questions you may have before the discussion concludes.

Best Regards,

Authors of Paper7871

We thank the reviewer for the valuable questions and suggestions.

W1. A global convergence proof would be nicer.

The local optima proof guarantees that our algorithm could reach a local optimum after a sufficient number of iterations, which is useful in showing that the algorithm is well-behaved without the issue of divergence. Moreover, our proof demonstrates the training procedure satisifes the necessary conditions for convergence, maximizing the probability discrepency of generating the same output by different input.

We agree that a thorough theoretical analysis on the global convergency can greatly strengthen our work. However, we ackowledge that it is very challenging and still an open problem, considering the complicated optimization landscape. To the best of our knowledge, existing works for explaining LLMs have not yet established such theoretical analysis. We believe this is a valuable perspective and will leave this opportunity as our future work.

Although we cannot theoretically prove the global convergence of the algorithm, its effectiveness is empirically demonstrated through its strong performance across all metrics in Table 1 and the convergence plot in Figure 4. In Table 1, JoPA achieves the lowest probability ratio compared to other methods, indicating its ability to significantly reduce the generation probability of the same output by the masked input. Furthermore, as the number of training iterations increases, the generation probability of the output from the modified prompt stabilizes, as shown in Figure 4. These experimental results collectively highlight the effectiveness of our algorithm.

W2. "Aim to solve" rather than "solve" the optimization problem.

Thanks for the suggestion. We will change it to "aim to solve" to avoid possible misunderstanding.

W3. Minor tweaks.

Thank you for the detailed suggestions in improving our writing. We will revise them in our final version.

Questions:

Q1. The value of k; the number of tokens within a prompt.

Table 1 displays the performance of JoPA when $k=3$ (Line 363). We also show the variation of probability ration (PR) with different $k$ in Figure 3. These results demonstrate the effectiveness of JoPA under various $k$. The summary of the user prompt length for different dataset is in Appendix A.1. In addition, we use a system prompt for all datasets: "You are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe. Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure that your responses are socially unbiased and positive in nature.\n\nIf a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information.", which covers 117 tokens.

Q2. How often are the final mask positions chosen from the initial (say) top-2k? Could we safely reject all but the top-2k tokens from consideration?

We found that 14.8% of the data samples from the Alpaca dataset do not have at least one final mask position that appeared in the initial top-$2k$ positions. As there are some samples whose final masks are not selected from the top-$2k$ tokens, relying solely on the $2k$ tokens obtained from the initialization would not be sufficient to achieve the final optimal mask, which suggests the necessity of exploring other tokens via our probabilistic search mechanism.

Q3. Suppose we start with two different mask patterns (for a given k, say : 0..01.....1 and 1.....10..0). What percentage of the time would Algorithm 1 arrive at the same final mask for both? Figure 4 suggests that even a sequence of gradient-proposed local moves towards a global minimum slowly.

In order to verify the stability of our algorithm with various inistializations, we compare the log-prob $logp_{\theta}(y|m\odot x)$ of our gradient-guided initialization with the random initialization by using different $k$ across different iterations N.

The results above indicate that regardless of the initializations, the algorithm consistently converges to a similar local optimum after a sufficient number of iterations, highlighting its stability. The gradient-based initialization could provide a better starting point to achieve a faster convergence than the random initialization. Furthermore, our algorithm is computationally efficient. As shown in Table 2, it outperforms Captum in efficiency, requiring only about 15 seconds to identify the optimal mask.

We thank the reviwer for the constructive suggestions.

W1. More baselines, e.g., leave-one-out and attention-based methods.

We appreciate the suggestion, and will include the baselines in our experiments.

• Leave-one-out method: We have compared with Captum [7] in Table 1, which is essentially a leave-one-out baseline by masking the token sequentially and use the variations in the output generation probability to indicate the importance of the masked token (Line 112-115).
• Attention-based methods: We use the last layer attention and the average of all layer attentions of the for measuring the token importance and use the top-$k$ tokens as the most important joint attributions.

We extend Table 1 to include the performance of attention-based methods on Llama-2 (7B-Chat) model:

Compared to various attention-based baselines, JoPA achieves the best performance across all metrics on the Llama-2 (7B-Chat) model.

W2. Highlighting the novelty of the proposed method.

We believe the primary novelty of our work lies in our formulation and framework for explaining the entire generation of LLMs through the joint attribution of the input prompts. This is a largely under-explored area, as discussed in Line 038-061: while existing methods either focus on classification/next-token prediction tasks or overlook the joint effects of input attributions on the output generations, our approach aims to quantify the importance of tokens while considering their semantic interactions. To achieve this goal, we formulate and address this novel problem of measuring the combinatorial effects of tokens through a combinatorial optimization algorithm: the proposed algorithm follows a certain general principles and adapts for this unique problem of joint attribution. Additionally, our work demonstrates significant computational efficiency compared to the baseline method, Captum, with JoPA being approximately 100 times faster. The effectiveness and efficiency of our method are demonstrated through experimental results in Table 1 and Table 2.

Yes I have already read the response and increased my score to 6 since the response addressed my major concern on missing baseline comparisons.

We thank the reviewer for the insightful suggestions.

W1. Concept and logic of the counterfactual explanation not being fully explained. Need enough case studies to support this argument.

Counterfactual Explanation: We apologize for any confusions. Our work JoPA is to explain the generation of LLMs by identifying a small changes (i.e., masking few tokens) needed to alter the outcome (i.e., minimizing the probability of the original generation). This can be understood as counterfactual explanation, since it answers a "what if" question: if these input tokens are masked, the LLM would not generate the response. Namely, by applying the mask $m$ on the input prompt $x$, the generation probability of its original response $y$ would be largely decreased (i.e., $p_{\theta}(y|m\odot x) \ll p_{\theta}(y|x)$). This suggests that the newly generated output $y'$ resulting from the masked prompt would differ substantially from the original outcome $y$, as measured by our metrics in Table 1: the semantic similarity (SentenceBert) and word similarity (BLEU and ROUGE-L) between $y$ and $y'$ are much smaller than those from baselines.

Case Study: Thank you for the suggestion, and we have included the following case studies to show the outcome changes. For instance, given the input prompt to the model:

Follow the law of supply and demand, describe what would happen to the price of a good if the demand increased.

The model would generate the response: "\nIn economics, the law of supply and demand states that the price of a good is determined by the interaction between the quantity of the good that...". After masking the 3 explanatory tokens: 'supply', 'price', and 'demand', the generated response would change significantly: "However, I must inform you that the question you've provided is not factually coherent, and I cannot provide an answer that may not be accurate or safe...".

Similarly, given the input prompt for the model to fill in the blank:

Fill in the blank with a word or phrase The most successful team in the NBA is the ____

the original response is: "The most successful team in the NBA is the Golden State Warriors". If we masked out the explanatory tokens 'most', 'successful', and 'NBA' from the input, the output would change to: "\nBased on the context you provided, the word that best fits the blank is "best." So, the sentence would read:\nThe team in the office is the best."

The case studies of diverse tasks above demonstrate that the new responses generated from the masked inputs differ significantly from the original outputs, which highlights the impact of these explanatory tokens in influencing and altering the model's output.