Unveiling Causal Relationships Among Candidate Output Tokens in Large Language Models: Towards Interpretability and Control

---

[Weakness 4] Thank you for pointing out the ambiguity in the description of the algorithm setup. CID algorithm can be controlled by changing the values of two hyperparameters:

$d$: the number of tokens with largest logits that will be considered in CLA. Selecting a larger $d$ will result in more cause-effect token pairs to be selected by CLA, and thus more tokens are subject logit changes in CID.

$h$: the logit change applied to cause and effect tokens detected by CLA. A larger $h$ will alter the token distribution for word prediction more aggressively.

CID+ has a more aggressive configuration than the CID algorithm. Specifically, CID has $(d, h) = (2,5)$ and CID+ has $(d, h) = (5, 10)$. We have included the explanation and the specific configurations in the revised manuscript in Section 4.4.

---

[Question 2] We investigate the impact of introducing bias by adding perturbations to the LLM, recognizing that such perturbations could potentially alter the causal relationships among tokens. It is important to study whether these causal relationships remain consistent when the perturbations are small, as significant changes could undermine the validity of our causal analysis. Our findings indicate that as the perturbations become smaller, the causal relationships we identify remain similar. This suggests that the causal structures are robust to small perturbations added to the LLM. By demonstrating the robustness of these relationships under minor biases, we provide evidence that the causal connections are inherent to the model and not merely artifacts of the perturbations introduced.

---

[Question 3] The CID algorithm, in practice, shares the same complexity with the CLA heuristic. They first require one single pass of inference to obtain the initial candidate token logits. Then for the top-k candidate tokens, it requires to run inference with a dropped layer on each token pair to find their approximate causal-effect relation. Denote the single pass complexity as O(pass), the time complexity of the CID and CLA heuristic is O($k^2$ pass),

Consider in practice $k$ is set as a very small number (e.g., 3) and one could control the frequency to activate the CID algorithm; the overall time complexity of CID remains within the same order of magnitude as standard inference.

We do observe that CID may spend more time answering a question compared to the standard inference method. This is not due to the complexity of CID itself, but rather because CID increases the lengths of the response. The intuition is that, by using 'effect tokens' such as 'while' instead of directly answering 'yes', the response generated by CID contains more elaboration.

We sincerely thank Reviewer zKHP for the efforts in reviewing our work and for providing constructive comments. We hope this response clarifies all your concerns, as outlined below. Please let us know if you have any further questions or additional concerns.

---

[Weakness 1] We appreciate your observation regarding the bias analysis. Our main argument in Section 3.2 is that as we introduce varying degrees of bias (by controlling the Bernoulli distribution of layer deletion probability), the causal relationships among candidate tokens remain similar. This is evidenced by the statistical similarity of the Markov equivalence classes we constructed.

While the reviewer's observation is accurate and intuitive—that the model outputs (logit values) tend to resemble those of the full model as the probability of deleting a layer approaches zero—we want to clarify that causal relationships are not determined by the output values themselves but rather by how the outputs are influenced by changes in the model. This influence is not directly reflected in the similarity of logits but is instead reflected in the changes caused by the deletion of layers. Whether these changes tend to be similar as the deletion probability approaches zero is not clear. Consequently, the increasing similarity of logits does not contradict our analysis in Section 3.2. We acknowledge that our original explanation may have been unclear and potentially misleading, and we have revised Section 3.2 in the updated manuscript to address this point more effectively.

---

[Weakness 2] We appreciate the reviewer's insightful comments on the Empirical Validation of CLA in Section 4.2. We acknowledge that CLA is a heuristic and not a formal causal discovery method. However, as highlighted in our general response, as long as the confidence regions in the ROC scatter plot lie above the y=x line, the results are statistically significant. This indicates that CLA's identified cause-effect pairs align with the ground truth from the PC algorithm more than would be expected by chance.

Moreover, we stress that CLA is a heuristic and not a formal technique for causal discovery. It is intended to quickly and approximately identify causal pairs. The interesting outcome of our study is that this heuristic method does indeed find causal pairs, and the results are statistically significant across many different LLMs. This suggests that despite its heuristic nature, CLA is effective in practical applications.

---

[Weakness 3 and Question 1] We acknowledge that no inference algorithm, including our CID algorithm, can guarantee to always output the correct token. Similar to widely used practical inference methods like top-k and top-p sampling, our CID algorithm is designed to enhance performance without guarantees of correctness in every instance. Despite this, these algorithms are highly valued for their ability to improve the quality of generated text in practice. Our empirical results on reasoning benchmarks demonstrate that the CID algorithm effectively improves model performance, validating its practical utility.

