Graph Structure Inference with BAM: Neural Dependency Processing via Bilinear Attention
Authors: Philipp Froehlich
Heinz Koeppl
Venue: NeurIPS 2024 poster (2024)
Keywords: Graph Structure Inference
Causal Inference
Supervised Deep Learning
Geometric Deep Learning
Primary area: causal_inference
TL;DR: We present a neural network with a novel bilinear attention mechanism for supervised graph structure inference, mapping observational data to their dependency structures with enhanced performance and robust generalizability.

Abstract: Detecting dependencies among variables is a fundamental task across scientific disciplines. We propose a novel neural network model for graph structure inference, which aims to learn a mapping from observational data to the corresponding underlying dependence structures. The model is trained with variably shaped and coupled simulated input data and requires only a single forward pass through the trained network for inference. Central to our approach is a novel bilinear attention mechanism (BAM) operating on covariance matrices of transformed data while respecting the geometry of the manifold of symmetric positive definite (SPD) matrices. Inspired by graphical lasso methods, our model optimizes over continuous graph representations in the SPD space, where inverse covariance matrices encode conditional independence relations. Empirical evaluations demonstrate the robustness of our method in detecting diverse dependencies, excelling in undirected graph estimation and showing competitive performance in completed partially directed acyclic graph estimation via a novel two-step approach. The trained model effectively detects causal relationships and generalizes well across different functional forms of nonlinear dependencies.

Graph Structure Inference with BAM: Neural Dependency Processing via Bilinear Attention
Philipp Froehlich
Heinz Koeppl
Department of Electrical Engineering and Information Technology
Technische Universität Darmstadt heinz. koeppl -arasstdd.t de
Abstract
Detecting dependencies among variables is fundamental task across scientific disciplines. We propose novel neural network model for graph structure infer- ence, which aims to learn a mapping from observational data to the corresponding underlying dependence structures. The model is trained with variably shaped and coupled simulated input data and requires only single forward pass through the trained network for inference. Central our approach is novel bilinear attention mechanism (BAM) operating on covariance matrices of transformed data while respecting the geometry of the manifold of symmetric positive definite (SPD) ma- trices. Inspired by graphical lasso methods, our model optimizes over continuous graph representations the SPD space, where inverse covariance matrices encode conditional independence relations. Empirical evaluations demonstrate the robust- ness our method in detecting diverse dependencies, excelling undirected graph estimation and showing competitive performance in completed partially directed acyclic graph estimation via a novel two-step approach. The trained model effee- tively detects causal relationships and generalizes well across different functional forms of nonlinear dependencies.
1 Introduction
The discovery and understanding of dependencies among variables are fundamental to scientific in- quiry across diverse disciplines, such biology [10), climate science [48], economics (8), and social studics [18]. These dependencics are commonly represented as edges within graphical models, par- ticularly directed acyelic graphs (DAGs), first introduced by Wright (70) study genetic inheritance, and later advanced for probabilistic inference by Lauritzen and Spiegelhalter [35]. Graph structure inference, the process of deriving such graphical representations from observational data, is essential for gaining insights into complex systems and their underlying causal relationships [49]. A prime example is the estimation of gene regulatory networks from experimental data [19), where identifying dependencies between genes is crucial for understanding biological processes and mechanisms.
Graph structure inference typically employs unsupervised learning methods to estimate the underlying graph through either score-based approaches, which rank graphs by predefined metrics, or constraint- based approaches that identify edges between variable pairs via conditional independence tests [66]. However, these methods face challenges. Score-based approaches encounter computational burdens due to the superexponential growth of potential graph structures with node count, and the necessity to balance fit and structural sparsity B1 Constraint- based methods usually require a large sample size
(66), rely on elusive optimal threshold hyperparameter, and Shah and Peters (58) proved that the
failure of Type error control in underlying conditional independence tests is unavoidable, which can have significant consequences in downstream analyses.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Supervised causal learning techniques, as presented by Lopez-Paz et al. [39 40, 41]. Li et al. [36, Ke al. [3T, Lorch et al. [42]. Dai al. [15], have recently emerged as an appealing alternative to unsupervised methods. In this rising approach, a neural network is typically trained on simulated matrix-shaped data, with corresponding graph structures serving as ground-truth labels for supervised learning. The paradigm capitalizes on the strengths of deep learning to discern complex patterns in data, thereby enabling accurate graph structure inference.
Supervised causal learning strongly relies on neural networks' ability to extract dependency infor- mation from observational data matrices. Current approaches for processing observational data, whether in causal learning [3] 42), or related fields [4). typically embed the data into expanded observational data spaces, attempting to implicitly capture dependencies within the Euclidean space in transformed data matrices. After marginalizing over the sample dimension, existing supervised
42 then predict edges separately from each other via compressed observational representation vectors, rather than estimating single graph object as in graphical lasso (GLASSO) algorithms 72 17]. GLASSO methods effectively optimize inverse covariance matrices in SPD space, where matrix entries directly encode the graph's conditional dependencies. This classical, highly influential approach is particularly powerful the Gaussian case, where the covariance matrix is sufficient statistic for the distribution. We extend this geometric insight from classical optimization to the neural network era by introducing SPD layers that learn and process dependency information directly. Instead of using pre- pre-derived mathematical operations as in GLASSO-bbased algorithms, our network learns these shape- inape-invariant transformations from data. Combined with learned transformations in the observational space, approach generalizes beyond the Gaussian assumption of classical methods while preserving both the theoretical guarantees and computational efficiency of GLASSO approaches. In contrast to GLASSO methods which must solve an optimization problem for each new dataset, our approach requires only a single forward pass.
causal learning methods [31
Analogous to GLASSO-based methods, the symmetry of covariance matrices constrains our approach to undirected graph estimation. We extend beyond this limitation via novel two-phase procedure. We first estimate the graph's skeleton and immoralities, then identify edge directions obtain completed partially directed acyclic graph (CPDAG) [12]. While this follows the general principle of the PC-algorithm [61], our deep learning approach distinguishes between skeleton edges and moralized edges through joint optimization in SPD space, instcad of relying on sequential conditional independence tests all possible immoralities for edge identification. By efficiently identifying immoralities in the first step, our method drastically reduces the number of required independence tests, minimizing both computational complexity and the propagation of false positive orientation errors.
channel
Riemannian manifold
embedding
dimar
ravarianee
LogEig
M
billnead
clemesstwiee activation
astention
Probablities classes: -mo edge -skeleton odge -moralized edge
observational
attention
atteurion
data
actributes
datapoista
Figure Neural network architecture: An input of arbitrarily shape (M, d) provided, which is then embedded into C channels. Attention between attributes and attention between datapoints are applied alternately. Covariance matrices are calculated, followed by alternating applications of bilinear attention and the custom activation function in the Riemannian manifold of SPD matrices. The matrices are then transformed back into Euclidean space using the Log-Eig layer. Output probabilities for each pair of variables being in the classes "no edge", "skeleton edge". and "moralized edge" are calculated using dense layers and applying a softmax layer on the channel axis.

---

2 Method
2.1 Supervised approach for learning graph structures
In supervised graph learning, simulated training data is generated in the form of input/label pairs (x, (X,9). Each pair consists of an observational data matrix X E R Mxd where denotes the number of samples and d represents the number of attributes, and its corresponding graph structure g (V, E). The graph G is defined by a set of nodes V = {ui...., va} and a set of edges E _ x V. Alternatively, graphs can be represented as adjacency matrices A E {0, 1 Iddd, where Aij 1 if edge exists between nodes e; and fj, and Aij 0 otherwise. The process begins by creating a random graph, followed by generating obserational data matrices using a structural equation model (SEM) X, fo (w)) Cp),
where fo is measurable coupling function of parent nodes pa(g) and zero-mean error ey. The random coupling in the SEM is designed to approximate broad class of functions, enabling the model to perform classification of the presence or absence of causal relationships, irrespective of their specific functional forms. neural network then learns the relationship between the graph and the generated data matrix as a mapping X - G, enabling the inference of graph structures from observational data. The supervised causal learning procedure shown in Figure
Generate random graph and coupling
SEM
X, fo (X)
Yo e v
Ground truth graphs
UL.... fa
D
g
On-the-lly training
Data matrices with random coupling and random number of samples and attributes
novel neural network architecture
Figure 2: Overview of our supervised graph learning approach. We generate synthetic training data by first creating random ground truth graphs. Then, observational data matrices are generated using SEMs with randomly parameterized polynomials. Employing our novel neural network architecture, we learn the mapping X + from observational data matrices the corresponding graph structures.
2.2 Generation of training data
For training data generation, we utilize Erdös-Rényi (ER) graphs, denoted as ERid, q), where the number of nodes d and the expected degree q are sampled from discrete uniform distributions d ~ U((d...,, d}}) and q U U[g,, 7), respectively. The shape-invariant design of our model enables it to train without restriction to fixed (d, q) pair, which proves advantageous when the graph density unknown. For each graph G;, we generate a data matrix X; E RMxdd using an SEM, where the sample size M drawn from a discrete uniform distribution, M ~ U(( M }). Although ER graphs used for training, demonstrate that our approach maintains its performance when tested on graphs generated using different random graph models, as detailed in Appendix G.6
As coupling functions in the SEM, we employ random multivariate Chebyshev polynomial functions to generate diverse continuous training data. Chebyshev polynomials effectively approximate real- world functions, showing factorial decay in their coefficients given that higher order derivatives are bounded (63), i.e., for the nith coefficient en holds lc,, S C for a constant C, making higher- degree terms negligible, such that all such smooth functions can be efficiently approximated with
n!

---

