Paper: Learning Identifiable Causal Structures with Pairwise Representation (ICLR-2024)

Reviewer comment (Writing):
Some sentences of the paper should be simplified as those will help readers better understand the concepts.

Author response:
Thanks for your careful reading and suggestions. We have revised the paper according to your suggestions.

Reviewer comment (Writing):
The authors mentioned about “standard forward sampling P(D|G)”. This should be more precise.

Author response:
As detailed in Appendix Sec. A3, we apply both the linear Gaussian mechanism and a general nonlinear mechanism for continuous data generation. The Random Fourier Function mechanism is used for the general nonlinear data following the previous paper (Lorch et al., 2022). For discrete data, the Bernoulli distribution is employed. We have revised Appendix Sec. A3 to further clarify it.

Reviewer comment (Reproducibility):
It is not explained properly how the trained mode on SF graphs was used to discover the Sachs causal graph.

Author response:
As demonstrated in Sec. 5.2, we trained the PAIRE model on randomly generated scale-free (SF) graphs and then used it to predict the causal graph based on observational data from the Sachs dataset. The observational data, sourced from the Bnlearn repository, is readily compatible and can be directly processed by our trained PAIRE model.

Reviewer comment (Experiments):
The authors mentioned in the appendix that all their synthetic graphs contain 30 nodes and Sachs contains 11 nodes. This does not show if their algorithm is suitable for discovering larger graphs. The author should show their algorithm performance on larger graphs (with more than 100 nodes) which can be found in the Bnlearn repository (https://www.bnlearn.com/bnrepository/)

Author response:
As evidenced in Sec. 5, we have verified that PAIRE outperform other methods on the 30-node graphs, which is a typical setting of the graph size in the literature (Ke et al., 2023)(Lorch et al., 2022). Currently we do not have the capability for experiments on the larger graphs due to the limited computational resource, and the improvement of the scalability is left as interesting future work.

Reviewer comment (Novelty):
How is PAIRE different from their closest baseline AVICI in terms of methodology? It appears to me that they also used transformer-based architecture.

Author response:
As we summarize in Sec. 1, we design pairwise representations to capture essential causal information, and propose to use a skeleton matrix with a v-tensor as the learning target. Our ablation study shown in Sec. 5.3 - 5.4 verify the effectiveness of these innovations.

Reviewer comment (Theory):
How does the proposed algorithm adapt to a dataset with higher-dimensional variables? For example, if each sample contains d variables each having k dimensions, i.e., the dataset matrix will be D \in R^{n x d x k}.

Author response:
In the standard framework of causal discovery, each variable constitutes a unidimensional probability distribution. When dealing with high-dimensional variables, the graph representation and the corresponding causal structures are unclear. We leave it as interesting future work.

Reviewer comment (Writing):
What does the node feature refer to?

Author response:
As elaborated in the final paragraph of Sec. 4.2, the node features refer to the features of each individual variable. These are derived from the node feature encoder module within our model. We have revised Appendix Sec. 4.2 to further clarify it.

Reviewer comment (Theory):
What is the theoretical guarantee that this structure will always learn correct causal relations among variables given any datasets?

Author response:
- **Theoretical Guarantee.** As outlined in Sec. 3 and further elaborated upon in the revised Appendix Sec. A1, our model is based on three assumptions which are commonly used in the causal discovery domain: causal sufficiency, Markovian property, and faithfulness. These assumptions are widely recognized and employed in related work (Ma et al., 2022)(Dai et al., 2023). Furthermore, by assuming that the neural network functions as a universal approximator, we provide a theoretical foundation for the asymptotic correctness w.r.t. sample size and consistency of our method in Appendix Sec. A1. 
- **Discussion on Practicability.** Even though our method has a theoretical guarantee of the correctness with infinite samples, it has distinct advantages as a kind of Supervised Causal Learning (SCL) method under the finite samples conditions. For example, traditional constraint-based methods usually suffer from the unreliablity of conditional independence test in practice. Additionally, search-based methods tend to underperform because of the exponential growth of the search space, which is also mentioned in related work (Ke et al., 2023). However, SCL methods can learn from supervised training data about how to make predictions from practical finite samples, which makes the causal structure learning more robust and general. We have elaborated the contents above in the revised Appendix Sec. A1.4.

Reviewer comment (Theory):
The authors mentioned that they used “unidirectional attention operation”. What would be the alternative and what information that might represent for causal learning?

Author response:
The unidirectional attention operation is motivated by two key insights:
  - As highlighted in Sec. 1, it is crucial to incorporate both pairwise and node information to effectively capture essential causal information, such as persistent dependency and orientation asymmetry. To this end, unidirectional attention is employed to utilize information from both pairwise features and node features and represent the essential causal informtion.
  - In terms of computational complexity, unidirectional attention is less demanding than an alternative self-attention operation involving pairwise features alone. Consequently, this approach also serves to reduce the overall computational costs.

We have further clarified it in the revised Sec. 4.2.

Reviewer comment (Experiments):
We know deep learning approaches take a lot of data to train; on the other hand, we do not generally have a large amount of real-world data for a specific causal graph. How will the proposed method adapt to those situations?

Author response:
- As mentioned in Sec. 1, the training data for SCL methods can be readily obtained via synthetic generation. Different with some other tasks, we are able to generate infinite observations for a given causal graph with an approximate free cost. Therefore, it offers the significant advantage of the acquisition of training data, allowing for extensive and sufficient training of the models. The effectiveness of this paradigm has been verified in previous SCL methods (Dai et al., 2023; Ke et al., 2023; Ma et al., 2022).

Reviewer comment (Experiments):
Did the authors apply the mentioned heuristic and the Meek rules for baseline algorithms as well?

Author response:
As detailed in Appendix Sec. A3, we have employed the implementation provided by the gCastle repository for multiple baseline methods. In the case of AVICI, the heuristic and Meek rules are not applicable, as AVICI directly predicts the Directed Acyclic Graph (DAG).

Reviewer comment (Evaluation):
What are the reasons behind classical algorithms’ worse performance for the chosen datasets?

Author response:
We adhere to the default settings provided in the gCastle library, and detailedly check their predictions to ensure correctness. The relevant codes and datasets will be made available for release.

Reviewer comment (Experiments):
How did the authors apply the trained model on the Sachs graph? How was the Sachs observational data used? These details should also be made precise.

Author response:
As demonstrated in Sec. 5.2, we trained the PAIRE model on randomly generated scale-free (SF) graphs and then used it to predict the causal graph based on observational data from the Sachs dataset. The observational data, sourced from the Bnlearn repository, is readily compatible and can be directly processed by our trained PAIRE model.

Reviewer comment (Theory):
How are the authors making sure that there would be no cycles during skeleton learning? Their mentioned baseline AVICI seemed to optimize for acyclicity while it is not clear how the authors are ensuring that.

Author response:
Firstly, in the task of skeleton learning, the predicted skeleton is an undirected graph, thereby precluding the presence of cycles. We assume the reviewer questions how to handle cycles in CPDAG prediction. Theoretically, our updated analysis in Appendix Sec. A1 shows that PAIRE is theoretically proven to be asymptotically correct and consistency. This implies that cycles are never predicted under ideal cases. Practically, the occurrence of cycles is also rare in our predictions and PAIRE performs well in our experiments.

Reviewer comment (Experiments):
If a model is trained on synthetic data that follows a specific distribution, that trained model might not perform well on any new dataset with a different data distribution due to any bias toward previous training datasets. In those situations, the existing constraint-based causal discovery algorithms might perform better. How does the proposed method in this paper deal with this issue?

Author response:
The PAIRE model, trained on randomly generated datasets, is specifically designed to adeptly capture causal relations through our pairwise representations. As detailed in the Appendix Sec. A3, in our experimental evaluations, we primarily focus on out-of-distribution (OOD) settings, where the test distribution is distinct from the training distribution. As the experimental results presented in Sec. 5 indicate, PAIRE exhibits strong generality under this OOD setting. It is worth to note that the transferability of Supervised Causal Learning is also evidenced in previous work (Dai et al., 2023).

Reviewer comment (Experiments):
Can the authors show their performance with varying edge density and sample size?

Author response:
**Varying Sample Size.** We assessed PAIRE across various quantities of observational samples per graph during testing (100, 200, ..., 1000). Owing to space limitations, we showcase four key metrics in the table below: AUC and ACC for skeleton prediction, along with V-structure F1 and Identifiable Edges F1 for CPDAG prediction. The visualized figures for all metrics are provided in the revised Appendix Sec. A4. The model's performance enhances with the augmentation of sample size. These consistent upward trends suggest that PAIRE exhibits stability and is not overly sensitive to changes in sample size.


- WS Dataset

|       | 1000  | 900   | 800   | 700   | 600   | 500   | 400   | 300   | 200   | 100   |
|-------|-------|-------|-------|-------|-------|-------|-------|-------|-------|-------|
| AUC   | 0.742 | 0.743 | 0.738 | 0.738 | 0.735 | 0.733 | 0.729 | 0.723 | 0.715 | 0.690 |
| ACC   | 0.753 | 0.752 | 0.751 | 0.751 | 0.750 | 0.748 | 0.746 | 0.744 | 0.739 | 0.727 |
| V-structure F1   | 0.343 | 0.343 | 0.342 | 0.332 | 0.328 | 0.323 | 0.314 | 0.304 | 0.283 | 0.230 |
| Identifiable Edges F1   | 0.407 | 0.405 | 0.405 | 0.397 | 0.393 | 0.388 | 0.382 | 0.372 | 0.352 | 0.303 |

- SBM Dataset

|       | 1000  | 900   | 800   | 700   | 600   | 500   | 400   | 300   | 200   | 100   |
|-------|-------|-------|-------|-------|-------|-------|-------|-------|-------|-------|
| AUC   | 0.981 | 0.981 | 0.981 | 0.981 | 0.980 | 0.979 | 0.977 | 0.974 | 0.965 | 0.934 |
| ACC   | 0.962 | 0.963 | 0.962 | 0.961 | 0.961 | 0.958 | 0.957 | 0.953 | 0.946 | 0.920 |
|  V-structure F1  | 0.800 | 0.802 | 0.798 | 0.793 | 0.795 | 0.780 | 0.775 | 0.753 | 0.719 | 0.607 |
| Identifiable Edges F1   | 0.822 | 0.825 | 0.823 | 0.817 | 0.818 | 0.804 | 0.797 | 0.780 | 0.749 | 0.639 |


**Varying Edge Density.** We evaluated PAIRE over a range of edge densities in the test graphs, utilizing the SBM dataset, as it allows for the direct setting of average edge densities. Owing to space limitations, we showcase four key metrics in the table below: AUC and Acc. for skeleton prediction, along with V-structure F1 and Identifiable Edges F1 for CPDAG prediction. The visualized figures for all metrics are provided in the revised Appendix Sec. A4.  It's apparent that the task is becomes more difficult as edge densities increase. However, the performance decline is not abrupt, indicating that PAIRE's performance remains relatively stable across various edge densities, thereby confirming its versatility.


| Avg. Edges per Vertex | 1.0   | 1.5   | 2.0   | 2.5   | 3.0   |
|-------|-------|-------|-------|-------|
| AUC   | 0.998 | 0.993 | 0.978 | 0.958 | 0.933 |
| ACC   | 0.994 | 0.982 | 0.959 | 0.933 | 0.909 |
| V-structure F1  | 0.940 | 0.878 | 0.791 | 0.720 | 0.658 |
| Identifiable Edges F1 | 0.944 | 0.887 | 0.809 | 0.749 | 0.696 |

Reviewer comment (Experiments):
For the experiments, the author(s) discussed the number of parameters in Section 5.3, but not in Section 5.1. I wonder if the number of parameters across methods in Table 2 are comparable.

Author response:
As outlined in Sec. 5.3, our experiments with PAIRE utilize a 6-layer node feature encoder, amounting to 2.8 million trainable parameters. In comparison, the AVICI method comprises 3.2 million trainable parameters. It's important to note that this comparison of parameter count is not directly applicable to other baseline methods, as they do not employ SCL techniques. Regarding the FLOPs, PAIRE requires 38M FLOPs while AVICI requires 33M FLOPs in processing one batch of graphs. Although PAIRE's computational cost is marginally higher, it remains within a similar range.

Reviewer comment (Evaluation):
It would be nice to have a FLOPs comparison between methods, since the pairwise decoder is more computationally intensive than regular self-attention.

Author response:
As outlined in Sec. 5.3, our experiments with PAIRE utilize a 6-layer node feature encoder, amounting to 2.8 million trainable parameters. In comparison, the AVICI method comprises 3.2 million trainable parameters. It's important to note that this comparison of parameter count is not directly applicable to other baseline methods, as they do not employ SCL techniques. Regarding the FLOPs, PAIRE requires 38M FLOPs while AVICI requires 33M FLOPs in processing one batch of graphs. Although PAIRE's computational cost is marginally higher, it remains within a similar range.

Reviewer comment (Presentation):
The figure captions could use more explanation as they should be self-contained.

Author response:
Thanks for your suggestion. We have revised the figure captions and the expression in Sec. 4.1 accrodingly.

Reviewer comment (Writing):
The phrase ‘pairwise relationship’ is not clear, especially in Section 4.1 where it is unclear if the pairs are (D,G) or pair of nodes.

Author response:
Thanks for your suggestion. We have revised the figure captions and the expression in Sec. 4.1 accrodingly.

Reviewer comment (Evaluation):
What are the number of parameters for methods in Table 1 and Table 2? Are there FLOPs comparisons?

Author response:
As outlined in Sec. 5.3, our experiments with PAIRE utilize a 6-layer node feature encoder, amounting to 2.8 million trainable parameters. In comparison, the AVICI method comprises 3.2 million trainable parameters. It's important to note that this comparison of parameter count is not directly applicable to other baseline methods, as they do not employ SCL techniques. Regarding the FLOPs, PAIRE requires 38M FLOPs while AVICI requires 33M FLOPs in processing one batch of graphs. Although PAIRE's computational cost is marginally higher, it remains within a similar range.

Reviewer comment (Theory):
What is the intuition behind using pairwise feature only for the Query, and what is the unidirectional attention mentioned in Section 4.2?

Author response:
The unidirectional attention means only using the pairwise features as the Query, and the node features as both Key and Value. This approach is motivated by two key insights:
  - As highlighted in Sec. 1, it is crucial to incorporate both pairwise and node information to effectively capture essential causal details, such as persistent dependency and orientation asymmetry. To this end, unidirectional attention is employed to utilize information from both pairwise features and node features.
  - In terms of computational complexity, unidirectional attention is less demanding than a self-attention operation involving pairwise features alone. Consequently, this approach also serves to reduce the overall computational costs.

We have further clarified it in the revised Sec. 4.2.

Reviewer comment (Theory):
The method claimed to resolve identifiability inconsistency, but there is no such guarantee even with infinite observed data.

Author response:
- **Theoretical Guarantee.** As outlined in Sec. 3 and further elaborated upon in the revised Appendix Sec. A1, our model is based on three assumptions which are commonly used in the causal discovery domain: causal sufficiency, Markovian property, and faithfulness. These assumptions are widely recognized and employed in related work (Ma et al., 2022)(Dai et al., 2023). Furthermore, by assuming that the neural network functions as a universal approximator, we provide a theoretical foundation for the asymptotic correctness w.r.t. sample size and consistency of our method in Appendix Sec. A1. 
- **Discussion on Practicability.** Even though our method has a theoretical guarantee of the correctness with infinite samples, it has distinct advantages as a kind of Supervised Causal Learning (SCL) method under the finite samples conditions. For example, traditional constraint-based methods usually suffer from the unreliablity of conditional independence test in practice. Additionally, search-based methods tend to underperform because of the exponential growth of the search space, which is also mentioned in related work (Ke et al., 2023). However, SCL methods can learn from supervised training data about how to make predictions from practical finite samples, which makes the causal structure learning more robust and general. We have elaborated the contents above in the revised Appendix Sec. A1.4.

Reviewer comment (Experiments):
Proposed method is not scalable (see below) and is limited to small problems.

Author response:
> Reviewer t5Qz concerns the scalability of PAIRE on larger graph.

In the whole paper, we have revealed the significance and potential of pairwise representations, as they are crucial for capturing essential information such as persistent dependency and orientation asymmetry. As specifically evidenced in Sec. 5.3, we have proved that models incorporating pairwise representations outperform the original models on the 30-node problems, which is a typical setting of graph size in the literature (Ke et al., 2023)(Lorch et al., 2022). The improvement of the scalability on larger graphs is left as interesting future work.

Reviewer comment (Experiments):
Experiments are not convincing: lacking setup details, lacking completive benchmark method, and some better results in the literature were not reported.

Author response:
We present the comparison between PAIRE and the aforementioned methods on the WS dataset, following the default settings in the gCastle library. 

| Model       | Skeleton F1 | Skeleton AUC | V-struc. F1 | Identfiable edges F1 | SHD     |
| ----------- | ----------- | ------------ | --------- | ----------------- | ------ |
| GRANDAG     | 0.163      | 0.502       | 0.116    | 0.115            | 169.48 |
| NOTEARS-MLP | 0.256      | 0.513       | 0.128    | 0.126            | 192.59 |
| GOLEM       | 0.293      | 0.539       | 0.158    | 0.191            | 172.63 |
|  PAIRE   |  0.479 +- 0.015   | 0.750 +- 0.003             |   0.298 +- 0.076     | 0.370 +- 0.062                  | 116.797 +- 7.253       |

PAIRE consistently demonstrates superior performance in comparison with these methods. We have included the above results in the revised Appendix Sec. A4. Note that the RL-BIC method is time-consuming and fails to produce results within a reasonable timeframe.

Reviewer comment (Theory):
The proposed method is mostly algorithmic and lacks guarantees or further useful information of the inferred results.

Author response:
- **Theoretical Guarantee.** As outlined in Sec. 3 and further elaborated upon in the revised Appendix Sec. A1, our model is based on three assumptions which are commonly used in the causal discovery domain: causal sufficiency, Markovian property, and faithfulness. These assumptions are widely recognized and employed in related work (Ma et al., 2022)(Dai et al., 2023). Furthermore, by assuming that the neural network functions as a universal approximator, we provide a theoretical foundation for the asymptotic correctness w.r.t. sample size and consistency of our method in Appendix Sec. A1. 
- **Discussion on Practicability.** Even though our method has a theoretical guarantee of the correctness with infinite samples, it has distinct advantages as a kind of Supervised Causal Learning (SCL) method under the finite samples conditions. For example, traditional constraint-based methods usually suffer from the unreliablity of conditional independence test in practice. Additionally, search-based methods tend to underperform because of the exponential growth of the search space, which is also mentioned in related work (Ke et al., 2023). However, SCL methods can learn from supervised training data about how to make predictions from practical finite samples, which makes the causal structure learning more robust and general. We have elaborated the contents above in the revised Appendix Sec. A1.4.

Reviewer comment (Theory):
The skeleton and v-structure are all inferred by a trained model, and there is no consistency guarantee, because we do not know how good the trained model is. In fact, in the experiment, authors only use 100 simulated graphs with 30-nodes in the training data. By no means this model would work well for all causal discovery settings.

Author response:
> Reviewer t5Qz concerns 100 simulated graph is insufficient for PAIRE training. 

As noted in Sec. 4.1, both the training graph $G$ and the corresponding dataset matrix $D$ are generated in real-time. They are not reused during training. Consequently, this approach potentially provides an infinite amount of training data for the model. A typical PAIRE model is trained for 300 epochs, with each epoch comprising 6,000 training graphs. This results in a cumulative total of 1,800,000 training graphs throughout the training process. As we mentioned in Appendix Sec. A3, each synthetic test datasets contain 100 graphs. We have further clarified this in the revised Appendix Sec. A3.

Reviewer comment (Experiments):
The pairwise representation used in this method may be beneficial, but it takes \(d^2h\)-dim where \(d\) is the number of nodes. If we have a 100-node problem and each original feature has \(h=4\) dim, this would be 40,000-dim feature. This design would limit the scalability of the problem. Clarification is needed.

Author response:
> Reviewer t5Qz concerns the scalability of PAIRE on larger graph.

In the whole paper, we have revealed the significance and potential of pairwise representations, as they are crucial for capturing essential information such as persistent dependency and orientation asymmetry. As specifically evidenced in Sec. 5.3, we have proved that models incorporating pairwise representations outperform the original models on the 30-node problems, which is a typical setting of graph size in the literature (Ke et al., 2023)(Lorch et al., 2022). The improvement of the scalability on larger graphs is left as interesting future work.

Reviewer comment (Experiments):
Important results were omitted. Existing works have better estimation results on Sachs data. For example, CSIvA has reported SHDs 5 and 6 on this dataset, while the proposed method has SHD=10. Although a method does not need to beat all others in all settings, reporting such results can be very informative for readers.

Author response:
> Reviewer t5Qz wonders why CSIvA reports a seemingly better result than PAIRE on Sachs.

As noted in Appendix Sec. A3, additional data of single interventions on each variable is available under the CSIvA settings hence CSIvA predicts the DAG. Our PAIRE is designed to learn from observational data and predicts CPDAG. Therefore, a direct comparison is not meaningful. Moreover, as noted in Appendix Sec. A3, since the code for CSIvA has not been made publicly available, we are unable to modify it for a comparative analysis with PAIRE.

Reviewer comment (Novelty):
Since the paper considers a transformer-based NN model with sampling methods to generate a causal graph, previous works earlier than CSIvA that also use transformer models and/or Bernoulli sampling should be discussed, such as:
- Causal discovery with reinforcement learning, ICLR 2020
- Neural Causal Structure Discovery from Interventions

Author response:
> Reviewer t5Qz suggests us to discuss two related papers [1, 2].
> - [1] causal discovery with reinforcement learning, ICLR 2020
> - [2] Neural Causal Structure Discovery from Interventions

Our work is more related to CSIvA compared with the two mentioned papers. Paper [1] is a score-based method rather than a supervised causal learning method. It suggests to use Reinforcement Learning (RL) in the search process, employing a transformer network in the encoder and outputs Bernoulli distributions for the adjacency matrix. Unlike CSIvA (Ke et al., 2023), AVICI (Lorch et al., 2022), and our approach, their transformer network lacks sample-wise and node-wise attention mechanisms. We have add a relevant discussion in Sec. 2 in our revised manuscript. Paper [2] introduces a continuous optimization method named Structure Discovery from Interventions (SDI). Its setup is different with PAIRE as it integrates both observational and interventional data.

Reviewer comment (Novelty):
Overall, while the paper proposes an improved method over CSIvA, the technical contribution is weak, there is no guarantee or other useful information of the inferred graph, and the experiments are not convincing. The reviewer recommends REJ for the current version.

Author response:
- **Theoretical Guarantee.** As outlined in Sec. 3 and further elaborated upon in the revised Appendix Sec. A1, our model is based on three assumptions which are commonly used in the causal discovery domain: causal sufficiency, Markovian property, and faithfulness. These assumptions are widely recognized and employed in related work (Ma et al., 2022)(Dai et al., 2023). Furthermore, by assuming that the neural network functions as a universal approximator, we provide a theoretical foundation for the asymptotic correctness w.r.t. sample size and consistency of our method in Appendix Sec. A1. 
- **Discussion on Practicability.** Even though our method has a theoretical guarantee of the correctness with infinite samples, it has distinct advantages as a kind of Supervised Causal Learning (SCL) method under the finite samples conditions. For example, traditional constraint-based methods usually suffer from the unreliablity of conditional independence test in practice. Additionally, search-based methods tend to underperform because of the exponential growth of the search space, which is also mentioned in related work (Ke et al., 2023). However, SCL methods can learn from supervised training data about how to make predictions from practical finite samples, which makes the causal structure learning more robust and general. We have elaborated the contents above in the revised Appendix Sec. A1.4.

We present the comparison between PAIRE and the aforementioned methods on the WS dataset, following the default settings in the gCastle library. 

| Model       | Skeleton F1 | Skeleton AUC | V-struc. F1 | Identfiable edges F1 | SHD     |
| ----------- | ----------- | ------------ | --------- | ----------------- | ------ |
| GRANDAG     | 0.163      | 0.502       | 0.116    | 0.115            | 169.48 |
| NOTEARS-MLP | 0.256      | 0.513       | 0.128    | 0.126            | 192.59 |
| GOLEM       | 0.293      | 0.539       | 0.158    | 0.191            | 172.63 |
|  PAIRE   |  0.479 +- 0.015   | 0.750 +- 0.003             |   0.298 +- 0.076     | 0.370 +- 0.062                  | 116.797 +- 7.253       |

PAIRE consistently demonstrates superior performance in comparison with these methods. We have included the above results in the revised Appendix Sec. A4. Note that the RL-BIC method is time-consuming and fails to produce results within a reasonable timeframe.