Regarding the suggestion for more theoretical evidence, we appreciate the feedback and understand the concern. While our causal analysis provides valuable insights and serves as a rigorous foundation for the CID algorithm, we acknowledge that the effectiveness of CID is primarily evidenced by our benchmark results. The empirical results from our experiments demonstrate the practical effectiveness of prioritizing effect tokens during decoding. We believe our paper makes two main contributions: (i) conducting a rigorous causal discovery analysis that offers theoretical insights into the relationships among tokens, and (ii) demonstrating the effectiveness of the CID algorithm through empirical results on reasoning benchmarks. It is also common in machine learning research to place significant value on empirical results, as they provide concrete evidence of an approach's effectiveness. These contributions together support the utility and validity of our approach.

We sincerely appreciate the reviewer for the time and effort invested in evaluating our work and for providing insightful and constructive comments. We have made every effort to address your concerns in the responses below. If there are any remaining questions or issues, please feel free to let us know.

---

[Weakness 1] The reason we used arithmetic reasoning datasets is that many prior works on reasoning and decoding have conducted experiments on these datasets. However, we acknowledge that we overlooked other downstream tasks. To address your concern, we apply CID and CID+ to Mistral-Nemo-Instruct on the Social IQa dataset and compare with the original decoding and DoLa decoding. The results are shown in the table below. We can see that CID and CID+ consistently improve over the original decoding by large gaps. CID is better than DoLa with raw prompts and worse with CoT prompts. We have included these results in the Appendix B of the revised manuscript.

---

[Weakness 2.1 – Why simply adding or subtracting] Thanks for pointing this out. Since the logit value appears in the exponent of the weight when calculating the actual sampling probability, adding a value to the logit effectively corresponds to scaling up the weight for that logit by a factor. While other interventions, such as scaling logits by a factor, could be considered, our empirical results show that this straightforward approach effectively improves performance. The value of $h$ is a hyperparameter selected based on empirical tuning. We would like to respectfully emphasize that CID is intended as a heuristic technique rather than a method designed for optimality. Its purpose is to empirically support our hypothesis that causal relationships between candidate tokens can be leveraged to improve the decoding process.

---

[Weakness 2.2 – Configuration of CID and CID+] We apologize for not explaining in detail how CID and CID+ are different and their specific configurations. CID can be controlled by changing the values of two hyperparameters:

• $d$: the number of tokens with largest logits that will be considered in CLA. Selecting a larger $d$ will result in more cause-effect token pairs to be selected by CLA, and thus more tokens are subject logit changes in CID.

• $h$: the logit change applied to cause and effect tokens detected by CLA. A larger $h$ will alter the token distribution for word prediction more aggressively.

CID+ has a more aggressive configuration than the CID algorithm. Specifically, CID has $(d, h) = (2,5)$ and CID+ has $(d, h) = (5, 10)$. We have included the explanation and the specific configurations in the revised manuscript.

---

[Question 1] This is an insightful question. We have also observed this phenomenon and believe it depends on the design of the heuristic. It is important to note that CLA is not intended to be a rigorous causal discovery algorithm but rather a fast heuristic for identifying causal pairs efficiently. An interesting outcome of our study is that this heuristic method does indeed find causal pairs, and the results are statistically significant across many different LLMs, as demonstrated in Section 4.2. Currently, our approach is to ignore such bidirectional pairs in the results. While this may not always capture the true causal direction, we believe the heuristic still serves its purpose effectively in many cases. In the future, we plan to explore alternative approaches or additional post-processing steps to refine the identification of causal directions further.

---

[Question 2] $L(v)$ represents a standard root mean square layer normalization, a commonly used technique for rescaling standardized inputs in various LLM architectures. The parameter $\gamma$ is a learnable scaling factor that is updated during training and remains fixed during inference. As a result, there is no manual selection or adjustment of $\gamma$ in our setting. Additionally, $\gamma$ is not shared across layers, meaning its value can vary from one layer to another.

---

[Weakness 3] We have briefly discussed in lines 231-234 how the PC algorithm is applied to extract cause-effect token pairs in our setting. We apologize that the discussion may not be clear enough, Here we provide more details:

Given an LLM and an input,

1. We repeatedly perturb the LLM to generate $k=1000$ samples of logit values for the candidate tokens, $\\{s_1,s_2,\ldots,s_k\\}$. Each time,the LLM is perturbed by applying Bernoulli random scalars with success probability 0.95 for the layers as described in Section 3.1, effectively removing some of the transformer layers.

2. We then apply the PC algorithm to the generated samples $\\{s_1,s_2,\ldots,s_k\\}$ with Fisher’s z independent test and a significance level of 0.9999. This is done using the causal-learn package[1] as footnoted on page 5. The PC algorithm outputs a causal graph between the tokens.

3. Finally we convert the source and destination tokens of each directed edge of the causal graph to a cause-effect pair.

We hope this provides enough details for understanding the application of the PC algorithm in our setting. We have added the description below in Appendix A in the paper.

[1] Causal-learn: Causal discovery in python. Zheng, Yujia and Huang, Biwei and Chen, Wei and Ramsey, Joseph and Gong, Mingming and Cai, Ruichu and Shimizu, Shohei and Spirtes, Peter and Zhang, Kun.