low polynomial degree Chebyshev polynomial. This choice of parameterization avoids unnatural, non-smooth coupling functions where a higher-order derivative is much larger than lower-order one, and where there is hard cutoff in the approximated order of the derivative. Chebyshev polynomials have been widely used in various applications due to their excellent approximation properties and computational efficiency [44] [3] Scatterplots of input/output values of four example Chebyshev polynomials are shown on the left side Figure? 2 The error terms e, in the SEM are modeled using Gaussian mixture models, which provide a flexible and expressive framework for capturing diverse range of error distributions
[55 Details about the parameterization of the SEM can be found in
Appendix E
2.3 Prediction task: three-class edge classification
Our neural network architecture designed process SPD matrices, which inherently contain exclusively symmetrical relationship information, thus excluding non-symmetrical information necessary for estimating edge directionality. While this configuration prevents the estimation of DAG or CPDAG in single run, it enables efficient estimation of symmetrical structures such as the graph skeleton. Additionally, it facilitates the identification of moralized edges. We differentiate between the following edge types:
Skeleton edges:
Undirected edges found in the underlying DAG.
Moralized edges: Not present in the DAG but emerge due conditional dependencies among nodes sharing a common child without a connecting edge between the parents.
No edge:
Conditionally independent variables given all other variables.
Skeleton and moralized edges are both represented as undirected edges in models that use partial correlations conditional independence tests, which the conditioning set includes all other variables, to estimate the moral graph, e.g.. in [17 59].
Using the assumption of faithfulness and the Markov condition enable us identify the Markov equivalence class [61] and thus uniquely solve the three-class classification problem. Appendix C provides explicit proof outlining the conditional independence relations that lead unique solution.
The prediction targets one-hot encoded adjacency matrices, denoted À E {0, 1] Here, for each i, the vector A.j,- E {0, 1}3 represents a one-hot encoded classification among the three classes: skeleton edge, moralized edge, and no edge.
extension of the binary adjacency matrix A {0, 1]dxd to set of
2.4 Neural network architecture
We use the multi-dimensional analogue of matrix multiplication: For a tensor [x and matrix B E RR X we denote
AB C C
with Capa = Σ Aija Bai
(1)
&
The network architecture is depicted in Figure In the following, we introduce the individual layers.
Channel embedding. We perform an embedding of the input X E representation with C' channels. For this, one axis for X is extended to X E IR Mxdx1 and then trainable weights W, 111 Wa RCC are used obtain
M
M xd to obtain a hidden
H + relu (Wi,) Wa E MxdxC
where broadcasting is used for the addition.
shape- and permutation-innvarinn invariant neural network is achieved by employing C x Ō weight matrices in conjunction with attention mechanisms for fixed channel dimensions C and of the current and the next layer. This enables information flow among elements both within each matrix and across the C' matrices, analogous to sequence models like the Transformer (65), where sequence of length I is embedded into an / x C matrix and processed via self-attention with C x C trainable weights.

---

Observational data self-attention. Self-attention for C' observational M x d data matrices, stored in tensor E MMxxx proposed by Kossen et al. [34 is based on axial attention [26] 67 and commonly used for processing observational data 4 (Figure 3 left) computes M attention matrices of shape d d. allowing attributes to attend to each other in parallel for each sample. The attention-betweent between-samples layer interchanges M and d, resulting in d attention matrices of shape M x M, enabling samples to attend to each other. These layers capture information flow along both axes of the observational data matrices. A mathematical description provided in Appendix B.1
31 42 'he
layer
A OX
H
H
6..4
K
K d. 09
Q
(M d.
TM
OM d. G
(6
W K
W Q
Wx
W Q
(C G
Σ
Figure 3: Left: Observational data self-attention layer across attributes. Gray denotes non-trainable tensors, and red represents trainable weights. Matrix multiplication is performed after necessary transposition to match axis dimensions. Right: Bilinear self-attention layer. The double arrow signifies the use of the matrix as bilinear operator.
Transition to SPD space via covariance calculation. Let a transformed data matrix X E R be given. To obtain explicit dependency information, calculate channelwise covariance matrices, i.e.:
== 1
M (I- 14x) T(3.2.1) ② (X- -xx
(4x))(1111 713.1.2) T(2,3,1) E ERxxxx
where we define for tensors A E RIXXXK B E R(XXXL the 1-parallel matrix multiplication A B = (Ai.,, Bi,,)s-i.... E RIXxXI and μx = M x E gxxxX is a tensor of sample means. We denote by saxd the cone of d symmetric positive semi-definite (SPD) matrices, and by SdxdxC = x... saxd the C-ary Cartesian power of saxd. It holds E S cdxdxC
The transpose notation T(ij.k) permutes the dimensions of a tensor according to the specified order (i, k), following the commonly used notation in deep learning [I] For example, A (223.1) rearranges the dimensions of A such that the second dimension becomes the first, the third becomes the second, and the first becomes the third.
Introducing the bilinear attention mechanism. To effectively process dependency information stored in the covariance matrices, which is crucial for causal discovery, we propose novel bilinear attention mechanism. Inspired by the success of SPD networks in image processing, our bilinear attention layer is designed to be shape-invariant and permutation-invarian
In contrast to existing SPD architectures (28, 68) that typically parameterize weights W E Rdxda to be applied as W'EW E dd to a matrix Σ E sdxé we parameterize weights W E RCxC to act as linear combinations EW E & SaxdxC on a set of covariance matrices Σ E SaxdxC We employ an attention
mechanism to generate attention matrices A E _dxxxxC, serving as adaptable d x d weights on each x d matrix of Σ E [dxxxCC Our bilinear attention mechansim is shown in (Figure right).
The Riemannian manifold of SPD matrices has nonpositive sectional curvature, and its geometric properties are not preserved under Euclidean operations [9), leading issues when using traditional neural networks [46]. Specialized architectures respecting the SPD manifold's geometry are necessary. Bilinea matrix multiplication Σ - W WWW, analogous to a dense layer in Euclidean space, preserves the space SPD matrices saxd and serves as a primary tool for SPD layers [68]
'The mapping A + AEAT is quadratic and hence nonlinear in A. It is often referred to as bilinear as it is special case of the mapping (A, BB AEB".
5

---

For an input Σ E S SáxdxC we obtain keys K = EWK E Stxdx and queries Q = EW q S SaxdxC; usiin non-negativity constraints on the weights WK, W E RC
Keys and queries are combined in bilinear fashion, parallel over the C channels by calculating K o (K ,Kc Q.,.e K (sse)) c Sdxdxc
resulting in tensor of SPD matrices.
We replace the softmax function for traditional attention with custom softmax function ō to preserve positive definiteness, i.e.,
sdxd σ : : + S o(S)= v A(()) / AA() where A(S) :- diag
(3)
expls
where [x denotes the elementwise application of the exponential function, 1 is a vector of length with all entries being 1, and diag transforms a vector of length d into a x d diagonal matrix. The quotient and root are also taken elementwise.
Using the definition the bilinear parallel tensor product in 2 we obtain the attention matrix A by applying ā channelwise, i.e.,
A = (a((K & Q).,..)
c=l....CC
and we obtain the output of the bilinear layer
H A E S sdxdxC
The elementwise exponential function preserves positive definiteness because SPD matrices are closed under addition and the Hadamard product [9]. and exp[S] >n= [S]", with elementwise exponentiation I.|". Theorems regarding the eigenvalue regularization and stability properties of our modified softmax function can be found in Appendix A
Log-Eig layer and output softmax. Data representation transitions from the SPD space to Eu- clidean space through the Log-Eig layer, as proposed by (28) This transformation, given an input matrix S = UDU via eigendecomposition can be expressed as
I: S ddddd 1(S) == log(S) == Ulog(D)U
The final layer, equipped with softmax activation function, consists of C' = 3 output units for generating the probabilities.
Interpretation. In standard attention, the (i, j)tt entry of the attention matrix A directly represents scalar influence of element i element In our bilinear attention framework, however, interdepen- dencies emerge. Specifically, for an output pair (i,j), the corresponding output value is given by the bilinear form Σx, Ai, Sk, At,j- Hence, the influence on the (i,j))-t J-th entry depends not only A but the entire i-th row and j-th column of A. This produces cross-shaped attention pattern within A and allows columns of S to attend to each other, capturing more complex interactions. The learned SPD representations can be interpreted end-to-ent learned kernel matrices encoding different kinds of dependencies across multiple channels. A detailed discussion of the intuition behind our bilinear attention mechanism is provided in Appendix
Implementation details. Additional details regarding the activation function, residual connections and normalization, the use of multiple heads in attention layers and the initialization of the C' x C weights in the BAM layers are provided in Appendices B.2 B.3 B.5 respectively.
CPDAG estimation from the graph skeleton and the set of moralized edges. To derive the CPDAG from the graph skeleton and identified immoralities, we train second neural network to infer v-structures from two parent nodes together with potential common child nodes that have edges to both parents, as well as other neighbor nodes related these parents. We iterate over all inferred immoralities, treating the two nodes involved in each immorality as parent nodes. By applying distinct layers to the data corresponding to the parent nodes, potential common children, and neighbors, we can break the symmetry among nodes and enable role-specific learning, thus facilitating edge directionality inference. The resulting CPDAG further refined using Meek rules 45) Details about the CPDAG estimation step can be found in Appendix D

---

3 Related Work
Causal discovery primarily relies unsupervised learning methods, including constraint baaseaddtt---eddddd approaches that infer conditional independencies [6]] [6] and score-based methods that optimize score function under acyclicity constraints [13 [1]. Some studies have estimated relations from multi-sequence alignment (MSA) data matrices, such as Rao et al. [54] using axial attention (26] and Li et al. (37) employing convolutional neural network on precision matrices. However, these models are designed for the specific conditions of the MSA prediction task and its fixed data dimension of 20 naturally-occurrigg residue types.
Our work closely aligns with Ke et al. [3I and Lorch et al. which also use attention mechanisms between samples and attributes to estimate adjacency matrices. The methodologies differ in their strategies for deriving an adjacency matrix from observational data representations. [31] employ an autoregressive transformer approach, whereas 442 utilize the dot product of embedding vectors derived from max-pout -ooolini across the sample axis of the observational data matrices. We demonstrate that deriving output adjacency matrices can be efficiently achieved through covariance matrices, allowing for direct processing of dependency information and enhanced learning efficiency through geometric learning on the SPD manifold.
Unlike [3], our model can be trained across different shapes of sample and attribute numbers, eliminating the need for re-training different datasets. While 42 's model can evaluate data with varying dimensions, their implementation has practical limitations in training across variable sample sizes, and it is restricted to handling a limited range of dimensions for training. Our model addresses identifiability challenges by first estimating an undirected graph, then proceeding to CPDAG estimation by testing immoralities and using Meek's ' rules. In contrast, BT and 42 deduce DAG using both observational and interventional data such that these models dependent on the availability of interventional data or making random guesses for edges that cannot be directly inferred.
Closely related to our approach, Li et al. (36) used permutation-invanne causal inference, but their method directly computes correlation matrices without learning high- dimensional representation of the observational data, which may limit its efficacy when covariance matrices are not sufficient. Despite preserving invariances, the small number of 5 free parameters per layer could restrict its representational power.
models for supervised
39 40 proposed supervised learning framework using kernel mean embeddings. 43 address identifiability with independence tests and cascade classifiers trained on vicinal graphs specific to observational data. Our method adopts a more general approach, training network that does not require re-training for new evaluation samples. [15)ss ss work on immoralities aligns with our second CPDAG estimation step.
4 Experiments
Training data hyperparameter settings. The hyperparameters for sampling training data, which
are defined in section 2.2 are set to d = 10, a = 100, = 1, q = min( 5). M = 50, and
M = 1000. This configuration allows for denser graphs than those explored in other studies
15 31 71).
Baseline algorithms. We evaluate BAM's performance against other algorithms: PC 61.. PC- HSIC 73,. reot and reit (62]. ccdr (3] GES 13 GIES (33) LiNGAM (60), MMPC 64 CAM [1],
SAM 30 all implemented in the causal discovery toolbox 29), GLASSO 17 as implemented by [50]. along with DAG-GNN [71]. Default hyperparameters as in (29] are used. For GLASSO, cross-validation for the sparsity parameter is performed as proposed by (50] As supervised causal discovery baseline, we use the AVICI model [2]. both pre-trained version (SCM-vO) and version trained from scratch on the same Chebyshev data as our method, using default hyperparameters and data dimensions d E {2, 5, 10, 20, 30, 40, 60, 80, 100). For the number of samples M, which can only be chosen as a single hyperparameter value for AVICI 42 due to memory allocation constraints, we selected M = 150. GLASSO and MMPC are excluded from CPDAG estimation as they focus on undirected graphs. For algorithms estimating directed edges, the skeleton computed for undirected prediction. We also compare SHD and accuracy results against a naive zero graph without edges.
7

---

PC
ccor
UNGAM
SAM
Auici SCM-VO Avtici Cheby
zero
HSIC
reit
GES
MMPC
CAM
DAGGNN
BAM
ABC
44-30 M7-300
Mx208
No-100
14-50
10-000
M4000
M-.00
M-3000
PH-5 M5110 M-200
M-d00
Figure 4: Undirected graph estimation results arranged from worst (left) to best (right) performance. (A, B) AUC values for different dependencies, with (A) d = 50, M 200, and (B) d 100, M 50. (C) Accuracy for the same dependencies as (B). (D, E) AUC vs. M for d = 100 with (D) Chebyshev and (E) cosine dependencies. (F) SHD vs. M for = 100 with Chebyshev dependency.
Avici SCMVOO
HSIC
relt
GES
MMPC
CAM
DAGGNN
Avici Cheby
BAM
Figure 5: CPDAG estimation results arranged from worst (left) to best (right) performance. (A-C) SHD for different dependencies, with (A) d = 100, M 500. (D) AUC vs. M for d = 100 with Chebyshev dependency. (E, F) SHD vs. M for d 100 with (E) Chebyshev and (F) sine dependencies.
20, M 200, (B) 50, M 200, and (C)
Performance indicators. We assess the performance of both undirected and CPDAG graph estima- tions using the Area Under the Precision-Recal Curve (AUC), which is well-suited for imbalanced binary classification tasks like sparse graph detection 244 For CPDAG estimations, we also employ
the Structural Hamming Distance (SHD), a standard metric in structure learning [7]] 31]. In the undirected graph estimation task. SHD is equivalent accuracy, defined as the percentage of correctly inferred edges. which we also report. Additional Structural Intervention Distance (SID) 52 results for CPDAG estimation are provided in Appendix G.5
Graph estimation results. Figure 4 showcases BAM's efficacy in undirected graph prediction. Trained on synthetic Chebyshev polynomial data, BAM consistently outperforms other methods

---

Expoesed
Expoeted
0.544
0 S4B
BBC
2.676
850
0 567
BCC
0 5.54
0.799
0 8.7
40,022
8.801
0
03:40
0
1.799
0.715
0 8.8
00
A0.000
A0:344
A0.071
050
AD.002
A0.000
0.000
Graph Diemiiutto
Cert CTCM
10/001
846
3.804
0.710
0090
80 003
8.800
S0/077
$0/000
0.833
0.728
0.717
0000
20/077
0.0
40/012
+0/043
60.301
60/399
+0:327
+0.000
40 000
600000
600007
+00000
40/004
0001
65 00001
60,000
+0.000
-07
Cordh
40/940
40/050
40/047
40/201
A0/001
40/775
ADU000
4D 100
40/054
40.035
AD/0H
40/256
AC 109
AC 200
40/000
AD
400077
40/090
10:045
A0.022
60/411
10:001
25
30 090
50 000
100-3
10,001
10/002
25 12:00
40 05000
10.074
55 007
10/002
-05
Chuda
622
Chedy
Line
C22
Chudy
Linerr
622
Figure 6: Heatmaps of mean AUC values with standard deviations across graph distributions and dependency functions for expected degrees 2, 3, and 4 obtained from 10 independent simulation runs, Graph distributions include Kénós (ER), Watts-Strogatz (WS) [69], Exponential Random Graph Model (ERGM) [56], configuration models with both homogeneous (Conf) and heterogeneous (Conf2) degree distributions [47], and geometric random graphs (Geom) [51].
across various dependency relations, dimensions, and sample sizes (Figure (A C)). It excels in capturing intricate non- non-monotonic dependencies (Figure 4 (A, B, E) and demonstrates superior performance over Avici on the training set (Figure (D, F)).
Figure 5 displays the results for CPDAG estimation tasks. In high-dimensional scenarios (d =d= 100, M = 50), no algorithm surpassed the baseline of zero-graph in SHD, as shown in Appendix [G.4] Therefore, our SHD analysis focuses on the low-dimensional setting shown in (A) and two moderate- dimensional (M > d) settings in (B) and in in (C) across various dependencies. Across these scenarios, the two-step method of our algorithm remains competitive, though its advantage is less pronounced than in the task of undirected graph estimation. Panels (D) and (E) depict the AUC and SHD across varying M values for the training datasets of BAM and the re- re-trained Avici, illustrating BAM's competitive performance
learning CPDAG-structure. Specificallyy, panel (E) presents SHD values for sine dependency across different M values for d = 100, further demonstrating the robustness of BAM in non-linear settings. Additional evaluation results for more complex multivariate dependencies using random Fourier features and MLPs are provided in Appendix G.3 and Figure 9 respectively, Additional SID results for CPDAG estimation in Appendix [G.5 demonstrate that our method achieves the highest performance for the best DAG of the estimated equivalence class. Compared to competing algorithms, it infers fewer edge orientations.
Graph distribution robustness. To demon- strate that our architecture generalizes beyond the Méndos-Rénii graphs it was trained on, we evaluate its performance on graphs generated using various random graph models. Our re- sults, shown in Figure 6] demonstrate consis- tent performance across all graph distributions, suggesting that our method robust to vari- ations in underlying graph structure. Details about the configuration of the distributions can be found in Appendix G.6
Table 1: Ablation study results featuring loss values. Data generated under Chebyshev dependencies. indicates the removal of corresponding layer. Param quantifies the difference the number of parameters to the full model.
model
loss
A Param
FULL
0.173 + 0.007
-BAM
0.202 + 0.005
80 K
- -AMM -LogEig
0.271 + 0.012
100 K
Ablation study. Table presents our abla- tion study results, underscoring the crucial role of the bilinear layer, whose removal sig- nificantly increases loss metrics. Omitting the I LogEig layer leads to a substantial deteriora- tion loss, indicating that direct predictions from the SPD-space are problematic. Notably, bilinear
-bss. att.
0.189 + 0.006
120 K
obs. att. -Dense
0.205 + 0.006
160 K
data processing alone yields relatively good results, possibly due to the embedding layer's ability to

---

decode non- -oneiinesssties in the data, which are lost when solely relying on covariance matrices. Further details are provided in Appendix F2
Effectiveness of novel two-step approach for CPDAG estimation. To validate our approach of identifying immoralities through the first network step, we ran additional experiments where we maintained the undirected edge estimation network but modified the second step. Specifically, we first obtained an undirected graph as the union of skeleton and moralized edges, then trained another neural network with identical hyperparameters to test each possible immorality in this graph, not just those identified by our first network. Table shows that our selective testing strategy consistently out- performs this exhaustive approach across different settings, particularly for Chebyshev dependencies used in training. This suggests that conditional dependencies are more accurately identified when leveraging information from the entire graph structure, rather than examining estimated local Markov blankets independently.
Table CPDAG estimation performance (AUC scores) with M = 200 samples. The baseline (' ('All Moralizations' AM) tests all possible moralizations, while BAM only tests those identified by the first network. Mean + standard deviation over 10 runs.
Dependency Function
Model
d = 10
d = 20
d = 50
d = 100
BAM
0.79 = 0.09
0.85 = 0.05
0.78 = 0.05
0.76 = 0.03 0.70 = 0.01
Chebyshev
AM
0.73 = 0.03
0.69 = 0.02
0.71 = 0.02
BAM
0.71 = 0.06
0.72 = 0.02
0.69 = 0.06
0.70 = 0.02
Linear
AM
0.67 = 0.03
0.63 = 0.06
0.63 0.04
0.68 0.03
BAM
0.69 = 0.16
0.68 = 0.09
0.62 = 0.07
0.66 = 0.04
AM
0.63 = 0.08
0.62 = 0.07
0.60 = 0.06
0.64 = 0.04
BAM
0.71 0.04
0.74 0.05
0.65 = 0.11
0.59 0.04
CO8
AM
0.67 0.07
0.64 = 0.05
0.62 - 0.07
0.59 = 0.05
BAM
0.77 0.05
0.77 + 0.09
0.71 + 0.08
0.63 - 0.11
x²
AM
0.68 + 0.06
0.64 + 0.05
0.64 0.05
0.61 + 0.05
BAM
0.56 + 0.03
0.55 + 0.17
0.54 + 0.09
0.59 + 0.06
2
AM
0.58 + 0.05
0.50 + 0.10
0.56 + 0.09
0.57 0.05
Efficiency. Training neural network takes approximately 6 hours
A-100 GPU with
81,920 MiB of graphical memory, while inference typically requires less than a few seconds and can be O(CMd + Cd2 + Md + 2 C²). The runtime complexity of our approach is OCC² Md + CMd CMdd C - C We present detailed breakdown of the single components in Appendix H Empirical evaluations show that our model achieves significantly lower computational times compared to most unsupervised approaches, shown in Appendix G.7
normal computer. The overall memory complexity of our proposed method is
5 Conclusion
In this study, we introduced novel neural network model for supervised graph structure learning that addresses identifiability issues in observational data through a two-step approach. Our approach extends the classical GLASSO paradigm to the neural network era by incorporating novel bilinear attention mechanism operating in both Euclidean and SPD spaces, effectively learning shape-invariant transformations directly from data rather than relying on pre-derived mathematical operations. The trained model effectively detects causal relationships and generalizes well across different functional forms of nonlinear dependencies, while requiring only a single forward pass for inference. Compre- hensive empirical evaluations demonstrate that our approach consistently outperforms state-of-the-art methods across diverse scenarios. Future directions include extending approach to handle more complex settings and investigating its performance under various noise and confounding conditions, well as scaling through local attention mechanisms. Limitations and future directions are further discussed in Appendix(]
10

---

Acknowledgments and funding disclosure
This project (HA project no.: 1304/22-09) is funded by the State of Hesse and the HOLM funding programme part of the 'Innovations in the field of logistics and mobility' initiative of the Hessian Ministry of Economies, Energy. Transport and Housing. Additional funding was provided by the Federal Ministry of Education and Research (BMBF, project KI4ELM 011S24022C) and the European Union's Horizon 2020 programme iPC (undriddalizzePPaaiarriccur) project under grant agreement No. 826121.
HESSEN
Bundesministerium für Bildung Forschung
Hessisches Ministerium für Wirtschaft, Energie, Verkehr und Wohnen
FRANKFURT HOLM
HESSEN
Funded by the European Union
HessenAgentur HA Hessen Agentur GmbH
The computational work was supported by the High-Performance Computing facility Lichtenberg at the NHR Centers NHR4CES TU Darmstadt.
Competing interests: The authors declare no competing interests.

---

References
[1] Martín Abadi, Ashish Agarwal, Paul Barham, Eugene Brevdo, Zhifeng Chen, Craig Citro, Greg S. Corrado, Andy Davis, Jeffrey Dean, Matthieu Devin, al. TensorFlow: Lcrge-scale machine learning on heterogeneous systems, 2015. BRLLhttps:///e URLlattpp//wwwtttoorrroo tensorflow..ooooo. org/ org/ Software available from tensorflow.org.
[2] Andrea Apicella, Francesco Donnarumma, Francesco Isgrò, and Roberto Prevete. A survey on modern trainable activation functions. Neural Networks, 138:14-32, 2021.
[3] Bryon Aragam and Qing Zhou. Concave penalized estimation of sparse Gaussian Bayesian networks. JMLR, 16(1):2233-3228, 2015.
[4] Sercan Ö. Arik and Tomas Pfister. TabNet: Attentive interpretable tabular learning. In AAAI, volume 35, 2021.
[5] Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton. 1607.0600 2016.
Layer normalization.
[6] Thomas Bachlechner, Bodhisattwa Prasad Majumder, Henry Mao, Gary Cottrell, and Julian McAuley. ReZero is all you need: Fast convergence large depth. In UAI, 2021.
[7] Onureena Banerjee, Laurent El Ghaoui, and Alexandre d' d'Appemmont. Model selection through sparse maximum likelihood estimation for multivariate Gaussian or binary data. JMLR, 9: 485-516, 2008
[8] Wolfram Barfuss, Guido Previde Massara, Tiziana Di Mattco, and Tomaso Aste. Parsimonious modeling with information filtering networks. Physical Review E, 94, 2016.
[9] Rajendra Bhatia. Positive definite matrices. In Positive Definite Matrices. Princeton University Press, 2009.
[10] Peter Bühlmann, Markus Kalisch, and Lukas Meier. High-dimensional statistics with view toward applications biology. Annual Review of Statistics and Its Application, 1(1):255-278, 2014.
[11] Peter Bühlmann, Jonas Peters, and Jan Ernest. CAM: Causal additive models, high-dimensional order search and penalized regression. The Annals of Statistics, 2014.
[12] David M. Chickering. Learning equivalence classes of Bayesian 445 1455498,22 2002.
Bayesian-network network structures. JMLR, 2:
[13] David M. Chickering. Optimal structure identification with greedy search. JMLR, 3:507-554, 2002.
[14] Hoon Chung, Sung Joo Lee, and Jeon Gue Park. Deep neural network using trainable activation functions. In IJCNN, 2016.
[15] Haoyue Dai, Rui Ding, Yuanyuan Jiang, Shi Han, and Dongmei Zhang. ML4C: Seeing causality through latent vicinity. In Proceedings of the 2023 SIAM International Conference on Data Mining (SDM), 2023.
[16] Mathias Drton and Marloes H. Maathuis. Structure learning in graphical modeling. Annual Review of Statistics and Its Application, 4(1):365-33) 2017.
[17] Jerome Friedman, Trevor Hastie, and Robert Tibshirani. Sparse inverse covariance estimation with the graphical lasso. Biostatistics, 933):4324444 2008
[18] Tobias Gerstenberg, Noah D. Goodman, David A. Lagnado, and Joshua B. Tenenbaum. A counterfactual simulation model of causal judgments for physical events. Psychological Review, 128(5):936, 2021.
[19] Clark N. Glymour, Kun Zhang, and Peter Spirtes. Review of causal discovery methods based on graphical models. Frontiers in Genetics, 10:524, 2019.
12

---

[20] Gene H. Golub and Charles F. Van Loan. Matrix computations. JHU press, 2013.
[21] Dominique Guillot and Bala Rajaratnam. Functions preserving positive definiteness for sparse matrices. Transactions of the American Mathematical Society, 367(():6277-44) 2015.
[22] Aric Hagherg, Pieter Swart, and Daniel Chult. Exploring network structure, dynamics, and function using NetworkX. Technical report, Los Alamos National Lab.(LANL), Los Alamos, NM (United States), 2008.
[23]
Alain Hauser and Peter Bühlmann. Characterization and greedy learning of interventional Markov equivalence classes of directed acyclic graphs. JMLR, 13(1):2409-2444 2464, 2012.
[24] Haibo He and Edwardo A. Garcia. Learning from imbalanced data. IEEE Transactions Knowledge and Data Engineering, 21(9): (203)):223-122000 1284,, 2009.
[25] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In CVPR, 2016.
[26] Jonathan Ho, Nal Kalchbrenner, Dirk Weissenborn, and Tim Salimans. Axial attention in multidimensional transformers. arXiv:/9/2.12//00 12800, 2019.
[27] Kurt Hornik. Approximation capabilities of multilayer feedforward networks. Neural networks, 4(2):251-277, 1991.
[28] Zhiwu Huang and Luc Van Gool. Riemannian network for SPD matrix learning. In AAAI, 2017.
[29] Diviyan Kalainathan, Olivier Goudet, and Ritik Dutta. Causal discovery toolbox: Uncovering causal relationships in Python. JMLR, 21(1):1406-144 2020.
[30] Diviyan Kalainathan, Olivier Goudet, Isabelle Guyon, David Lopez-Paz, and Michèle Sebag. Structural agnostic modeling: Adversarial learning of causal graphs. JMLR, 23(219):1-62, 2022.
[31] Nan Rosemary Ke, Silvia Chiappa, Jane X. Wang, Jorg Bornschein, Anirudh Goyal, Melanie Rey, Theophane Weber, Matthew Botvinick, Michael Curtis Mozer, and Danilo Jimenez Rezende. Learning to induce causal structure. ICML: Workshop on Spurious Correla- tions, Invariance and Stability, 2022.
[32] Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In ICLR, 2015.
[33] Daphne Koller and Nir Friedman. Probabilistic graphical models: Principles and techniques. MIT press, 2009.
[34] Jannik Kossen, Neil Band, Clare Lyle, Aidan N. Gomez, Thomas Rainforth, and Yarin Gal. Self-Attention between datapoints: Going beyond individual input-output pairs in deep learning. In NeurlPS, 2021.
[35] Steffen L. Lauritzen and David J. Spiegelhalter. Local computations with probabilities graphical structures and their application to expert systems. Journal of the Royal Statistical Society: Series (Methodological), 50(2):157-194, 1988.
[36] Hebi Li, Qi Xiao, and Jin Tian. Supervised whole DAG causal discovery. 2XXx:20006006997, 2020.
[37] Yang Li, Jun Hu, Chengxin Zhang, Dong-Jun Yu, and Yang Zhang. ResPRE: high-accuracy protein contact prediction by coupling precision matrix with deep residual neural networks. Bioinformatics, 35(22):4647-4655, 2019.
[38] Min Lin, Qiang Chen, and Shuicheng Yan. Network in network. arXiv preprint arXiv: 1312.4400, 2013.
[39] David Lopez-Paz, Krikamol Muandet, and Benjamin Recht. The randomized causation coeffi- cient. JMLR. 16:2901-2907, 2015.
13

---

[40] David Lopez-Paz. Krikamol Muandet, Bernhard Schölkopf, and lliya Tolstikhin. Towards learning theory of cause-effect inference. In ICML, 2015.
[41] David Lopez-Paz, Robert Nishihara, Soumith Chintala, Bernhard Schölkopf, and Leon Bottou. Discovering causal signals in images. In CVPR, 2017.
[42] Lars Lorch, Scott Sussex, Jonas Rothfuss, Andreas Krause, and Berhhard Schölkopf. Amortized inference for causal structure learning. In NeurIPS, volume 35, 2022.
[43] Pingchuan Ma, Rui Ding, Haoyue Dai, Yuanyuan Jiang, Shuai Wang. Shi Han, and Dongmei Zhang, ML4S: Learning causal skeleton from vicinal graphs. In SIGKDD, 2022.
[44] John C. Mason and David C. Handscomb. Chebyshev polynomials. Chapman and Hall/CRC, 2002.
[45] Christopher Meek. Causal inference and causal explanation with background knowledge. In UAI, 1995.
[46] Maher Moakher and Mourad Zéraï. The Riemannian geometry of the space of positive-definie matrices and its application to the regularization of positive-definite matrix-valuce data. Journal of Mathematical Imaging and Vision, 40(2):171-188) 2011.
[47] Michael Molloy and Bruce Reed. A critical point for random graphs with given degree sequence. Random structures algorithms, 6(2-3):161-180, 1995.
[48] Peer Nowack, Jakob Runge, Veronika Eyring, and Joanna D. Haigh. Causal networks for climate model evaluation and constrained projections. Nature communications, 11(1):1415, 2020.
[49] Judea Pearl. Probabilistic reasoning in intelligent systems: Networks of plausible inference. Elsevier, 2014.
[50] Fabian Pedregosa, Gaël Varoquaux, Alexandre Gramfort, Vincent Michel, Bertrand Thirion, Olivier Grisel, Mathieu Blondel, Peter Prettenhofer, Ron Weiss, Vincent Dubourg, al. Scikit- learn: Machine learning in Python. JMLR, 12(85):2855-2830 2011.
[51] Mathew Penrose. Random geometric graphs, volume 5. OUP Oxford, 2003.
[52] Jonas Peters and Peter Bühlmann. Structural intervention distance for evaluating causal graphs. Neural computation, 27(3):771-79,, 2015.
[53] William H. Press. Numerical recipes 3rd edition: The art of scientific computing. Cambridge university press, 2007.
[54] Roshan M. Rao, Jason Liu, Robert Verkuil, Joshua Meier, John Canny, Pieter Abbeel, Tom Sercu, and Alexander Rives. MSA transformer. In ICML, volume 139, 2021.
[55] Douglas A. Reynolds. Gaussian mixture models. Encyclopedia of biometrics, 741(659-663), 2009.
[56] Garry Robins, Pip Pattison, Yuval Kalish, and Dean Lusher. An introduction to exponential random graph (p*) models for social networks. Social networks, 29(2): 1911 2007.
[57] Isaac J. Schoenberg. Positive definite functions spheres. Duke Mathematical Journal, 9: 96-108, 1942. MR0005922 (3,232c).
[58] Rajen D. Shah and Jonas Peters. The hardness of conditional independence testing and the generalised covariance measure. The Annals of Statistics, 48(3), 2020.
[59] Gal Shalom, Eran Treister, and Irad Yavneh. PISTA: Preconditioned iterative soft thresholding algorithm for graphical lasso. arXiv:2205. /0027. 2022.
[60] Shohei Shimizu, Patrik O. Hoyer, Aapo Hyvärinen, Antti Kerminen, and Michael Jordan. linear non- non- Gaussian acyclic model for causal discovery. JMLR, 7(10), 2006.
14

---

[61] Peter Spirtes, Clark N. Glymour, Richard Scheines, and David Heckerman. Causation, predic- tion, and search. MIT Press, 2000.
[62] Eric V. Strobl, Kun Zhang, and Shyam Visweswaran. Approximate kernel-based conditional independence tests for fast non-parametric causal discovery. Journal of Causal Inference, 7(); 201 20180017, 2019.
[63] Lloyd N. Trefethen. Is Gauss quadrature better than (lenshaw-Curtis! SIAM Review, 50(1): 67 87, 2008.
[64] Ioannis Tsamardinos, Laura Brown, and Constantin Aliferis. The max-min hill-climbing Bayesian network structure learning algorithm. Machine learning, 65:31-78, 2006.
[65] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, ukasz Kaiser, and Illia Polosukhin, Attention All you Need. In NeurlPS, volume 30, 2017.
[66] Matthew Vowels, Necati Cihan Camgoz, and Richard Bowden. D' like DAGs? A survey on structure learning and causal discovery. ACM Computing Surveys, 55(4):82:1-82330 2022.
[67] Huiyu Wang. Yukun Zhu, Bradley Green, Hartwig Adam, Alan Yuille, and Liang-Chies Chigh Chen. Axial-DeepLab: Stand-alone -xial-attention for panoptic segmentation. ECCV, 2020.
[68] Rui Wang, Xiao-Jun Wu, Tianyang Xu, Cong Hu, and Josef Kittler. Deep metric learning on the SPD manifold for image classification. IEEE Transactions on Circuits and Systems for Video Technology, 2022.
[69] Duncan J. Watts and Steven H. Strogatz. Collective dynamics of 'small-world" networks. Nature, 393(6684):440-444 1998.
[70] Sewall Wright. Correlation and causation. Journal of Agricultural Research, 20(7):557-585, 1921.
[71] Yue Yu, Jie Chen, Tian Gao, and Mo Yu. DAG-GNN: DAG structure learning with graph neural networks. In ICML, 2019.
[72] Ming Yuan and Yi Lin. Model selection and estimation in the Gaussian graphical model. Biometrika, 94(1):19-35, 2007.
[73] Kun Zhang, Jonas Peters, Dominik Janzing, and Bernhard Schölkopf. Kernel-based conditional independence test and application in causal discovery. In UAI, 2011.
15

---

A Properties and theorems regarding the custom softmax function
Similar to the standard softmax, softmax for SPD matrix returms positive values summing to d, but unlike the standard softmax, the rows do not sum to We demonstrate that our modified softmax function ē regularizes the eigenvalues:
Theorem 1. For any S E S the largest eigenvalue of ( is 1.
Proof. Let 3 = eS. By similarity transformation, the eigenvalues of (SS) are equal to the eigenvalues of SA(S). It holds that SA(S)S1 $1, which demonstrates that S1 is the Perron eigenvector corresponding to the eigenvalue 1. The assertion now follows from the Perron-Frobenine theorem.
Proposition The custom softmax is invariant to additive shifting, i.e., ö(S + a) f(S) each a E I, S sdxd.
Proposition 2 shows that or unlike standard softmax, does not need scaling constant, and it can easily manage exploding exp values via maximum-value scaling.
B
Architectural details
B.1 Observational attention
This layer accepts an input X E xxXx from which generates keys K XW K Rxdxx queries Q = XWQ E TRdxxx and values V = XWv E Mxxx with W Wq E RCx, and W, E RCXC
For each ,,.... M, keys and queries combined in parallel along the inner axis, leading to
K @ Q = (Km.. m....QQuu Qm,;t) mm I. E IR M xdxd
(4)
This results in the attention weights
K Q T(L3.2)
A
R Mxxd
Ve
where (pperm) denotes a permutation of the axes according to a permutation perm and is the softmax operation along the last axis. For example, if perm = [1,3, 2], the operation would rearrange the second and third axes of the tensor while the first axis stays unchanged. Finally, for each m, the output computed as
H V
B.2 SPD activation function
To obtain an activation function for the SPD net, we leverage the following theorem, a direct consequence of (Theorem 4.11 Guillot and Rajaratnam 2015). which is based the work of [57]: Theorem Any continuous function from 1(--1,1),[- 1)) acting elementwise on a matrix preserves positive definiteness if it can be expressed as series f [-1,1] + [-1,1] with
f(x)
wk
subject to
Wk < with wy 20 0, N.
We employ this theorem to construct an activation function for the SPD neural network, using relatively small maximal polynomial degree value of Nm = 3. This function employs trainable weights wP and is applied after' corrrllation normalization", i.e., the conversion covariance matrices into correlation matrices. Trainable activation functions using low-degree Taylor polynomials were also proposed in 14] for general neural networks, not focusing on SPD data. Additionally, provides various types of trainable activation functions for Euclidean neural networks.
16

---

B.3 Residual connections and normalization.
We adopt normalization and residual connection [25] on the imput of all attention layers, expressed as Particularly when M the residual connection can enhance the rank of covariance matrices, so full-rank representations can be attained even for under-determined problems. For learnable residual scaling, we utilize recently established methods (6].
Normalization critical in attention networks, with various methodologies available 25 51.
S (Layerooo
Within the SPD manifold, we employ correlation normalization alongside residual connections. This approach preserves positive definiteness and corresponds intuitively standard normalization in Euclidean space.
B.4
Multiple heads
In all attention layers, employ multihead attention, process that divides the input tensor along the channel axis into several smaller tensors. Each of these is then subjected to attention independently.
B.5 Weight initialization
For initialization of the C' x C weight matrices W+ in each bilinear attention layer, we draw initial weights from U (0, )) Given tensor of SPD matrices E E SdxdxC in expectation, the diagonal entries each of the d x d matrices in EW+ should stay within a comparable range, while off- diagonal entries are expected to be drawn to zero due to the symmetric distribution of positive and negative values, resulting in perturbed identity matrices for keys and queries.
17

---

C Theoretical foundation for the three-class edge classification problem
Building upon the foundational principles of Markov and faithfulness [33 demonstrate that the three-class classification problem can be theoretically deduced from the distribution of the nodes by examining the following independence relations: If there is directed edge from node X to Y in the DAG, then and are dependent given any set from the power set of the other nodes, i.e.:
X = X y Y
AC E P(V (X,Y)).
For an immorality between X and Y, there exists a set of nodes C' E P (X,Y)) within the power set of all other nodes such that X and Y are conditionally independent given this set (e.g., the set of all common ancestors of and Y, or the set of all parents X or Y), but dependent given all other nodes in the graph, i.e.:
X + 2 + Y, + Y, AY = 3C P(V (X X IL Y | M.Y V (X,)).
If there is edge between X and Y, and if X and Y do not have common child, then X and Y are conditionally independent given all other nodes, i.e.,
X+Y, Y,X + Y, A2 X Z
XXLYXXXXX
Note that testing for the no-edge class is cost-effective, as one only needs test for a single set, rather than checking for any 7 E V (X,Y) Y) if there is a v-structure X - Z t- Y. Testing for sepset to differentiate between the skeleton and moralized edge classes is more intricate. The neural network is tasked with learning an approximation for this distinction.
D Details of the CPDAG estimation model
To estimate CPDAG from the graph skeleton, along with the of immoralities between pairs of nodes, we test each estimated immorality to determine which potential common child nodes indeed common children. The parent nodes, denoted by pa, are the nodes between which immorality was first estimated. Potential common child nodes, denoted by ec, are nodes that have an edge both parent nodes. Neighbor nodes, denoted by ne, are nodes that have one edge to exactly one parent.
We take the columns Tas R2Xx2 Fee R Tne E R.M x/Mx of the data matrix X corre- sponding to pa, dimensionality expansion to Ta E IR Mxx1 Zoe E MM Tne E Mx Next, three separate feed-forward subnetworks lpa: le: Ine are applied the three inputs T Tees Tne- Each of these layers has the same architecture: For an input ã IRMx Mxd11 weight matrices W, RIx and W1 E CxC, together with a bias vector by E RC are used to embed æ to
inputs for the neural network. Each of these submatrices undergoes
h tanh + b,) E MxdxC Then, two residual layers of the form
hi+1 h; + tanh (h, W. :V +b,W,
with W, E ROx W E RCxC b; E RC, 1, 2 are applied to obtain representations hpu E Mx2X hec E &RMxlexx hac E Mx For the addition of the bias terms, broadcasting is used, i.e.
EERC b RxxxXC with bm.Le- b. m m= M, 1=...,44 d, c I.....C.
Now, the representations are concatenated along the dimension axis to obtain M M [2+|cc|+ +|nel)xxC tensor. We use d = 2 + |cc| + |ne]. We employ the same dnservation-toodepeeedeed network as before, but instead of using softmax on the output of the LogFig-Laaer we use mean-pooling
X (+ T3221
to inflate of the variable axes of the dx output of LogEig obtain dxC batch of C vectors dimension d. After applying a dense C' x 1 layer, we obtain a vector of length d. Now, the entries corresponding to the potential common children be sliced and backpropagated for training.
The network architecture is shown in figure
18

---

N1
P2
C1
C2
P1
N2
N2
Samples (M))
Variables (d)
Potential common
children
neighbours
Parents
Feed
Common Children
Feed- Fead-Forwate Forward
Feed Feed-Forwatd Forward Nelghbours
Riemannian manifold
mean pooling
denwe siggmid
alice out childrwn
covariancs
LogEig
Probabilities common childrrm
elementte activation
otteation
attention
datapointa
Figure 7: CPDAG Estimation Architecture: In the graph, black edges represent undirected connec- tions, whereas the red edge signifies an immorality. Columns in the data matrix corresponding to parent nodes, potential common children, and neighboring nodes are taken as input for the neural network and are processed through three distinct feed-forward networks. The embeddings are con- catenated along the variable axis and processed through network with data matrix attention, bilinear attention, and a LogEig layer. Mean pooling and sigmoid activation are applied to output probabilities for potential common children.
19

---

E
Parameterization of the SEM
E.1 Chebyshev polynomials for training
For the Chebyshev polynomial, utilize the following parameterization:
fe C1, c2)
Bw
On T (w)) + +1
")Epag())
+ Om
Σ
S, Tn.TTlla (m
+
Σ wEpa(g)
(2)
vu E V,
where 5 the degree, and Tn denotes the Chebyshev polynomials of the first kind (scaled for the input to have maximum absolute value of 1), and bivariate polynomials are given by:
(M v) = (x Mx) (gy
(6)
+ Mm (1 [1yy
where in the index stands for "multidimensional."
Here, Pz U|-1, 1]. By ~ U[-1, I]. and the coefficients a,,..., Cry, O' are calculated from 0g En On tum for 1,.... F, with ã 7. 7s ~ UI-1, 1], and om ~ U1-1, Par pal
8..
with Au U10.7,1.33. 1.3). 6,, are random weights with 6,,t te [6,el" Ja,t ~ U 11,1,
8.,
This parameterization is motivated by the observation that for a smooth function-whert higher- order derivatives are not significantly larger than the lower-order ones-the coefficients of the Chebyshev approximation decrease neural network 'typical' smooth functional dependencies. Furthermore, this suggests that using Chebyshev polynomials of degree r = 5 is not significant limitation, as the coefficients of higher orders are already negligibly small.
factorial manner. This provides a rationale for training the
In order to prevent the values from exploding and to properly account for the common domain of Chebyshev polynomials, we implement several measures. Firstly, we scale each input to the SEM by the maximum value within the batch. Secondly, we standardize all variables; we subtract the mean and divide by the standard deviation for each batch. This ensures the variables are both centered and scaled. To further improve stability and robustness our model, we introduce thresholds for any absolute values exceeding 5, thereby mitigating the potential impact of outliers.
E.2 Gaussian mixture error terms
The additive етrоr term follows Gaussian mixture distribution. We randomly determine the number components L ~ U,,,..., from a discrete uniform distribution. Each component has randomly assigned parameters for the means fr U[-1. 1]. standard deviations such that E ~ Et= wiNV (M of). The multplicatiiv error term is uniformly distributed, with E2 ~ 1[,1,1] We scale the error to zero mean and variance 1 after simulation.
U (0.05, 1],
and weights कn U00.3,, 1], wi Σ. we
E.3
Testing dependencies
To create testing data, we create synthetic data according to an SEM equipped with different depen- dency function, while the error term follows Gaussian mixture distribution as before. We use the following dependencies for testing: Chebyshev, linear, sine, cosine, x², x² multidimennionall multi- plicative dependency. The chebyshev-dependene used was the same as in the training procedure. We used the following testing dependency functions: x, sin(x), cos(x:), x² x as g(x)
Fiest pap 8)
Σ wEpa(()
ong(Tw) +
20

---

For the multi-dimensional multiplicative test dependency, we used
62) - Otm Σ
21 + Σ (m( 62)
wo E V,
wCEa(()
with
(m(2,)) = - Mx) pyy
(1+ + Mx [1 + ply
with Ma ~ U[-1, 11, My U1- dett Enncculrr 10,,ol Japt U[-
3,
21

---

Training details
E1
Model hyperparameters
The implementation was performed using TensorFlow [. We employed the ADAM optimizer by Kingma and Ba 322 For training, used the hyperparameters stated in Table 3
Table Hyperparameters for the undirected graph estimation
Table 4: Hyperparameters of the CPDAG es- timation model
Hyperparameter
Value
Hyperparameter
Value
Layer Parameters
Layer Parameters
Number of channels
100
Number of channels C
100
Number of inner channels
100
Number inner channels e
100
Maximal degree activation function
Maximal degree activation function
Attention heads
5
Attention heads
Number of layers
Number of layers
Attention between attributes
Attention between attributes
10
Attention between samples
10
Attention between samples
x C dense observational layers
C' x C dense observational layers
10
bilinear attention + SPD activation
10
bilinear attention + SPD activation
10
Training Schedule epochs
Training Schedule epochs
1000
1000
samples per epoch Initial learning rate
128
matrices per epoch
0.0005
Initial learning rate
0.0005
Learning rate decrease factor Minibatchsize
( (10) 1/500
Learning rate decrease factor Minibatchsize
(10 1/1000
Additionally, we generated data with random number of samples M U150, 51,.... 1000} and random variable dimension d ~ U100 11...., 1003.
Ablation study setting
In the ablation study, evaluate the performance of the full model in comparison models with reduced complexities. The full model is comprised of two attention between attributes layers, two attention between samples layers, two dense layers, and four bilinear layers equipped with SPD activation functions. This setup maintains parity between the number of attention layers operating on observational data and those focusing on covariance data, while also ensuring a comparable parameter count across different configurations. All models in the study utilize C' trained over 500 epochs, with each epoch comprising 128 data matrix adjacency label pairs. Again, we generated data with random number of samples M (50, 51, ... 1000} and random variable dimension U (10, 11, 1003.
100 channels and are
F.3 Loss function for the three-class edge classification problem
We employ the categorical cross cross-entropy loss function for classifying edges into one of three cate- gories: no-edge, skeleton edge, and moralized edge. Additionally, to enforce the condition that moral edge between nodes X and Y should only be predicted if there is potential common child 7 (i.e., Z and Z Y), we introduce penalty term, fp-
The overall loss function is defined as La + L. + Lp. Here, Lg the categorical crossentropy of the three categories given by
C. C. H(A.
ΣΣ> 2=
A log (r
22

---

denotes the categorical crossentropy of the three categories no-edge, skeleton edge, and moralized edge with
with Asj.. {
if (i,j) is in category c in the ground-truth DAG
E R R
clse
and Augic denotes the estimation by the algorithm on it. A and A are symmetric along its first two axes, i.e., Aij.e Aji,,, i, = -....,d, d,
c ,.......
Lo denotes the binary loss of no-edge vs. any edge present (present edges = skeleton edges U moralized edges):
()) " = (( A
A log( (og(A)) (1 A! log(1 A33
with
A E Rdxd
with
1
1
if an edge is estimated between i and j
else
being the adjacency matrix of no-edge (direct edge moralized edge). The penalty term, Cp is defined as:
Cp(A): = max (Ās [AgAg]
0).
Å Ag, E Rdxd are estimates of A.,, A,, and A by the algorithm respectively. This term penalizes the prediction of moralized edge in the absence of potential common child edges. The square root operation is applied element-wise.
23

---

G
Further experiments
G.1 Error bars
In our experiments, each algorithm was evaluated on five distinct trials, each involving a unique data matrix and corresponding ground-truth graph. The bars in the figures represent the mean performance values across these trials, while the error bars indicate the standard deviation. It is important to bars, is predominantly due to
note that the observed variability, manifested as relatively large the random sampling of graph degrees, which has a substantial influence on estimation accuracy. Despite this inherent variability, the comparison across algorithms remains valid, each algorithm is tested for the same graphs. Therefore, the magnitude of the error bars should not be interpreted as undermining the reliability of our findings.
G.2 Additional results undirected graph estimation
GES
0.7
CAll
SCM-0
Arii Chetz
Chbbyter
Liner
4
Mulldim
Chibyshav
4
Holldom
(a)
(b)
Figure 8: AUC values for undirected graph estimation in low-dimensional regimes: (a) d 10, M 200 (b) d 20, M 500.
codr
SAM DAGGNN
Avici SCM-VO
HSIC
ralt
GES
MMPC
CAM
Avici Cheby
BAM
(
bb)
(c)
Figure 9: AUC values for undirected graph estimation for random MLP dependency: (a): d 20.
(b): d = 50, (c): d = 100
Figures 8 and 9 provide supplemental data on the task of undirected graph estimation. Figures g showcases performance in low-dimensional settings characterized by 10, A =00 and d = 20, M = 500. In this scenario, our method (BAM) also outperforms competing graph inference algorithms. To further assess its capability to recognize multidimensional dependencies. we extended our tests to cases where the dependency function within the SEM is modeled via a randomly initialized multilayer perceptron (MLP) with a random number layers ~ U([,,.... 5). a random number of hidden layers ~ (4, 64), and relu or tanh activation with probability 0.5 each. Despite these complexities, our algorithm maintained state-of-the-are performance, delivering AUC scores competitive to the por-eerforning existing methods such depicted in Figure 9
24

---

G.3 Additional results on random Fourier feature dependencies
We evaluate our method's performance on more complex dependency structures using random Fourier features, which represent multidimensional non-linear relationships between variables. For undirected graph estimation (Figure A-),, our approach demonstrates strong performance across different graph sizes and evaluation metrics, particularly outperforming baselines for most sample sizes. The performance slightly decreases for very large sample sizes (M 500, 1000), potentially due to overfitting to the training distribution when encountering dependencies that deviate significantly from [F))F)
the training patterns. For CPDAG estimation (Figure D-F), while our method shows reasonable performance measured by AUC and SID, there remains room for improvement in handling multidimensional dependencies. These results complement our findings on MLP dependencies (Figurel)) indicating that approach can effectively handle various types of con-linear relationships while highlighting specific scenarios for future enhancements.
ccdr
glasso
Avici SCM-VV
zero
HSIC
reit
GES
MMPC
CAM
DAGGNN
Avici Cheby
BAM
B
M4-50
54-000
P4500
M4-00
M4200
M+500
MM-300
M4-30
M-330
M-3006
200
10:00
180
100
Figure 10: Additional computational results using random Fourier feature dependency. (A, B) AUC for undirected graph estimation with (A) d = 20, M 200 and (B) d = 50, MM 200. (C) Accuracy for undirected graph estimation with d 50, .M = 200. (D) AUC for CPDAG estimation with d 50, M 200. (E, F) SID for CPDAG estimation with (E) d 20, M 200 and (F) d 50, M 200.
G.4 Additional results on CPDAG estimation
rost
0
CES
0.77
CACOAN
50
0.99
2310
Cwly
Clebystev
Lineer
Muitdin
Chsbyshov
Linea
55-
Motdden
(a)
(b)
Figure 11: (a) SHD values for the high-dimensional CPDAG estimation d = 100, M = 50. (b) AUC values for the high-dimessional CPDAG estimation 100, M
50.
CPDAG estimation in high-dimensional settings presents significant challenges. the specific case of d 100 and M 50, none of the algorithms we evaluated could outperform zero-graph (i.c.,
25

---

graph with no edges) baseline in terms of Structural Hamming Distance (SHD). The results, depicted in Figure (a),, substantiate this observation and suggest that CPDAG estimation remains a difficult problem under these conditions, at least with our chosen graph density setup. Given the complexities encountered in high-dimensional contexts, our analysis primarily emphasizes the evaluation of AUC, as illustrated in Figure bb). In this evaluation, only the BAM and PC algorithms demonstrated AUC values exceeding 0.5 in certain instances. However, the AUC metrics remain low. This underscores the utility of undirected graph methods for high-dimensional (d > M) problems for such problems, as directed approaches may not only be inefficient but also risk yielding misleading interpretations.
G.5 Additional results regarding SID metrics
We performed additional evaluations using the structural intervention distance (SID) metric for CPDAG estimation and graph dimensions of d 10, as shown in Figure 12 For higher graph dimensions, the implementation of the SID metric within the causal discovery toolbox was leading to computational errors due to instability and high complexity of the evaluation metrie. The minimal and maximal SID values, representing the best and worst DAGs within the equivalence classes, are shown. Our method has the best performance for the best DAG of the estimated equivalence class. However, the results were less favorable for the worst DAG configurations. These findings suggest that our model estimates a broad equivalence class and tends to be conservative the directionality of edges.
G.6 Results with varying graph distributions
We investigate the robustness of our method to different graph distributions by testing various distributions in combination with different dependency functions. Specifically, we consider the Watts- Strogatz (WS) graph [69,, the Exponential Random Graph Model (ERGM) [56), the configuration model [4), and the geometric random graph [51] for generating graph structures. We report the area under the precision recall curve (AUC) values for the undirected graph estimation task for the case
M = 200, d = 20.
The WS graph produces small-world graphs with high clustering and short average path lengths by randomly rewiring edges of a regular ring lattice. We use the implementation in the NetworkX package [22] create WS graphs with specified expected degree. The rewiring probability p is sampled uniformly from 0.1 to 0.9 to cover a wide range of graph structures.
ERGM generates random graphs capturing various structural properties based probability distribution defined by sufficient statistics. We use the ERGM implementation in NetworkX to generate random DAGs with specified expected degree, iteratively adding edges that maintain acyclicity.
The configuration model 47 generates random graphs with a specified degree sequence. We use two variations of the configuration model in our experiments:
1. In the first variation (conf), the degree sequence is set to a constant value equal to the expected degree for all nodes, resulting in graphs with homogeneous degree distribution.
2. In the second variation (conf2), we use a variable degree sequence sampled from a Poisson distribution with a mean of qa= where d is the number of nodes and q is the expected degree. This allows for generating graphs with heterogeneous degree distributions while maintaining the expected average degree. We ensure that the sum of the sampled degree sequence is even by resampling until this condition is met.
The geometric random graph model [51] generates graphs based on the spatial proximity of nodes in metric space. Nodes are randomly placed in a unit square, and edges are formed between pairs of nodes if their Euclidean distance is below a specified threshold radius. We calculate the radius using
the formula for the radius r V q where d is the number of nodes and q the expected degree.
This choice of radius ensures that the expected degree is approximately achieved in the resulting graph.
To focus on the impact of graph distributions, we keep the expected degree fixed for evaluation. Comparing results across different graph generations, we find no significant difference in performance, suggesting robustness to variations in graph structure.
26

---

PC
ocor
LINGAM
SAM
Auici SCK-VO
zero BAM
HSIC
rcit
GES
MMPC
CAM
DAGGNN
AviCi Clebbyy
05
a
(g)
(h)
(j)
Figure 12: Structural Intervention Distance (SID) results for CPDAG estimation, arranged by SID
values from highest (left) to lowest (right) for dimension d = 10. Panels (a), (b), (c), (g), and (h) present the minimal SID values, while panels (d), (e), (f), (i), and (i) display the maximal SID values within their respective Markov equivalence classes. The panels correspond to different sample sizes: M 50 for (a) and (d); M 100 for (b) and (c); M 200 for (c) and (f); M 500 for (g) and (i); and M 1000 for (h) and (j).
G.7 Time comparison
depicts the average runtimes per evaluation step, accompanied by their corresponding standard deviations. We compared BAM with various unsupervised methods, noting that the eval- uation times for other supervised approaches, such as Avici, are comparable to those observed for BAM. Specifically, the results are presented in the form of mean + one standard deviation. The x-axis enumerates various sample sizes, denoted as M, while both mean and standard deviation were computed based on 5 independent inference tests for each configuration with fixed sample size M and graph dimension d.
Figure 13
27

---

These empirical observations substantiate the computational efficiency of supervised approaches in the inference phase.
cartime
rcot
50-50
Vo- 100
Mo203
N-5m
No300
14-50
100
Me.20
M-503
Mo 1010
ccdr
GES
(b)
NHPC
glasso
CAM
DAGGNN
rundire
We500
5-1030
(d)
Figure 13: Algorithm runtime for (a) d = 10, (b) d = 20, (c) d 50, (d) 100 in seconds per
M x d data matrix inference.
28

---

H
Computational complexity analysis
H.1 Memory complexity
The overall memory complexity of our proposed method is O(C'd + Cd Md + M + C²). where M is the number of datapoints. d the data dimension, and the number of channels. This complexity arises from the following components: O(CMd + d),,, covariance matrix calculation OC + d )
O(CMd + Md
attention-betwenn netween-ddtappinng bilinear attention OCC²)), and matrix logarithm O(Cd²). These complexities are derived from the matrix shapes used in Figure Additionally. C' x C' weight matrices are used, which require ( memory.
H.2 Runtime complexity
The overall runtime complexity of our approach is M(C CM Oad Cd.). We present a detailed breakdown of the single components:
We first consider the computational complexity of the tensor multiplication defined in I a tensor A E IMMx MxxCC is multiplied by a matrix B E RCXC This operation involves multiplying each of the M slices E R of A by B. The multiplication of d C' matrix by C x C' matrix has a time complexity of OCC2d). Since this multiplication is performed in parallel over the M axis, the total complexity of this tensor multiplication is O Md). Similarly, the parallel computation of the attention matrix K @ Q for peaettion--htweee-attrire
(defined in (4). matrix A in Figure 3 left) consists of M parallel matrix multiplications of d x C with C' matrices, resulting complexity of O(CMd²). Multiplying the attention matrix with the values in (5) is an M-parallel computation of M d X d tensor with a M x d x C tensor, having complexity OCMMd). Thus, the overall complexity of attention-between netween-atribuues attributes is O Md + C. By switching axes, attention-between-
has time complexity of OC Md CMdd).
Calculating C' covariance matrices has a complexity of O(CM M, d)). In a BAM layer, multiplying the d x d x C tensor with C. C weight matrix has complexity of OCC²d)). The bilinear operation in (2) (for calculating the attention matrix and the output of the BAM layer) is defined C-parallel computation of two d x d matrix multiplications, resulting in complexity of O(Cd)). Computing the custom softmax can be expressed as d x d matrix multiplications in C channels, so its complexity is OCCd)). Calculating the matrix logarithm equivalent computing an eigendecomposition which has a complexity of O(Cd)) when performed over C' channels 20). In summary, the overall time complexity of our approach is MC CMd +CM2 of &Z +Cd)).
H.3 Computational efficiency discussion
While our method demonstrates strong performance across various settings, the computational resources required for optimal training (approximately 82 GiB of GPU memory) may present limitations in resource-constrained environments. This is primarily due to the attention mechanisms in our architecture, which exhibit quadratic complexity with respect to sequence length.
For datasets with large sample sizes M, the attention- between-datapoinns layer becomes the primary computational bottleneck. with memory complexity ((M2d) and runtime complexity OCCM2d) To address this limitation and enhance scalability, we propose a local attention strategy: dividing the sample axis M ms into subsamples of size and applying attention within these segments. This modification reduces the runtime complexity to OCCm² ds) O CMfd) ) and memory complexity to sdm By maintaining fixed subsample size m, the approach achieves línear scalability in M through linear increments in Similar principles can be applied to both the coantiinn-betweer datapoints and bilinear attention mechanisms.
Furthermore, as demonstrated in Section G.7 our trained model offers significant practical advantages through fast inference times compared to traditional unsupervised approaches, which require iterative optimization for each new dataset.
29

---

I Interpretation
L1 Shape-agnostic architecture and the role of attention layers
When employing shape-agnostin architecture for matrices RMxx, it is crucial to ensure that all elements within the M x matrix can interact and influence one another. Consider a scenario where one axis of the matrix is expanded to the shape R xxx11 followed by dense layers with 1 C and C' C'CC C weights. In this configuration, the dense layers carry out element-wise operations on the M x elements, processing them in isolation from each other. This is because each hidden representation is essentially a linear combination of C' matrices of shape M x d prior to the application of an element-wise activation function.
This limitation is addressed by incorporating attention layers into the architecture. These layers adaptively compute non-tree trainable M x M and d x d attention matrices based on trainable C' x C weights. This approach allows for a permutation- and shape-agnoste architecture, as the same set of trainable weights can be employed for any M X d matrix, while still enabling the matrix entries to influence each other. In this way, the attention mechanism becomes an essential component of our model. Although also experimented with Settwork-i--eewooeee methods [38], we found that the attention mechanism offers a more stable, efficient, and straightforward computation of non-trainable M x M and d attention matrices, using only trainable C x C' weights.
Our proposed bilinear attention mechanism is, to our knowledge, the first SPD layer to enable shape- agnostic computations. It uses trainable C' x C' weights to calculate non-trainable d x d attention matrices, allowing for adaptive weighting across different d SPD matrix sizes. This flexibility makes it unique and essential component of architecture. Additionally, this construction ensures the desired permutation invariance among the variables. Our approach essentially learns matrix operations that should be applicable to any input matrix with arbitrary M x d input.
1.2 Attention scores in the BAM layer
Consider the setting as in Figure 3 (right) and the computation of the output H by A @ S, where A SdxxxC are the attention scores and S E &exxax are the input matrices into the BAM layer. consider for simplicity the output of
Since A S is processed parallel across the channels, d to be quadratic, positive definite d x d single channel here and assume A Sdxd and S saxd matrices.
While traditional self-attention computes scores to assess the importance of one data point another, our bilinear attention mechanism extends this by exploring the interdependence of variable pairs. Specifically, for an output pair j), its associated output value determined not merely by direct scalar relationship but by the bilinear form: Σx,1 As, Thus, instead of singular focus on the relation "How does affect i?", quantified in the score matrix A in classical attention, the score matrix in bilinear attention shows the interaction strengths of pair sets {(i,k) k == I..... .d and ((,5)) |[ = 1,....d) d). The "receptive field" adopts a cross-form within the scores A instead of being Ar.j only, in the sense that relevant scores for the output at position (i, are not limited to A but {A,, U J.
1.3 Keys and queries
Continuing with the single-channel assumption due to parallel channel processing, consider the quadratic form gaxd Saxd + Sdxd, (K,Q) K" QK of the key-query interaction. The (ih entry of KTT QK is KT TKK for the columns K i..... K of K, which are often referred to keys. Using the eigendecomposition of Q Uº DU one obtains for the (i,j)-th entry the bilinear form (UK,)* D(UK,), Note that UK is similarity measure between K and Q analogous to standard attention. So, for bilinear attention, similarity scores are calculated between the keys K and the eigenvectors of the queries Q. Afterwards, the D ---ieihhted bilinear-form (D²UK,) KT) (D²UK, is used to create covariance matrices by combining the similarity scores between U and K. Hence, in bilinear attention, the similarity scores are functions of both the i-th and j-th keys as well as all queries. This is consistent with the attention-score behavior, where the
30

---

interaction strengths of all pair sets k(),k) k = ......d d) and ((.5) 1= 1,... I.....d] df collectively influence the output.
This is in contrast to standard attention, which uses the untransformed dot product (k: q) E IRd x
I + k" q E R for the columns of key and query matrices K, Q.
31

---

Limitations
The proposed model effectively captures smooth dependence relations using Chebyshev polynomials. While this approach demonstrates excellent performance across various types of dependencies, it may face challenges when dealing with data structures that significantly deviate from the generated synthetic data, particularly those with rapidly increasing higher-order derivatives. However, adapting the synthetic data generation process to accommodate these structures is straightforward and can mitigate this limitation.
The Log-Fig layer performs efficiently within moderate dimensional range, but may encounter computational constraints when scaling to higher dimensions. Training our model with parameters similar to those used in this study requires substantial memory resources; in our experiments, approximately 80 GB of GPU memory was necessary.
The attention mechanism, although effective, be computationally expensive for high-dimensional inputs (both in terms M and d). Local attention can alleviate this issue, but it may introduce its own set of challenges. Despite these limitations, the model's architecture allows for easy extensions to address specific applications. For instance, incorporating separate embedding layers for observational and interventional data could enable the model to leverage interventional data effectively.
However, it is important to note that the high computational cost is primarily associated with the training phase, and inference is much faster. While our experiments utilized substantial computational resources (approximately 80 GB GPU memory), the model architecture allows for effective scaling to more modest computational environments. By reducing key parameters such as the maximum number of samples (M), maximum data dimension (d), and number of channels (C), training can be successfully conducted on standard computers with suitable GPU support, though potentially with reduced performance in large dataset regimes.
An end-to-end approach for CPDAG estimation could potentially offer further benefits. However, the current model loses directional information during the covariance computation, making an end-to-end approach unfeasible with the existing architecture. possible extension facilitate directional inference could involve using two separate embeddings for each variable, one for the variable as parent and another for the variable as a child, effectively doubling the dimensionality to 2d.
Another limitation of our approach is the difficulty in theoretically proving the identifiability of causal effects using our neural network method. Identifiability of causal effect, when having access to the data distribution via a neural network, relies on the fundamental approximation theorem of neural networks, which guarantees that a manglehiddeen-aye continuous function on a compact subset of the input space, given sufficient hidden units (27). such that any mapping X + from the data distribution to graph structures can be learned. However, the theorem might not be directly applicable to the observational attention layers and the bilinear attention layers employed in our approach. Furthermore, for identifiability be assured, the test data must be sufficiently close to the training data for out-of out-of-sample approximation, which adds another layer of complexity to the problem.
neural network can approximate any
32

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We introduced novel neural network model for supervised graph structure learning. Our empirical evaluation demonstrates ability to detect the presence of causal relationships, regardless of their specific parameterizations.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
Limitations are discussed in the conclusion and detail in Appendix J
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g.. independence assumptions, noiseless settings, model sellsspeciiiationn asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex, system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
33

---

Answer:
[Yes]
Justification:
The proof of Theorem | provided, while Proposition 2 is obvious.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
The paper provides detailed information for reproducing the main experimental results, including the parameterization the SEM (Appendix E) training details (Appendix
Justification:
F), and the model architecture (Section and Appendices BD
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, other means that appropriate to the research performed.
also be provided via detailed
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reprodueibility may be tricky in some cases, in which authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing or verifying the results.
5. Open access to data and code
34

---

Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: provide source code in the supplementary material, together with a conda environment file that specifies the required dependencies and their versions and README file with instructions to run the code.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.c cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
wwwp://wweeeeeeeeeeeeeeeee http://oooSSSSisseeeeeeeeeeee wwriiisssiiipcyooioo for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Hyperparameters stated in the main text in the "Training data hyperpa- rameter settings" paragraph within the Experiments section. Additionally, a full list of hyperparameters is provided in Appendix ET
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can provided either with the code, appendix, or as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: All results are reported as mean + standard deviation, indicating the statistical significance of the experiments. We discuss the error bars in Appendix G.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emoor bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
35

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative еπτог rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: The paper provides information on the computational resources required in the "Efficiency" paragraph of the Experiments section in the main text. Additional details available in Appendix G.7
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported the paper (e.g., preliminary or failed experiments that didn't 't make into the paper).
9. Code Of Ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics
Answer: [Yes]
Justification:
No ethical concerns or deviations from the NeurlPS Code of Ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: Broader impacts were discussed in the conclusion.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
36

---

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
The paper poses no such risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: All existing assets used in this work are properly cited and credited.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
37

---

If assets are released, the license, copyright information, and terms of use in the thcode. has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
package should be provided. For popular datasets, paperswi
latasets.paperssi
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information not available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification:
The provided code is documented and explained in the paper.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
38

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
39


Reviews:
[{"review_id": "hCRXO1xW4R", "reviewer": "Reviewer_wqVr", "summary": "This paper proposes a novel neural network model for supervised graph structure inference. The model aims to learn the mapping between observational data and their underlying dependence structure using a bilinear attention mechanism (BAM). The BAM operates on the level of covariance matrices of transformed data and respects the geometry of the manifold of symmetric positive definite (SPD) matrices. The paper demonstrates the robustness of the proposed method in detecting dependencies and inferring graph structures.", "questions": "1. Provide a more detailed analysis of the theoretical computational complexity and practical scaling ability of the proposed method. \n\n2. Discuss potential strategies for optimizing the computational efficiency and handling larger datasets effectively.", "limitations": "N/A", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The introduction of the bilinear attention mechanism for processing dependency information in the SPD space is novel and well-motivated. This approach leverages the geometric properties of SPD matrices, offering a fresh perspective on graph structure inference.\n\n2. The paper provides extensive experimental results, showing that the proposed method outperforms state-of-the-art algorithms in various scenarios.", "weaknesses": "Although the paper claims that the proposed method is computationally efficient compared to some unsupervised approaches, the actual computational cost, especially during training, is unclear.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer wqVr,\n\nThank you for the insightful and constructive feedback. We greatly appreciate your recognition of the potential of our approach.\n\n**Efficiency:**\nWe acknowledge that the level of computational resources required for our method (approximately 82 GiB, as stated in lines 309-312) may be a limiting factor in cases where such resources are not readily available. This is indeed an advantage of traditional causal learning algorithms. Deep learning models, particularly attention-based architectures, which exhibit quadratic complexity with respect to sequence length, require significant computational resources to achieve optimal performance. However, when coupled with large computational resources, deep learning has demonstrated superior performance across various domains, such as natural language processing (e.g., GPT-4) and computer vision (e.g., state-of-the-art object detection and segmentation models).\n\n\n**Memory complexity:**\nThe overall memory complexity of our proposed method is $O(CMd + Cd^2 + Md^2 + M^2d+C^2)$, where $M$ is the number of datapoints, $d$ is the data dimension, and $C$ is the number of channels. This complexity arises from the following components: attention-between-attributes $O(CMd+Md^2)$, attention-between-datapoints $O(CMd+M^2 d)$, covariance matrix calculation $O(C(Md+d^2))$, bilinear attention $O(Cd^2)$, and matrix logarithm $O(Cd^2)$. These complexities are derived from the matrix shapes used in Figure 3 of the paper. Additionally, $C\\times C$ weight matrices are used, which require $O(C^2)$ memory. We thank the reviewer for drawing our attention to the incorrect memory complexity given in line 311, which will be corrected in the final version of the paper.\n\n\n**Computational time complexity:** \nWe thank the reviewer for the suggestion to also add an analysis of the computational complexity. The runtime complexity of our approach is $O(C^2 M d+CMd^2+CM^2 d +C^2 d^2 + Cd^3)$. We present a detailed breakdown of the single components:\n\nWe first consider the computational complexity of the tensor multiplication defined in lines 127-128: a tensor $\\boldsymbol{A}\\in\\mathbb{R}^{M\\times d \\times C}$ is multiplied by a matrix $\\boldsymbol{B}\\in \\mathbb{R}^{C\\times C}$. This operation involves multiplying each of the $M$ slices $\\in\\mathbb{R}^{d\\times C}$ of $\\boldsymbol{A}$ by $\\boldsymbol{B}$. The multiplication of a $d\\times C$ matrix by a $C\\times C$ matrix has a time complexity of $O(C^2 d)$. Since this multiplication is performed in parallel over the M axis, the total complexity of this tensor multiplication is $O(C^2 M d)$.\nSimilarly, the parallel computation of the attention matrix $\\boldsymbol{K}\\odot\\boldsymbol{Q}$ for attention-between-attributes (line 498, matrix $\\boldsymbol{A}$ in Figure 3 left) consists of $M$ parallel matrix multiplications of $d\\times C$ with $C\\times d$ matrices, resulting in a complexity of $O(CMd^2)$. Multiplying the attention matrix with the values (line 503) is an M-parallel computation of a $M\\times d\\times d$ tensor with a $M\\times d \\times C$ tensor, having complexity $O(CMd^2)$. Thus, the overall complexity of attention-between-attributes is $O(C^2 Md+CMd^2)$. By switching axes, attention-between-datapoints has a time complexity of $O(C^2 Md+CM^2d)$.\n\nCalculating $C$ covariance matrices has a complexity of $O(CMd(\\min{M,d}))$. In a BAM layer, multiplying the $d\\times d\\times C$ tensor with a $C\\times C$ weight matrix has a complexity of $O(C^2 d^2)$. The bilinear operation in line 174 (for calculating the attention matrix and the output of the BAM layer) is defined as $C$-parallel computation of two $d\\times d$ matrix multiplications, resulting in a complexity of $O(Cd^3)$. Computing the custom softmax can be expressed as $d\\times d$ matrix multiplications in $C$ channels, so its complexity is $O(Cd^3)$. Calculating the matrix logarithm is equivalent to computing an eigendecomposition, which has a complexity of $O(Cd^3)$ when performed over $C$ channels.\nIn summary, the overall time complexity of our approach is $O(C^2 M d+CMd^2+CM^2 d +C^2 d^2 + Cd^3)$.\n\n**Handling larger datasets:**\nWe acknowledge the importance of reducing computational complexity for large, high-dimensional datasets. We propose using local attention instead of global attention to address those cases.\n\nConsider a dataset with a large sample size $M$. In such cases, the attention-between-datapoints layer becomes a bottleneck due to its $O(M^2 d)$ memory complexity and $O(CM^2 d)$ computational complexity. To alleviate this, we propose dividing the sample axis $M = ms$ randomly into $s$ subsamples, each containing $m$ samples. Applying local attention to these smaller segments reduces the runtime complexity to $O(Cm^2 ds)=O(\\frac{CM^2d}{s})$ and the memory complexity to $O(dm^2s)$ for calculating the attention matrix between samples. By maintaining a fixed subsample size $m$, our approach achieves linear scalability with respect to $M$ by incrementing linearly in $s$. This technique can be similarly applied to the attention-between-datapoints and our bilinear attention methods, enhancing overall efficiency.\n\n\nIn the revised manuscript, we will incorporate the computational complexity analysis and provide a detailed discussion on efficiency in the appendix. Furthermore, we will address the limitations related to computational complexity, including the potential of local attention as a mitigating approach.\n\nBest,\n\nAuthors of Submission15606", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "rtZxu9Gcwd", "reviewer": "Reviewer_rbW5", "summary": "This paper proposes to use supervised causal learning approach to learn causal structure. It takes a dataset as input, and outputs a moral graph. The moral graph is an undirected graph with two types of edges: skeleton edge and moralized edge.\n\nIn technical space:\narchitecture: The approach adopts alternating node-wise and sample-wise attention + covariance in embedded space + bilinear attention via SPD net\ntraining data: ER sampling to generate graphs, Chebyshev polynomial functions as basis to generate mechanism and with Gaussian distribution for noise terms.\nlearning task: The main task is a three-class edge classification to output a moral graph. \n\nexperiment shows good performance by comparing with other SOTA competitors, for both the moral graph prediction and CPDAG prediction.", "questions": "See my comments on the weakness part above.", "limitations": "N/A", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. Supervised causal learning is an interesting area, and it has potentially more impact to the causal discovery community.\n\n2. The learning task respects the identifiability, and the main task is to predict a moral graph, which is new to the best of my knowledge.\n\n3. A novel NN architecture is interesting and could make potential impact to the community.", "weaknesses": "1.\tSuperiority of model architecture: What is the significance and necessity of using semi-definite (SPD) matrix space? How to \"enhancing graph structure inference\" should be explained in more detail. Compared with only using alternating attention, what advantages does BAM have and what scenarios are more suitable? The theoretical discussion and experiments in the current paper are not enough to explain the above problems.\n\n2.\tThe necessity of introducing Moralized edges: Overall I like the idea of representing the undirected graph by the moral graph, which is new and sound. However, given that the paper also has a second prediction model to do orientation from moral graph to yield the final CPDAG, so, what are the advantages compared to learning skeleton and v-structure separately? So, I would like to see more discussion and some empirical evidence to demostrate the advantages of choosing moralized graph instead of skeleton as the key learning task.\n\n3.\tThe motivation for choosing Chebyshev as the mechanism of training data: For example, in \"E.1 Chebyshev polynomials for training\", this Chebyshev polynomial form does not seem to be able to describe complex multivariate coupling effects. In addition, there should be corresponding experiments for the hyperparameter setting of Chebyshev to prove that these training effects are robust or sensitive to this hyperparameter.\n\n4.\tIn the experimental part, a large amount of synthetic data is used to test performance, but the functions (mechanisms) of these synthetic data are relatively simple. More complex or more representative mechanisms should be used for testing to enhance persuasiveness, such as MLP, Random Fourier features, etc.\n\n5.\tIn addition, the experimental part is all tested on synthetic data, and should be tested on real benchmarks, such as Sachs, etc.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer rbW5,\n\nThank you for the insightful and constructive feedback. We sincerely appreciate your recognition of the potential impact of our approach.\n\n**Significance of SPD layers:**\nWe agree that the paper will benefit from a clearer justification of the integration of SPD layers and the Bilinear Attention Mechanism (BAM) in our neural network architecture. We have addressed this in the overall author rebuttal.\n\n**Evidence for the effectiveness of moralized edge estimation:**\nThank you for acknowledging our contribution in distinguishing between skeleton and moralized edges in the undirected edge estimation task. We agree that empirical evaluation would further demonstrate its effectiveness. To address this, we ran additional experiments where we maintained the undirected edge estimation network but modified the second step for estimating immoralities. We trained another neural network to test each possible immorality in the graph, not just those identified by the first neural network, using the same hyperparameters for training as before. The results, shown in Table 1 of the rebuttal PDF (at the bottom of the author rebuttal), indicate that our approach of only testing immoralities detected by the first neural network consistently outperformed the case where all possible moralizations were tested.\n\nWe observed that the neural network testing all possible moralizations had a strong bias towards estimating no moralization, likely because this occurs more frequently in sparse graphs used for training. While this bias could potentially be addressed using a weighted loss function, it highlights the efficiency of our approach in avoiding these issues. Information about conditional dependencies might be easier to detect \"globally\" (i.e., having information about the whole graph, not only the Markov blanket estimated beforehand), akin to the Gaussian case, where this information is contained in the precision matrix.\n\n**Multidimensional couplings:**\nWe appreciate the reviewer's suggestion to test our approach on more complex and representative mechanisms. We agree that this would enhance the persuasiveness of our experimental results. To address this point, we conducted additional experiments, as shown in Figure 1 of the rebuttal PDF. In these experiments, we tested our approach for detecting Random Fourier feature dependencies.\n\nFor the undirected case, shown in (A, B, C) of Figure 1, our algorithm outperforms the baselines for all sample sizes except the two highest, $M=500, 1000$, where our method still produces very good results. We believe that for this sample size, our approach might overfit the training data, becoming more confident in not estimating an edge when it behaves differently from what was seen during training. \n\nRegarding the directed estimation results shown in (D, E, F), our algorithm still exhibits reasonable performance, although it is not the best when compared to the baselines. This highlights opportunities for improvement for multidimensional dependencies, which we will discuss in the limitations section.\nAdditionally we have already tested our approach on MLP dependencies, as shown in Appendix G.2 of our paper, where we demonstrate competitive performance of our approach.\n\n**Rationale behind using Chebyshev polynomials:**\n\nChebyshev polynomials exhibit factorially decreasing coefficients when approximating functions with bounded derivatives of any order. By sampling the coefficients from a uniform distribution with factorial decay, we can generate coupling functions that are characterized by being smooth and well-behaved with a large degree of variation. This property ensures that the generated coupling functions are diverse and well-behaved, without being dominated by high-frequency oscillations or other unstable behavior.\n\nThe orthogonality of Chebyshev polynomials is a desirable property for our application, as it allows us to easily evaluate the magnitude of individual effects and ensures that the effects do not interfere with each other in an unstable manner. This is particularly important to control the magnitude of the causal effects to ensure that an edge in the adjacency matrix corresponds to a meaningful, controllable effect in the generated data. Training with disrupted training data, where an edge is to be predicted without being predictable is causing problems for learning.\n\nIn contrast, using e.g. a randomly initialized MLP as coupling may not provide the same level of control and interpretability. While the universal approximation theorem suggests that MLPs can represent a wide range of functions, the expressive power of a randomly initialized MLP is limited. The weights are typically drawn from simple distributions like Gaussian or uniform, which can lead to similar patterns in the network's behavior across different initializations. As a result, the simulated relationships may be opaque and limited expressive.\n\nOur experiments demonstrate that Chebyshev polynomials can approximate a wide range of simple and complex dependencies, including random MLP and RFF dependencies. However, we acknowledge our algorithm's limitations in detecting multidimensional dependencies, particularly in the RFF dependency case for directed edge estimation.\n\n**Evaluation on real benchmarks:**\nWe acknowledge the suggestion to test on real-world benchmark datasets. While we recognize the importance of testing on real data, due to time constraints and the scope of this submission, we focused on synthetic data to validate our method. \n\nFor real-world data, the true graph structure is often uncertain, and there are many random effects, making principled evaluation challenging. Synthetic data allows for a more controlled evaluation of detection accuracy. Nevertheless, we agree that real-world data are crucial and will be the next step in our future studies.\n\nWe will incorporate all points discussed above in the final version of our paper.\n\nBest,\n\nAuthors of Submission15606", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "dBFp4xEBLe", "reviewer": "Reviewer_2mp1", "summary": "This paper studies the problem of graph structure inference using a neural network with bilinear attention mechanism.", "questions": "See Weaknesses.", "limitations": "Limitations have been adequately discussed.", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The proposed framework is novel and can operate in Euclidean and positive semi-definite matrix spaces. Attention mechanism is utilized in an innovative manner to reveal interdependencies. Experiments demonstrate that the proposed framework achieves competent and robust performance.", "weaknesses": "1. Could you clarify whether BAM has any theoretical guarantees on the recovery of correct structure? In the absence of such guarantees, it is not clear how BAM might perform in settings beyond those explored in the paper.\n\n2. The advantages of BAM over the state-of-the-art approaches could be explained better from a conceptual perspective.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer 2mp1,\n\n\nThank you for recognizing the novelty and innovation of our framework. We appreciate the constructive and insightful feedback, on clarifying the advantages of our BAM layer in the SPD space and the theoretical guarantees of our approach.\n\n**Significance and advantages of SPD layers:**\nWe agree that the paper will benefit from a clearer justification of the integration of SPD layers and the Bilinear Attention Mechanism (BAM) in our neural network architecture, highlighting their advantages. We have addressed this in the overall author rebuttal.\n\n**Theoretical guarantees:**\nThank you for raising the important question of whether our BAM framework provides theoretical guarantees on the recovery of the correct graph structure. We acknowledge that the current paper does not include formal proofs of identifiability or consistency. We have addressed the challenges in theoretically proving identifiability and consistency in the Limitations section. The universal approximation theorem for neural networks provides some assurance that, in principle, a sufficiently expressive neural network can learn any continuous mapping from the data distribution to graph structures. However, several assumptions need to be fulfilled to ensure convergence to such a solution. \n\nMoreover, the theorem pertains to performance on the training data, and generalization to out-of-sample data is a different challenge. Rigorously extending these results to our specific architecture with observational attention and bilinear attention layers is non-trivial. \n\nThat said, we believe our extensive empirical evaluation provides compelling evidence that BAM is able to reliably recover meaningful graph structures in practice. Across a range of settings, BAM achieved strong performance competitive with or exceeding state-of-the-art baselines. This situation, where empirical performance outpaces theoretical understanding, has been a common phenomenon in the deep learning field. Many deep learning approaches have achieved state-of-the-art results in various domains despite limited theoretical foundations. While we absolutely agree that developing a rigorous theoretical understanding is crucial and the lack thereof is a clear limitation, we also believe that research in deep learning and its applications to new domains should not entirely be avoided due to the absence of a complete theoretical framework. \n\nWe will incorporate a more detailed discussion on the advantages of SPD layers and the BAM architecture, as well as an extended analysis of the limitations regarding theoretical guarantees, in the revised manuscript and appendix, following your suggestions.\n\nBest,\n\nAuthors of Submission15606", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "rbDxqhmAWC", "reviewer": "Reviewer_wDfW", "summary": "The paper presents a novel neural network model designed for supervised graph structure inference. The core contribution is the introduction of a Bilinear Attention Mechanism (BAM) which processes dependency information at the level of covariance matrices of transformed data. This approach respects the geometry of the manifold of symmetric positive definite matrices, leading to robust and accurate inference of both undirected graphs and partially directed acyclic graphs (CPDAGs). The model is trained on variably shaped and coupled simulated input data and requires only a single forward pass for inference. Empirical evaluations demonstrate the model's robustness in detecting a wide range of dependencies, outperforming several state-of-the-art methods.", "questions": "1. How does the performance of the model scale with larger and more complex graph structures?\n2. Please summarize the research motivation and major contributions of this work. \n3. Can the authors provide more details on the computational requirements and potential optimizations for the training process?\n4. Considering the high computational cost and significant memory resources required during the training phase (e.g., approximately 80 GB of GPU memory), what strategies or architectural modifications do you plan to implement to mitigate these computational demands without compromising the model's performance in the future?", "limitations": "The authors have adequately discussed the limitations and potential negative societal impact of their work. However, more detailed discussions on the scalability and computational requirements could be beneficial. Constructive suggestions for improvement include exploring optimizations for the training process and providing additional real-world case studies to validate the method's applicability.", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The introduction of the Bilinear Attention Mechanism (BAM) is novel and aims to address key challenges in graph structure inference.\n- The empirical evaluations are comprehensive, demonstrating the robustness and effectiveness of the proposed method.\n- The paper is well-written. The figures depicted contribute to clearly demonstrating the pipeline proposed.\n- The method provides a significant improvement over existing approaches.", "weaknesses": "- The method's complexity, particularly in terms of the bilinear attention mechanism and its application, may limit its accessibility to a broader audience.\n- The training process, requiring extensive computational resources, may pose a challenge for researchers with limited access to high-performance computing facilities.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer wDfW,\n\nThank you for the insightful and constructive feedback. We sincerely appreciate your recognition of the novelty and effectiveness of our Bilinear Attention Mechanism (BAM) for graph structure inference. \n\n**Research motivation:** We agree that the paper will benefit from a clearer explanation of the research motivation, particularly regarding why we introduce the BAM layer in the SPD space for graph structure inference. We have addressed this in the overall author rebuttal.\n\n**Major contributions:** Our primary contribution lies in the introduction of the novel Bilinear Attention Mechanism (BAM) layer. Additionally, we propose the innovative application of geometric deep learning methods, inspired by the computer vision field, in the SPD space for supervised causal learning. Moreover, we introduce a novel learning task that respects identifiability and distinguishes between moralized and skeleton edges, enabling more precise causal structure inference. In the final version of our paper, we will enhance our presentation to ensure that these essential contributions are more readily accessible to the reader.\n\n**Computational complexity and resource requirements**\nWe acknowledge the potential limitations of our method in terms of its computational complexity and the requirement for extensive computational resources during training, which is a common challenge in deep learning, particularly when compared to the vast computational architectures employed by large companies in fields such as computer vision and GPT models.\n\nNevertheless, we would like to highlight that our approach can be effectively scaled down to accommodate more modest computational resources. By reducing model and training generator parameters such as the maximum number of samples ($\\overline{M}=500$), maximum data dimension ($\\overline{d}=50$), and number of channels ($C=50$), we have successfully conducted training on a standard MacBook equipped with an M1 chip and 16 GB memory. We utilized Apple Metal to facilitate GPU-like operations on the M1 chip. This approach is also feasible on computers equipped with a suitable NVIDIA graphics card, offering an alternative for efficient training without the need for a dedicated computing server.\nAlthough this scaled-down version may not achieve the same level of performance as the high-resource setup for large dataset regimes with large $M$ and $d$ values, it showcases the flexibility and adaptability of our approach to various computational constraints.\n\n**Memory complexity:**\nThe overall memory complexity of our proposed method is $O(CMd + Cd^2 + Md^2 + M^2d+C^2)$, where $M$ is the number of datapoints, $d$ is the data dimension, and $C$ is the number of channels. This complexity arises from the following components: attention-between-attributes $O(CMd+Md^2)$, attention-between-datapoints $O(CMd+M^2 d)$, covariance matrix calculation $O(C(Md+d^2))$, bilinear attention $O(Cd^2)$, and matrix logarithm $O(Cd^2)$. These complexities are derived from the matrix shapes used in Figure 3 of the paper. Additionally, $C\\times C$ weight matrices are used, which require $O(C^2)$ memory. We thank the reviewer for drawing our attention to the incorrect memory complexity given in line 311, which will be corrected in the final version of the paper.\n\n\n**Computational time complexity:** \nWe thank the reviewer for the suggestion to also add an analysis of the computational complexity. The runtime complexity of our approach is $O(C^2 M d+CMd^2+CM^2 d +C^2 d^2 + Cd^3)$. We present a detailed breakdown of the single components:\n\nThe complexity of attention-between-attributes is $O(C^2 Md+CMd^2)$. By switching axes, attention-between-datapoints has a time complexity of $O(C^2 Md+CM^2d)$.\n\nIn a BAM layer, multiplying the $d\\times d\\times C$ tensor with a $C\\times C$ weight matrix has a complexity of $O(C^2 d^2)$. The bilinear operation in line 174 (for calculating the attention matrix and the output of the BAM layer) is defined as $C$-parallel computation of two $d\\times d$ matrix multiplications, resulting in a complexity of $O(Cd^3)$. Computing the custom softmax can be expressed as $d\\times d$ matrix multiplications in $C$ channels, so its complexity is $O(Cd^3)$. Calculating the matrix logarithm is equivalent to computing an eigendecomposition, which has a complexity of $O(Cd^3)$ when performed over $C$ channels.\n\n\n**Handling larger datasets:**\nWe acknowledge the importance of reducing computational complexity for large, high-dimensional datasets. We propose using local attention instead of global attention to address those cases.\n\nConsider a dataset with a large sample size $M$. In such cases, the attention-between-datapoints layer becomes a bottleneck due to its $O(M^2 d)$ memory complexity and $O(CM^2 d)$ computational complexity. To alleviate this, we propose dividing the sample axis $M = ms$ randomly into $s$ subsamples, each containing $m$ samples. Applying local attention to these smaller segments reduces the runtime complexity to $O(Cm^2 ds)=O(\\frac{CM^2d}{s})$ and the memory complexity to $O(dm^2s)$ for calculating the attention matrix between samples. By maintaining a fixed subsample size $m$, our approach achieves linear scalability with respect to $M$ by incrementing linearly in $s$. This technique can be similarly applied to the attention-between-datapoints and our bilinear attention methods, enhancing overall efficiency.\n\nIn the revised manuscript, we will incorporate the computational complexity analysis and provide a detailed discussion on efficiency in the appendix. Furthermore, we will address the limitations related to computational complexity, including the potential of local attention as a mitigating approach.\n\nBest,\n\nAuthors of Submission15606", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We appreciate the time and effort the reviewers have invested in evaluating our work. We are grateful for your constructive feedback and insightful suggestions, which have helped us identify areas for improvement. 

We have added a PDF at the end of the author rebuttal, which includes a comparison of the identification of moralizations and evaluations on data using Random Fourier feature dependencies. 

In response to the reviewers' requests for further elaboration on the motivation and rationale behind the advantages of our approach, we provide a more detailed explanation of the advantages of our approach.

**Motivation for Optimizing in the SPD Space for Graph Estimation: GLASSO optimization:**

We acknowledge the importance of clearly justifying the integration of SPD-layers in the neural network architecture.

The use of SPD layers is inspired by the success of traditional GLASSO-based methods for undirected graph estimation, such as the influential works of Banerjee et al. (2008), Friedman et al. (2008), and Yuan and Lin (2007). These methods optimize a penalized likelihood over the SPD space using techniques like block coordinate descent, ensuring that the estimate remains positive definite at each iteration. These approaches have proven to be efficient for learning graph structures when the input SPD matrix encodes sufficient information, particularly in the Gaussian case where the covariance matrix is a sufficient statistic for the distribution (neglecting a shift by the mean). The main reason optimization is performed solely in the SPD space is that inverse covariance matrices naturally represent graphs. Zero entries indicate missing edges and non-zero entries correspond to the strength of conditional dependencies.  This characteristic makes them a powerful tool for estimating well-conditioned and interpretable structures within the SPD space, particularly when combined with $l_1$ regularization to promote sparsity.  

However, for general SEMs with non-linear dependencies and non-Gaussian error terms, pure SPD models cannot capture the full information of the data. Neural networks, on the other hand, excel at learning and decoding complex non-linear relationships. The observational attention layers in our model are trained to nonlinearly transform the data into a batch of $C$ covariance matrices, learning an end-to-end observational transformation that filters the information useful for graph prediction and provides sufficient information for efficiently solving the undirected graph estimation task.

Thus, our BAM-network learns data-driven matrix transformations akin to those used in solving the GLASSO optimization problem. Our network implicitly encodes the solution to a parametric optimization problem. For some objective function $F$ (e.g. penalized negative log-likelihood for GLASSO), we optimize $\min_{\Sigma} F(\Sigma, \mathcal{D})$ with respect to $\Sigma$, where the solution is parametric in the data $\mathcal{D}$, or any sufficient statistics thereof. The solution describes a mapping $\mathcal{D} \rightarrow \hat{\Sigma}$, with $\hat{\Sigma} = \arg\min_{\Sigma} F(\Sigma, \mathcal{D})$ being the space of all optimization solutions for different data $\mathcal{D}$. Hence, instead of solving the SPD GLASSO optimization for each new dataset $\mathcal{D}$ to find the corresponding $\hat{\Sigma}$, the network learns the entire mapping $\hat{\Sigma}(\mathcal{D})$, requiring only a single cheap forward pass through a neural network for evaluation on a specific $\mathcal{D}$. Additionally, a transformation in the observational data space to handle non-Gaussian data is learned at the beginning. It is worth noting that the operations for solving GLASSO problems by a traditional algorithm can be applied to input data of any dimension $d$ and any sample size $M$. Our network also achieves this by applying only attention-based, shape-invariant operations.

Another advantage of our end-to-end learning approach over traditional likelihood-based methods is that it does not require a sensitive sparsity hyperparameter. 

**Advantages of Optimizing in the SPD Space for Graph Estimation: Enhanced expressiveness and natural representation:**

In traditional supervised causal discovery, the adjacency matrix values are directly calculated from a learned $M \times d\times C$ transformed observational data matrices. This requires decoding the $d \times d$ pairwise relationships within each of the $M$ samples that far, that probabilities can be obtained after a few simple operations. Before the output is calculated, the $M$ axis is max-pooled out (Lorch et al., 2022) or cut out (Ke et al., 2022), resulting in a $C\times d$ data representation in both cases. However, extracting $d(d-1)/2$ distinct pairwise interactions from a $C \times d$ representation by one operation on the $C$-axis lacks expressiveness especially when the data dimension $d$ varies. Also, using Max-Pooling for the sample axis might be problematic for varying sample sizes $M$, because the magnitude of the output values scales with $M$ alone. 

In contrast, using a $d \times d$ representation naturally accommodates the varying dimensionality and allows for a more direct and efficient extraction of the pairwise dependencies. We argue that using $C \times d$ representations for inferring $d \times d$ pairwise interactions is structurally misaligned, leading to inefficiencies compared to directly operating on $d \times d$ matrices.

**References**:

Banerjee et al. (2008). Model selection through sparse maximum likelihood estimation for multivariate Gaussian or binary data. JMLR

Friedman et al. (2008). Sparse inverse covariance estimation with the graphical lasso. Biostatistics

Ke et al. (2022). Learning to Induce Causal Structure. In ICLR.

Lorch et al. (2022). Amortized inference for causal structure learning. In NeurIPS

Yuan \& Lin (2007). Model selection and estimation in the Gaussian graphical model. Biometrika

Meta-review: This paper introduces a new neural network model for supervised graph structure inference, featuring a novel Bilinear Attention Mechanism (BAM) that effectively processes dependency information at the level of covariance matrices of transformed data. The proposed method shows promise for learning causal structures, with experimental results sufficient to support its effectiveness. While there were initial concerns regarding several aspects, including the algorithm's complexity, the authors' rebuttal has adequately addressed many of the reviewers' questions. Given these considerations, I believe the paper has sufficient merit to be accepted for this venue.
Decision: Accept (poster)