We sincerely thank you for your thorough review and for the clear and detailed suggestions. We believe these suggestions further highlight the significance of the causality-based approach presented in the paper compared to correlation-based approaches.

Answer to questions

Re Question 1 and 2

Firstly, under the assumption made in the paper that hidden confounders may be present, a correlation-based method may not be applicable. The method we presented in the paper employs a causal discovery algorithm to extract a graph, accounting for the possible presence of latent confounders, from the attention matrix.

As described below and empirically demonstrated, accuracy in generating legal tokens is more strongly associated with the causality-based approach presented in the paper than to the correlation-based methods, as measured by the confidence score (Section 4.3).

You suggested 2 correlation-based methods to derive an adjacency matrix. The first is thresholding the attention matrix, and the second is thresholding the precision matrix. The first method is related to testing marginal independence (empty conditioning set), which was examined in Section 5.2.1, Figure 3a. In the second method the precision matrix represents pair-wise independence conditioned on all other nodes (only the largest possible conditioning set size, $n-2$). In relation to these methods, constraint-based causal discovery tests independence conditioned on a range of conditioning set sizes from 0 to the maximal possible, as needed. In Section 5.2.1 Figure 3 we demonstrate that statistical significance results obtained using empty conditioning sets are different than those obtained using size 1 (Figure 3a-3c). The aggregation of results from CI tests required by causal discovery improves the number of statistically significant cases (Figure 3d) differentiating legal from illegal move generation. Furthermore, given a graph $\mathbf{G}$, nodes' values are $\boldsymbol{X}=\boldsymbol{X}\mathbf{G}+\boldsymbol{U} \implies \boldsymbol{X}=(\mathbf{I}-\mathbf{G})^{-1}\boldsymbol{U}$ (Equation 4). Treating attention as $\mathbf{G}$, instead of $(\mathbf{I}-\mathbf{G})^{-1}$ imposes restrictions on the graph.

Empirical Evaluation

For each game length we calculated confidence (Section 4.3). As in Figure 5, we binned confidence values and calculated mean accuracy. Results for Chess are in the tables below. It is evident that, except for a few cases, for both correlation-based methods there is no clear trend of accuracy with respect to confidence, as clearly evident for causality-based approach (Figure 5).

Attention thresholding

Precision-matrix thresholding

Re Question 3

A graph constructed using the last layer attention represents a graph over the output tokens. The last layer's attention output represents tokens (pre-training loss compares output tokens to input tokens), up to an MLP non-linear mapping, where MLP has no effect on inter-token relations and therefore does not influence causal structure learning. Earlier layers represent context (such as exogenous nodes in the SCM) for their following layer. See also the explanation following Equation 8. For example, in Othello intermediate layers may represent the board (as exhaustively tested and found by Li et al., 2023) which serve as context for the causal graph over game moves represented by the last layer. Furthermore, recent studies empirically show that the last layers have more weight in determining the next token, specifically in Othello GPT [1] and in LLMs in general [2,3,4].

[1] How GPT learns layer by layer. Du et al., 2025. [2] Adaptive Large Language Models by Layerwise Attention Shortcuts. Verma et al., 2024. [3] The Mechanics of Conceptual Interpretation in GPT Models: Interpretative Insights. Aljaafari et al., 2024. [4] Exploring Concept Depth: How Large Language Models Acquire Knowledge at Different Layers? Jin et al., 2024.

Re Other comments or suggestions:

• In Eq 9 both RHS and LHS are matrices. A hat over index $i$ represents the omission of the $i$-th row/column, as implicitly mentioned in the sentence before the equation. We will explicitly describe this notation.
• We will clearly define the computation of the entropy values $H_\textrm{ind}$ and $H_\textrm{dep}$.

We sincerely thank you for your review and your perspective on the paper. We value your feedback and believe the following answers your concerns.