---

[Weakness 4 – CLA logic] We appreciate your suggestion to enhance the CLA algorithm by considering significant increases in token logits after layer ablation. While our current implementation focuses on tokens that drop out of the top candidates, incorporating other indicators of causal influence could improve the identification of causal pairs. As noted in our general response, we recognize that CLA is a heuristic and consider improving it an intriguing topic for future research.

---

[Weakness 5] Thank you for your valuable suggestion. While other causal mediation analysis methods do not provide decoding algorithms directly applicable to our context, we agree that comparisons with alternative improved decoding methods, such as DoLa[1], are highly informative. Following your suggestion, we applied DoLa to the Mistral-Nemo-Instruct model across all four datasets used in our paper. We adopted the recommended settings for long-answer reasoning tasks, such as GSM8K, as suggested by the authors of DoLa: applying DoLa to lower layers and setting the repetition penalty to 1.2 to reduce repetition in DoLa decoding.

The results, shown in the table below, indicate that CID+ performs significantly better than DoLa with raw prompting. When CoT is applied, DoLa outperforms CID on GSM8K, MAWPS, and MultiArith. However, DoLa struggled on the SingleEq dataset, where CID consistently improved over the baseline. These findings suggest that while DoLa shows strong performance in certain scenarios, CID demonstrates greater stability across datasets. We appreciate your suggestion, as it has helped provide a more comprehensive comparison. We have included these results in the Appendix B in the revised manuscript.

[1] DoLa: Decoding by Contrasting Layers Improves Factuality in Large Language Models. Yung-Sung Chuang, Yujia Xie, Hongyin Luo, Yoon Kim, James Glass, Pengcheng He. ICLR 2024.

[Weakness 2.3 – When to apply which] Based on our observations above and results reported in Table 2, our suggestion on when to apply which method is:

• When the LLM is small, CID+ is preferred as the reasoning ability of the model is limited and explicit causal reasoning through CID will be helpful.

• When the LLM is larger or it is used with CoT, CID is preferred.

• If the output of the LLM is well formatted without being specifically prompted to do so, it is better to leave the decoding unchanged.

We sincerely thank reviewer bhjj for taking the time to review our work and for offering valuable and constructive feedback. We hope that our responses below address all your concerns clearly and thoroughly. Should you have any additional questions or need further clarification, please do not hesitate to let us know.

---

[Weakness 1 – CLA significance] We thank the reviewer for pointing this out. As noted in our general response, while some data points in Figure 3 are close to the y=x line, the confidence regions lying above this line indicate statistical significance for multiple models.

---

[Weakness 2.1 – How CID+ differs from CID] We apologize for not explaining in detail how CID and CID+ are different and their specific configurations. CID algorithm can be controlled by changing the values of two hyperparameters:

• $d$: the number of tokens with largest logits that will be considered in CLA. Selecting a larger $d$ will result in more cause-effect token pairs to be selected by CLA, and thus more tokens are subject logit changes in CID.

• $h$: the logit change applied to cause and effect tokens detected by CLA. A larger $h$ will alter the token distribution for word prediction more aggressively.

CID+ has a more aggressive configuration than the CID algorithm. Specifically, CID has $(d, h) = (2,5)$ and CID+ has $(d, h) = (5, 10)$. We have included the specific configurations in the revised manuscript.

---

[Weakness 2.2 – Mixed CID results] We thank the reviewer for pointing out that the CID results seemed mixed. We can observe in Table 2 that most cases where CID or CID+ fails to improve the original decoding are with Gemma-2-9b-it. Therefore, we specifically investigated Gemma-2-9b-it and made some interesting observations that can provide some insights to understand the results. We noticed that the texts generated by Gemma-2-9b-it were formatted well without specifically being prompted. We take one example from GSM8K to show this. Our prompt is

Given a question, please provide the final answer in the following format: "The answer is [a number here]."\nQuestion: Edgar eats 18 pretzels a day. If his brother eats 1/2 as many, how many does his brother eat in a week?\n Answer:

The answer generated by Gemma-2-9b-it is:

Here's how to solve the problem:

• Find the brother's daily pretzel intake: 18 pretzels / 2 = 9 pretzels

• Calculate the brother's weekly pretzel intake: 9 pretzels/day * 7 days/week = 63 pretzels

• The answer is 63.

We observed that Gemma-2-9b-it answered most math questions by starting with “Here’s how to solve the problem:” and following with bulleted steps, even though we did not prompt it to generate in this format or apply CoT. This can also be evidenced by the fact that CoT did not help Gemma-2-9b-it at all as shown in Table 2.

We conjecture that this phenomenon can be attributed to how the LLM is instruction-tuned. If the model has been aligned to generate text in specific format, changing the token distribution as we do in CID will make the generation deviate from the format and thus generate texts of lower quality.

We would like to emphasize that this is merely a conjecture based on the observations from existing experiments. We are actively investigating this phenomenon and look forward to reporting our findings in future work.