• Re first point. We would like to clarify that autoregressive generation is not inherently causal. Often the attention in GPT is called 'causal' but it only means that a token is a function of previous tokens in the sequence. That is, the upper triangular part of the attention matrix is masked (zero). However, this is an oversimplification of the broader term "causal". In the paper we show that an attention value $\mathbf{A}_{i,j}$ describes the sum of all directed paths from node $j$ to node $i$, that is the effect node $j$ has on node $i$ accounting for all directed paths. Another difference between the attention and causal graphs is the set of encoded conditional independence relations. For example, consider the causal relations $X$ causes $Z$ and $Y$ causes $Z$, and no other relations: a graph: $X$ --> $Z$ <-- $Y$. The graph entails that $X$ and $Y$ are marginally independent, but dependent conditioned on $Z$. However, note that several causal graphs may entail the exact same independence relations (e.g., a hidden confounders may be present between $X$ & $Z$ and $Y$ & $Z$ instead of causal relations). The attention matrix only represents the correlation between $X,Y$ and $Z$ (the level of 'attention' $Z$ gives to $X$ and $Y$) and does not represent causality. The causal interpretation in this paper demonstrates that a causal graph, including hidden confounders, could be extracted, where the attention matrix represents the total effect a node has on another through all directed causal paths. That is, the $(i,j)$ element in the attention matrix $\mathbf{A}$ represents the sum of all the directed paths from node $j$ to node $i$ in the causal graph. We also introduce a novel confidence score for the learned graph using entropy of p-values of statistical tests used during the causal structure learning.
• Re second point. We will clarify in the paper that a 'causal world model' describes the causal mechanism underlying the observations as well as probability distribution of hidden variables, as often used in the causal inference literature. Specifically, we will note that it is assumed that underlying each input sequence there exists a corresponding structural causal model (SCM, Section 3.2). That is, the causal world model consists of the causal mechanism that generates tokens.
• Re third point. We discuss implications of the findings in Section 6 (Conclusions) and in Section 'Impact statement" (after Section 6. Conclusions; part of ICML 2025 template which does not count towards page limit). The implications include a) the ability for zero-shot causal discovery which can be beneficial in various scientific domains (e.g., understanding the mechanism by which a GPT trained on the protein space generates novel protein sequences [1]), b) measuring uncertainty of attention heads per input sequence, c) calibration during training (accuracy vs. causal confidence), and d) better human supervision by examining the causal mechanism by which a token is generated.

[1] Ferruz, N., Schmidt, S., & Höcker, B. (2022). ProtGPT2 is a deep unsupervised language model for protein design. Nature communications, 13(1), 4348.

Thank you for your detailed feedback, insights, and important questions. Addressing your review improves the over clarity of the paper and emphasizes the significance of the contribution.

Re Questions for Authors

1. The relation $\mathbf{D}^{-1}\mathbf{A} = (\mathbf{I}-\mathbf{G})^{-1}$ is not an assumption, but rather a result by considering each token as a node in an SCM (the sentence just before equation 8). Since attention calculates $\mathbf{Z}=\mathbf{A}\mathbf{V}$ (Equation 1) and the SCM calculates $\boldsymbol{X}=(\mathbf{I}-\mathbf{G})^{-1}\boldsymbol{U}$ (Equation 4), we equate the outputs: token embedding and SCM node value. This is described two sentences before Equation 9, but we will clarify this point more explicitly in the paper.

2. A ground truth of causal graphs is not available for most domains. In the strategy games of Othello and Chess, since there are multiple next legal moves, a causal graph over a given set of game moves may contain information about the strategy of the player (in case of a real game) in addition to the game rules. Moreover, it is unclear how to modify the game rules such that a sequence of game moves is legal for both the original rules and the modified rules, while making sure the set of next legal moves entailed by the modified rules does not overlap with that entailed by the original moves. That is, for a fair comparison the input sequence should be legal for both the original and intervened rules, and for correctly testing, the next legal move should be different. Nevertheless, at the request of Reviewer g9Sr, we evaluated the accuracy of legal move generation as the function of confidence score (Section 4.3) calculated for correlation-based methods they suggested. See also Figure 3a vs. Figure 3d in the ablation study (Section 5.2). We believe your concern is additionally addressed in our reply to their Questions 1 and 2. We did not find an increase in legal move generation accuracy as a function of the confidence score for correlation-based methods (as found for the learned causal graph).

3. For a given sequence of tokens and a legal next token generation, there may be multiple, equally minimizing the loss, attention matrices. During pre-training, these attention matrices were not restricted to have high confidence for determining high-order conditional independence (CI) relations between tokens (in comparison to marginal pair-wise correlations). We examined confidence of attention matrices generating legal vs illegal token generation. In all our experiments we found a correlation between the causal confidence score, which was calculated during causal structure learning, and the accuracy of generating legal moves. Confidence decreases when uncertainty increases in CI tests decisions used for constructing the causal graph structure (e.g., an edge is removed if conditional independence is found, and edge orientation is determined based on the conditioning set that disconnected two nodes). We will add a corresponding clarification this in the paper.

4. We provided a qualitative examples of causal graphs learned for the first few moves in Othello and Chess games in Figure 1 and Figure 6, respectively. These graphs are actual results. After these few moves it becomes challenging to follow the true causal relations. For example, during the Othello game, previously placed pieces may flip their color multiple times affecting the set of legal moves. Next, we regard the possibility of having high $R(\mathbf{A})$ due to the last few moves effecting the next move. In our Ablation study in Figure 3 we compare (a) marginal (in)dependence relations to (d) conditional independence relations that are required to construct causal graph. It is demonstrated that the difference in $R(\mathbf{A})$ between legal and illegal token generation is more prominent when considering independence relations required for constructing a causal graph. As mentioned in our reply to your Question 2, we also examined the experiment in Figure 5 for correlation-based methods.

5. The main difference between the causal interpretation and other interpretation is the incorporation of conditional independence relations having various conditioning set sizes rather that treating the attention values as weights. As mentioned in our answer in Question 1, by relating output embeddings computed using the attention matrix to the node values in an SCM we obtain the relation in Equation 9. By equating the covariance matrices (Equation 8) we can employ a constrain-based causal discovery algorithm. In our experiments we demonstrated that a confidence score computed from conditional independence relations provides a better differentiation between legal and illegal token generation compared to marginal independence relations.

Finally, we found the papers you suggested to include the the introduction to strengthen the significance of this paper's results. As suggested, we will add a related discussion.

Thank you for the detailed review and suggestions for improvement. Your suggestions will improve the clarity and emphasize the significance of the proposed approach and findings presented in the paper. We also thank you for the many suggestions for future work.

Re Other Comments or Suggestions

Upon your interesting suggestions, in our future work we will extend the analysis for OOD data in different domains. Specifically, we plan to explore improvements for pre-training relying on the findings in the paper. For example, we plan to explore the effect of OOD generalization after calibrating token-generation accuracy with respect to causal confidence (Equation 10). Finally, we will improve the clarity of figures' captions as requested and correct grammatical errors and typos.

Re Questions for Authors

• The full set of CI tests having all possible conditioning set sizes has a one-to-one correspondence with a causal graph (up to Markov equivalence). A causal-discovery algorithm uses a sub-set of this exhaustive set of CI tests. A sound and complete algorithm, like the one we use in the paper, creates a one-to-one correspondence between this subset and a Markov equivalence class. We will add this clarification along with a citation (Spirtes et al., 2009).

• Yes. Positional encoding specifically for the board was not used (in accordance with the work of Li et al., 2023). Positional encoding was (also) not used by Li et al., 2023 who trained a GPT without domain information, such as the presence of a board and its 2D alignment on which the game moves take place. We will clarify this.