Graph Structure Inference with BAM: Neural Dependency Processing via Bilinear Attention

Authors: philipp  froehlich
heinz  koeppl
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 10:34:05

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.

Main paper: 
 introduction: 
 The discovery and understanding of dependencies among variables are fundamental to scientific inquiry across diverse disciplines, such as biology [10] , climate science [48] , economics [8] , and social studies [18] . These dependencies are commonly represented as edges within graphical models, particularly directed acyclic graphs (DAGs), first introduced by Wright [70] to 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 constraintbased 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 [31] . Constraint-based methods usually require a large sample size [66] , rely on an elusive optimal threshold hyperparameter , and Shah and Peters [58] proved that the failure of Type I error control in underlying conditional independence tests is unavoidable, which can have significant consequences in downstream analyses. Supervised causal learning techniques, as presented by Lopez-Paz et al. [39, 40, 41] , Li et al. [36] , Ke et al. [31] , Lorch et al. [42] , Dai et 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 information from observational data matrices. Current approaches for processing observational data, whether in causal learning [31, 42] , or related fields [34, 54] , 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 causal learning methods [31, 42] then predict edges separately from each other via compressed observational representation vectors, rather than estimating a single graph object as in graphical lasso (GLASSO) algorithms [72, 7, 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 in the Gaussian case, where the covariance matrix is a 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-derived mathematical operations as in GLASSO-based algorithms, our network learns these shape-invariant transformations from data. Combined with learned transformations in the observational space, our 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.Analogous to GLASSO-based methods, the symmetry of covariance matrices constrains our approach to undirected graph estimation. We extend beyond this limitation via a novel two-phase procedure. We first estimate the graph's skeleton and immoralities, then identify edge directions to obtain a 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, instead of relying on sequential conditional independence tests over 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. d ) is 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. 
method: 
 2.1 Supervised approach for learning graph structuresIn supervised graph learning, simulated training data is generated in the form of input/label pairs (X, G). Each pair consists of an observational data matrix X ∈ R M ×d , where M 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 = {v 1 , . . . , v d } and a set of edges E ⊆ V × V . Alternatively, graphs can be represented as adjacency matrices A ∈ {0, 1} d×d , where A ij = 1 if an edge exists between nodes v i and v j , and A ij = 0 otherwise. The process begins by creating a random graph, followed by generating obserational data matrices using a structural equation model (SEM)X v = f v (X pa G (v) , ϵ v ), where f v is a measurable coupling function of parent nodes pa(G) and zero-mean error ϵ v . The random coupling in the SEM is designed to approximate a broad class of functions, enabling the model to perform classification of the presence or absence of causal relationships, irrespective of their specific functional forms. A 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 is shown in Figure 2 .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 → G from observational data matrices to the corresponding graph structures.For training data generation, we utilize Erdős-Rényi (ER) graphs, denoted as ER(d, 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 ([q, q]), respectively. The shape-invariant design of our model enables it to train without restriction to a fixed (d, q) pair, which proves advantageous when the graph density is unknown. For each graph G i , we generate a data matrix X i ∈ R M ×d using an SEM, where the sample size M is drawn from a discrete uniform distribution, M ∼ U({M , . . . , M }). Although ER graphs are used for training, we 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 realworld functions, showing factorial decay in their coefficients given that higher order derivatives are bounded [63] , i.e., for the n-th coefficient c n it holds |c n | ≤ C n! for a constant C, making higherdegree terms negligible, such that all such smooth functions can be efficiently approximated with a low polynomial degree Chebyshev polynomial. This choice of parameterization avoids unnatural, non-smooth coupling functions where a higher-order derivative is much larger than a lower-order one, and where there is a 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, 53] . Scatterplots of input/output values of four example Chebyshev polynomials are shown on the left side of Figure 2 . The error terms ϵ v in the SEM are modeled using Gaussian mixture models, which provide a flexible and expressive framework for capturing a diverse range of error distributions [55] . Details about the parameterization of the SEM can be found in Appendix E.Our neural network architecture is designed to 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 a DAG or CPDAG in a 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 to 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 or conditional independence tests, in which the conditioning set includes all other variables, to estimate the moral graph, as e.g., in [17, 59] .Using the assumption of faithfulness and the Markov condition enable us to identify the Markov equivalence class [61] and thus uniquely solve the three-class classification problem. Appendix C provides an explicit proof outlining the conditional independence relations that lead to a unique solution.The prediction targets are an extension of the binary adjacency matrix A ∈ {0, 1} d×d to a set of one-hot encoded adjacency matrices, denoted as A ∈ {0, 1} d×d×3 . Here, for each i, j, the vector A i,j,• ∈ {0, 1} 3 represents a one-hot encoded classification among the three classes: skeleton edge, moralized edge, and no edge.We use the multi-dimensional analogue of matrix multiplication: For a tensor A ∈ R I×J×K and a matrix B ∈ R K×L we denoteEQUATIONThe network architecture is depicted in Figure 1 . In the following, we introduce the individual layers.Channel embedding. We perform an embedding of the input X ∈ R M ×d to obtain a hidden representation with C channels. For this, one axis for X is extended to X ∈ R M ×d×1 and then trainable weightsW 1 ∈ R 1×C , W 2 ∈ R C×C are used to obtain H = X + relu XW 1 W 2 ∈ R M ×d×C ,where broadcasting is used for the addition.A shape-and permutation-invariant neural network is achieved by employing C × C weight matrices in conjunction with attention mechanisms for fixed channel dimensions C and C 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 a sequence of length l is embedded into an l × C matrix and processed via self-attention with C × C trainable weights.Observational data self-attention. Self-attention for C observational M × d data matrices, stored in a tensor X ∈ R M ×d×C , proposed by Kossen et al. [34] , is based on axial attention [26, 67] and commonly used for processing observational data [4, 31, 42] . The attention-between-attributes layer (Figure 3 Transition to SPD space via covariance calculation. Let a transformed data matrix X ∈ R M ×d×C be given. To obtain explicit dependency information, we calculate channelwise covariance matrices, i.e.:(d, d, C) W K (C, C) K (d, d, C) W Q (C, C) Q (d, d, C) A (d, d, C) H (d, d, C) X (M, d, C) W K (C, c) K (M, d, c) Q (M, d, c) V (M, d, C) A (M, d, d) H (M, d, C) W Q (C, c) W V (C, C) < l a t e x i t s h a 1 _ b a s e 6 4 = " 7 1 A H A g O 1 F I X 8 u + T Z t u 7 B S f U 2 9 K 4 = " > A A A B / H i c b V A 7 T 8 M w G H T K q 5 R X o C O L R Y X E V C W I 1 1 j B w l g E f U h N V D m O 0 1 p 1 7 M h 2 k K K o / B U W B h B i 5 Y e w 8 W 9 w 2 g z Q c p L l 0 9 3 3 y e c L E k a V d p x v q 7 K y u r a + U d 2 s b W 3 v 7 O 7 Z + w d d J V K J S Q c L J m Q / Q I o w y k l H U 8 1 I P 5 E E x Q E j v W B y U / i 9 R y I V F f x B Z w n x Y z T i N K I Y a S M N 7 b o X C B a q L D Z X 7 t 3 T U Y y m Q 7 v h N J 0 Z 4 D J x S 9 I A J d p D + 8 s L B U 5 j w j V m S K m B 6 y T a z 5 H U F D M y r X m p I g n C E z Q i A 0 M 5 i o n y 8 1 n 4 K T w 2 S g g j I c 3 h G s 7 U 3 x s 5 i l W R z 0 z G S I / V o l e I / 3 m D V E d X f k 5 5 k m r C 8 f y h K G V Q C 1 g 0 A U M q C d Y s M w R h S U 1 W i M d I I q x N X z V T g r vΣ = 1 M -1 (X -µ X ) T (3,2,1) ⊙ (X -µ X ) T (3,1,2) T (2,3,1) ∈ R d×d×C ,where we define for tensors A ∈ R I×J×K , B ∈ R I×K×L the I-parallel matrix multiplication The transpose notation T (i,j,k) permutes the dimensions of a tensor according to the specified order (i, j, k), following the commonly used notation in deep learning [1] . For example, A T (2,3,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.A ⊙ B := (A i,•,• B i,•,• ) i=1,...,I ∈ R I×J×L , and µ X = 1 M 1 T M X ∈ R 1×d×CIntroducing the bilinear attention mechanism. To effectively process dependency information stored in the covariance matrices, which is crucial for causal discovery, we propose a 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-invariant. In contrast to existing SPD architectures [28, 68] that typically parameterize weights W ∈ R d×dout to be applied as W T ΣW ∈ S d×d ⪰ to a matrix Σ ∈ S d×d ⪰ , we parameterize weights W ∈ R C×C to act as linear combinations ΣW ∈ S d×d×C ⪰ on a set of covariance matrices Σ ∈ S d×d×C ⪰ . We employ an attention mechanism to generate attention matrices A ∈ R d×d×C , serving as adaptable d × d weights on each d × d matrix of Σ ∈ R d×d×C . Our bilinear attention mechansim is shown in (Figure 3 , right).The Riemannian manifold of SPD matrices has nonpositive sectional curvature, and its geometric properties are not preserved under Euclidean operations [9] , leading to issues when using traditional neural networks [46] . Specialized architectures respecting the SPD manifold's geometry are necessary. Bilinear 1 matrix multiplication Σ → W T ΣW , analogous to a dense layer in Euclidean space, preserves the space of SPD matrices S d×d ⪰ and serves as a primary tool for SPD layers [68] .For an input Σ ∈ S d×d×C ⪰ , we obtain keys K = ΣW K ∈ S d×d×C ⪰ and queries Q = ΣW Q ∈ S d×d×C ⪰ , using non-negativity constraints on the weightsW K , W Q ∈ R C×C + .Keys and queries are now combined in a bilinear fashion, parallel over the C channels by calculatingEQUATIONresulting in a tensor of C SPD matrices.We replace the softmax function for traditional attention with a custom softmax function σ to preserve positive definiteness, i.e., Using the definition of the bilinear parallel tensor product in 2,we obtain the attention matrix A by applying σ channelwise, i.e.,A := σ((K ⊗ Q) •,•,c ) c=1,...,C ∈ S d×d×C ⪰ ,and we obtain the output of the bilinear layer asH = A ⊗ Σ ∈ S d×d×C ⪰ .The elementwise exponential function preserves positive definiteness because SPD matrices are closed under addition and the Hadamard product [9] , and exp[S] = ∞ n=0 1 n! [S] n , with elementwise exponentiation [•] n .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 Euclidean space through the Log-Eig layer, as proposed by [28] . This transformation, given an input matrix S = U DU T via eigendecomposition can be expressed asl : S d×d ⪰ → R d×d , l(S) := log(S) := U log(D)U T .The final layer, equipped with a softmax activation function, consists of C = 3 output units for generating the probabilities.Interpretation. A detailed discussion on the necessity, attention scores, and keys and queries of our novel attention mechanism is available in Appendix I. Unlike traditional attention where the (i, j)-th score in the attention matrix A indicates the influence of element i on j, our bilinear attention mechanism reveals interdependences. For an output pair (i, j), its associated output value is determined not merely by a direct scalar relationship but by the bilinear form k,l A i,k S k,l A l,j . Consequently, the influence on the (i, j)-th entry in the output is based on A i,• ∪ A •,j rather than just A i,j , leading to attention scores forming a cross shape within the matrix A. This approach allows columns of the matrix S to attend to each other, highlighting a more complex interaction pattern. 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 a 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 to 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 is further refined using Meek rules [45] . Details about the CPDAG estimation step can be found in Appendix D. 
related work: 
 Causal discovery primarily relies on unsupervised learning methods, including constraint-based approaches that infer conditional independencies [61, 16] and score-based methods that optimize a score function under acyclicity constraints [13, 71] . 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 a 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-occurring residue types.Our work closely aligns with Ke et al. [31] and Lorch et al. [42] , 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 [42] utilize the dot product of embedding vectors derived from max-pooling 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 [31] , our model can be trained across different shapes of sample and attribute numbers, eliminating the need for re-training on 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, [31] and [42] deduce a 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-invariant models for supervised causal inference, but their method directly computes correlation matrices without learning a highdimensional 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. [39, 40] proposed a 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 a network that does not require re-training for new evaluation samples. [15] 's work on immoralities aligns with our second CPDAG estimation step. 
experiments: 
 Training data hyperparameter settings. The hyperparameters for sampling training data, which are defined in section 2.2, are set to d = 10, d = 100, q = 1, q = min( d 3 , 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] , rcot and rcit [62] , ccdr [3] , GES [13] , GIES [23] , LiNGAM [60] , MMPC [64] , CAM [11] , 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 a supervised causal discovery baseline, we use the AVICI model [42] , both a pre-trained version (SCM-v0) and a version trained from scratch on the same Chebyshev data as our method, using default hyperparameters and data dimensions d ∈ {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 is computed for undirected prediction. We also compare SHD and accuracy results against a naive zero graph without edges. Performance indicators. We assess the performance of both undirected and CPDAG graph estimations using the Area Under the Precision-Recall Curve (AUC), which is well-suited for imbalanced binary classification tasks like sparse graph detection [24] . For CPDAG estimations, we also employ the Structural Hamming Distance (SHD), a standard metric in structure learning [71, 31] . In the undirected graph estimation task, SHD is equivalent to 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. [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 4 (A -C )). It excels in capturing intricate non-monotonic dependencies (Figure 4 (A, B, E ) and demonstrates superior performance over Avici on the training set (Figure 4 (D, F )).Figure 5 displays the results for CPDAG estimation tasks. In high-dimensional scenarios (d = 100, M = 50), no algorithm surpassed the baseline of a 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 moderatedimensional (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-trained Avici, illustrating BAM's competitive performance in learning CPDAG-structure. Specifically, 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 demonstrate that our architecture generalizes beyond the Erdős-Rényi graphs it was trained on, we evaluate its performance on graphs generated using various random graph models. Our results, shown in Figure 6 , demonstrate consistent performance across all graph distributions, suggesting that our method is robust to variations in underlying graph structure. Details about the configuration of the distributions can be found in Appendix G.6.Ablation study. Table 1 presents our ablation study results, underscoring the crucial role of the bilinear layer, whose removal significantly increases loss metrics. Omitting the LogEig layer leads to a substantial deterioration in loss, indicating that direct predictions from the SPD-space are problematic. Notably, bilinear data processing alone yields relatively good results, possibly due to the embedding layer's ability to decode non-linearities in the data, which are lost when solely relying on covariance matrices. Further details are provided in Appendix F.2.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 2 shows that our selective testing strategy consistently outperforms 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. O(CM d + Cd 2 + M d 2 + M 2 d + C 2 ). The runtime complexity of our approach is O(C 2 M d + CM d 2 + CM 2 d + C 2 d 2 + Cd 3). We present a 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, as shown in Appendix G.7.Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)? Answer: [Yes] Justification: The paper provides detailed information for reproducing the main experimental results, including the parameterization of the SEM (Appendix E), training details (Appendix F), and the model architecture (Section 2 and Appendices B,D). 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 a 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 a novel architecture, describing the architecture fully might suffice, or if the contribution is 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 to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed. • While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. , with an open-source dataset or instructions for how to construct the dataset). (d) We recognize that reproducibility may be tricky in some cases, in which case 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 to have some path to reproducing or verifying the results. 
conclusion: 
 In this study, we introduced a 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 a 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. Comprehensive empirical evaluations demonstrate that our approach consistently outperforms state-of-the-art methods across diverse scenarios. Future directions include extending our approach to handle more complex settings and investigating its performance under various noise and confounding conditions, as well as scaling through local attention mechanisms. Limitations and future directions are further discussed in Appendix J.• 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. • If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration 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.• 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 a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an 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 a 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 to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?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 a 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 is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?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 or dataset.• The authors should state which version of the asset is used and, if possible, include a 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.• If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • 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 is not available online, the authors are encouraged to reach out to the asset's creators. 
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/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was 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. 
crowdsourcing 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 as 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 in the supplemental material is fine, but if the main contribution 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. 
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 an equivalent approval/review based on the requirements of your country or institution) were obtained?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. • Depending on 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.• We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to 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. 
 
 Appendix: Similar to the standard softmax, our softmax for an SPD matrix returns positive values summing to d, but unlike the standard softmax, the rows do not sum to 1. We demonstrate that our modified softmax function σ regularizes the eigenvalues: Theorem 1. For any S ∈ S d×d ⪰ , the largest eigenvalue of σ(S) is 1.Proof. Let S := exp[S]. By similarity transformation, the eigenvalues of σ(S) are equal to the eigenvalues of SΛ(S). It holds that SΛ(S) S1 = S1, which demonstrates that S1 is the Perron eigenvector corresponding to the eigenvalue 1. The assertion now follows from the Perron-Frobenius theorem.Proposition 2. The custom softmax σ is invariant to additive shifting, i.e., σ(S + α) = σ(S) for each α ∈ R, S ∈ S d×d ⪰ .Proposition 2 shows that σ, unlike standard softmax, does not need a scaling constant, and it can easily manage exploding exp values via maximum-value scaling.This layer accepts an input X ∈ R M ×d×C , from which it generates keysK = XW K ∈ R M ×d×c , queries Q = XW Q ∈ R M ×d×c , and values V = XW V ∈ R M ×d×C , with W K , W Q ∈ R C×c , and W V ∈ R C×C .For each m = 1, . . . , M , keys and queries are combined in parallel along the inner axis, leading toEQUATIONThis results in the attention weightsA = σ K ⊙ Q T (1,3,2) √ c ∈ R M ×d×d ,where T (perm) 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 is computed asEQUATIONTo 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 on the work of [57] :Theorem 3. Any continuous function from C([-1, 1], [-1, 1]) acting elementwise on a matrix preserves positive definiteness if it can be expressed as a series f :[-1, 1] → [-1, 1] with f (x) = ∞ k=1 x k w k , subject to ∞ k=1 w k ≤ 1 with w k ≥ 0, ∀k ∈ N.We employ this theorem to construct an activation function for the SPD neural network, using a relatively small maximal polynomial degree value of N max = 3. This function employs trainable weights w k and is applied after 'correlation normalization', i.e., the conversion of 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, [2] provides various types of trainable activation functions for Euclidean neural networks.Normalization is critical in attention networks, with various methodologies available [25, 5] . We adopt normalization and a residual connection [25] on the input of all attention layers, expressed as S + Attention(LayerNorm(S)). Particularly when M < d, 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] .Within the SPD manifold, we employ correlation normalization alongside residual connections. This approach preserves positive definiteness and corresponds intuitively to standard normalization in Euclidean space.In all attention layers, we employ multihead attention, a process that divides the input tensor along the channel axis into several smaller tensors. Each of these is then subjected to attention independently.For initialization of the C × C weight matrices W + in each bilinear attention layer, we draw initial weights from U 0, 2 C . Given a tensor of SPD matrices Σ ∈ S d×d×C ⪰ , in expectation, the diagonal entries in each of the d × d matrices in ΣW + should stay within a comparable range, while offdiagonal 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.Building upon the foundational principles of Markov and faithfulness [33] , we 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 an directed edge from node X to Y in the DAG, then X and Y are dependent given any set from the power set of the other nodes, i.e.:X → Y =⇒ X ⊥ ⊥ ̸ Y | C ∀C ∈ P(V \ {X, Y }).For an immorality between X and Y , there exists a set of nodes C ∈ P(V \ {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 X and Y , or the set of all parents of X or Y ), but dependent given all other nodes in the graph, i.e.:X → Z ← Y, X ← - Y, X → - Y =⇒ ∃C ∈ P(V \ {X, Y }) : X ⊥ ⊥ Y | C, X ⊥ ⊥ ̸ Y | V \ {X, Y }.If there is no edge between X and Y , and if X and Y do not have a common child, then X and Y are conditionally independent given all other nodes, i.e.,X ← - Y, X → - Y, ∄Z : X → Z ← Y =⇒ X ⊥ ⊥ Y | V \ {X, Y }.Note that testing for the no-edge class is cost-effective, as one only needs to test for a single set, rather than checking for anyZ ∈ V \ {X, Y } if there is a v-structure X → Z ← Y .Testing for a 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.To estimate a CPDAG from the graph skeleton, along with the set of immoralities between pairs of nodes, we test each estimated immorality to determine which potential common child nodes are indeed common children. The parent nodes, denoted by pa, are the nodes between which an immorality was first estimated. Potential common child nodes, denoted by cc, are nodes that have an edge to both parent nodes. Neighbor nodes, denoted by ne, are nodes that have one edge to exactly one parent.We take the columns x pa ∈ R M ×2 , x cc ∈ R M ×|cc| , x ne ∈ R M ×|ne| of the data matrix X corresponding to pa, cc, ne as inputs for the neural network. Each of these submatrices undergoes a dimensionality expansion tox pa ∈ R M ×2×1 , x cc ∈ R M ×|cc|×1 , x ne ∈ R M ×|ne|×1. Next, three separate feed-forward subnetworks l pa , l cc , l ne are applied to the three inputs x pa , x cc , x ne . Each of these layers has the same architecture: For an input x ∈ R M ×d×1 , weight matrices W 1 ∈ R 1×C and W 1 ∈ R C×C , together with a bias vector b 1 ∈ R C are used to embed x toh 1 = tanh( xW 1 + b 1 ) W 1 ∈ R M ×d×CThen, two residual layers of the formh i+1 = h i + tanh(h i W i + b i ) W i with W i ∈ R C×C , W i ∈ R C×C , b i ∈ R C , i = 1, 2 are applied to obtain representations h pa ∈ R M ×2×C , h cc ∈ R M ×|cc|×C , h ne ∈ R M ×|ne|×C . For the addition of the bias terms, broadcasting is used, i.e., b ∈ R C → b ∈ R M ×d×C with b m,l,c = b c m = 1, . . . , M, l = 1 . . . , d, c = 1, . . . , C.Now, the representations are concatenated along the dimension axis to obtain a M ×(2+|cc|+|ne|)×C tensor. We use d = 2 + |cc| + |ne|. We employ the same observation-to-dependency network as before, but instead of using softmax on the output of the LogEig-Layer, we use mean-poolingX → 1 d X T (3,2,1) 1 d Tto inflate one of the variable axes of the R d×d×C output of LogEig to obtain a R d×C batch of C vectors of dimension d. After applying a dense C × 1 layer, we obtain a vector of length d. Now, the entries corresponding to the potential common children can be sliced out and backpropagated for training.The network architecture is shown in figure 7 . For the Chebyshev polynomial, we utilize the following parameterization:f v (x pa , ϵ 1 , ϵ 2 ) = w∈pa G (v) β w r n=1 α n T n (x w ) + ϵ 1 + α m   s,t∈pa(G),s<t δ s,t T m (x s , x t ) + w∈pa(G) T m (x w , ϵ 2 )   ∀v ∈ V,where r = 5 is the degree, and T n denotes the Chebyshev polynomials of the first kind (scaled for the input to have a maximum absolute value of 1), and bivariate polynomials are given by:EQUATIONwhere m in the index stands for "multidimensional".Here,µ x ∼ U [-1, 1], µ y ∼ U [-1, 1], and the coefficients α 1 , . . . , α r , α m are calculated fromα i = αi n αn+ αm for i = 1, . . . , r, with α i = γi i! , γ i ∼ U [-1, 1], and α m ∼ U [-1, 1]. β w = βw |pa| with β w ∼ U [0.7, 1.3]. δ s,t are random weights with δ s,t = δs,t s,t∈pa(G),s<t | δs,t| , δ s,t ∼ U [-1, 1]This parameterization is motivated by the observation that for a smooth function-where higherorder derivatives are not significantly larger than the lower-order ones-the coefficients of the Chebyshev approximation decrease in a factorial manner. This provides a rationale for training the neural network on 'typical' smooth functional dependencies. Furthermore, this suggests that using Chebyshev polynomials of degree r = 5 is not a significant limitation, as the coefficients of higher orders are already negligibly small.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 of our model, we introduce thresholds for any absolute values exceeding 5, thereby mitigating the potential impact of outliers.The additive error term ϵ 1 follows a Gaussian mixture distribution. We randomly determine the number of components L ∼ U {1, . . . , 5} from a discrete uniform distribution. Each component has randomly assigned parameters for the meansµ l ∼ U [-1, 1], standard deviations σ l ∼ U [0.05, 1],and weightsw l ∼ U [0.3, 1], w l = w l l w l such that ϵ 1 ∼ L l=1 w l N (µ l , σ 2 l ). The multiplicative error term ϵ 2 is uniformly distributed, with ϵ 2 ∼ U [-1, 1]. We scale the error to zero mean and variance 1 after simulation.To create testing data, we create synthetic data according to an SEM equipped with different dependency function, while the error term follows a Gaussian mixture distribution as before. We use the following dependencies for testing: Chebyshev, linear, sine, cosine, x 2 , x 3 , multidimensional multiplicative dependency. The Chebyshev-dependency used was the same as in the training procedure. We used the following testing dependency functions: x, sin(x), cos(x), x 2 , x 3 as g(x) inf test (x pa , ε) = w∈pa(G) α w g(x w ) + ϵ.For the multi-dimensional multiplicative test dependency, we usedf v (x pa , ϵ 1 , ϵ 2 ) = α m   s,t∈pa(G),s<t δ s,t T m (x s , x t ) + w∈pa(G) T m (x w , ϵ 2 )   ∀v ∈ V, with T m (x, y) := (x -µ x )(y -µ y ) (1 + |µ x |)(1 + |µ y |) with µ x ∼ U [-1, 1], µ y ∼ U [-1, 1], δ s,t = δs,t s,t∈pa(G),s<t | δs,t| , δ s,t ∼ U [-1, 1].The implementation was performed using TensorFlow [1] . We employed the ADAM optimizer by Kingma and Ba [32] . For training, we used the hyperparameters stated in Table 3 In the ablation study, we evaluate the performance of the full model in comparison to 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 = c = 100 channels and are trained over 500 epochs, with each epoch comprising 128 data matrix / adjacency label pairs. Again, we generated data with a random number of samples M ∼ U {50, 51, . . . , 1000} and a random variable dimension d ∼ U {10, 11, . . . , 100}.We employ the categorical cross-entropy loss function for classifying edges into one of three categories: no-edge, skeleton edge, and moralized edge. Additionally, to enforce the condition that a moral edge between nodes X and Y should only be predicted if there is a potential common child Z (i.e., X -Z and Z -Y ), we introduce a penalty term, L p .The overall loss function is defined asL b + L c + L p .Here, L c is the categorical crossentropy of the three categories given byL c (A, A) := H(A, A) = - d i=2 i-1 j=1 3 c=1 A i,j,c log( A i,j,c )denotes the categorical crossentropy of the three categories no-edge, skeleton edge, and moralized edge withA ∈ R d×d×3 with A i,j,c = 1 if (i, j) is in category c in the ground-truth DAG 0 else and A i,j,c denotes the estimation by the algorithm on it. A and A are symmetric along its first two axes, i.e., A i,j,c = A j,i,c , i, j = 1, . . . , d, c = 1, . . . , 3.L b denotes the binary loss of no-edge vs. any edge present (present edges = skeleton edges ∪ moralized edges):L b (A (b) , A (b) ) := H(A (b) , A (b) ) = - d i=2 i j=1 A (b) i,j log( A (b) i,j ) + (1 -A (b) i,j ) log(1 -A (b) i,j )withA (b) ∈ R d×d with A (b) i,j =1 if an edge is estimated between i and j 0 else being the adjacency matrix of no-edge vs. (direct edge ∪ moralized edge).The penalty term, L p is defined as:L p ( A) := max A 3 -[ A 2 A 2 ] 0.5 , 0 . A 1 , A 2 , A 3 ∈ R d×d are estimates of A •,•,1 , A •,•,2, and A •,•,3 by the algorithm respectively. This term penalizes the prediction of a moralized edge in the absence of potential common child edges.The square root operation is applied element-wise.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 note that the observed variability, manifested as relatively large error bars, is predominantly due to the random sampling of graph degrees, which has a substantial influence on estimation accuracy. Despite this inherent variability, the comparison across algorithms remains valid, as 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. 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 of layers ∼ U{1, . . . , 5}, a random number of hidden layers ∼ U{4, 64}, and relu or tanh activation with probability 0.5 each. Despite these complexities, our algorithm maintained state-of-the-art performance, delivering AUC scores competitive to the top-performing existing methods such as depicted in Figure 9 .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 10 A-C), 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 the training patterns. For CPDAG estimation (Figure 10 D-F), while our method shows reasonable performance as measured by AUC and SID, there remains room for improvement in handling multidimensional dependencies. These results complement our findings on MLP dependencies (Figure 9 ), indicating that our approach can effectively handle various types of non-linear relationships while highlighting specific scenarios for future enhancements. CPDAG estimation in high-dimensional settings presents significant challenges. In the specific case of d = 100 and M = 50, none of the algorithms we evaluated could outperform a zero-graph (i.e., a graph with no edges) baseline in terms of Structural Hamming Distance (SHD). The results, depicted in Figure 11 (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 11 (b). 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.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 metric. 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 in the directionality of edges.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 [47] , 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 caseM = 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] to create WS graphs with a 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 on a probability distribution defined by sufficient statistics. We use the ERGM implementation in NetworkX to generate random DAGs with a 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 a homogeneous degree distribution.2. In the second variation (conf2), we use a variable degree sequence sampled from a Poisson distribution with a mean of q d d-1 , 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 a 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 = q π(d-1) , where d is the number of nodes and q is 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. F k O m v i g 4 P F e F V X 1 g k R w b V z 3 2 1 l Z X V v f 2 M x t 5 b d 3 d v f 2 C w e H D R 2 n i r I 6 j U W s W g F q J r h k d c O N Y K 1 E M Y w C w Z r B 8 H b i N x + Z 0 j y W D 2 a U M D / C v u Q h p 2 i s d F / C s 2 6 h 6 J b d K c g y 8 e a k C H N U u 4 W v T i + m a c S k o Q K 1 b n t u Y v w M l e F U s H G + k 2 q W I B 1 i n 7 U t l R g x 7 W f T U 8 f k 1 C o 9 E s b K l j R k q v 6 e y D D S e h Q F t j N C M 9 C L 3 k T 8 z 2 u n J r z x M y 6 T 1 D B J Z 4 v C V B A T k 8 n f p M c V o 0 a M L E G q u L 2 V 0 A E q p Ma m k 7 c h e I s v L 5 P G e d m 7 K l / W L o q V 2 t M s j h w c w w m U w I N r q M A d V K E O F P r w D K / w 5 g j n x X l 3 P m a t K 8 4 8 w i P 4 A + f z B 7 Z u j e M = < / l a t e x i t > (a)< l a t e x i t s h a 1 _ b a s e 6 4 = " X Z D 2 B v x g z q H k f 9 J 4 B Y h H 0 U x 0 u U k = " > A A A B 6 n i c b V D J S g N B E K 1 x j X G L e v T S G I R 4 C T P i d g x 4 8 Z i g W S A Z Q k + n J m n S 0 z N 0 9 w h h C P g D X j w o 4 t U v 8 u b f 2 F k O m v i g 4 P F e F V X 1 g k R w b V z 3 2 1 l Z X V v f 2 M x t 5 b d 3 d v f 2 C w e H D R 2 n i m G d x S J W r Y B q F F x i 3 X A j s J U o p F E g s B k M b y d + 8 x G V 5 r F 8 M K M E / Y j 2 J Q 8 5 o 8 Z K 9 6 X g r F s o u m V 3 C r J M v D k p w h z V b u G r 0 4 t Z G q E 0 T F C t 2 5 6 b G D + j y n A m c J z v p B o T y o a 0 j 2 1 L J Y 1 Q + 9 n 0 1 D E 5 t U q P h L G y J Q 2 Z q r 8 n M h p p P Y o C 2 x l R M 9 C L 3 k T 8 z 2 u n J r z x M y 6 T 1 K B k s 0 V h K o i J y e R v 0 u M K m R E j S y h T 3 N 5 K 2 I A q y o x N J 2 9 D 8 B Z f X i a N 8 7 J 3 V b 6 s X R Q r t a d Z H D k 4 h h M o g Q f X U I E 7 q E I d G P T h G V 7 h z R H O i / P u f M x a VW c Q = " > A A A B 6 n i c b V D L S g N B E O y J r x h f U Y 9 e B o M Q L 2 F X f B 2 D X j x G N A 9 I l j A 7 O 5 s M m Z 1 d Z m a F s O Q T v H h Q x K t f 5 M 2 / c Z L s Q R M L G o q q b r q 7 / E R w b R z n G x V W V t f W N 4 q b p a 3 t n d 2 9 8 v 5 B S 8 e p o q x J Y x G r j k 8 0 E 1 y y p u F G s E 6 i G I l 8 w d r + 6 H b q t 5 + Y 0 j y W j 2 a c M C 8 i A 8 l D T o m x 0 k M 1 O O 2 X K 0 7 N m Q E v E z c n F c j R 6 J e / e k F M 0 4 h J Q w X R u u s 6 i f E y o g y n g k 1 K v V S z h N A R G b C u p Z J E T H v Z 7 N Q J P r F K g M N Y 2 Z I G z 9 T f E x m J t B 5 H v u 2 M i B n q R W 8 q / u d 1 U x N e e x m X S W q Y p P N F Y S q w i f H 0 b x x w x a g R Y 0 s I V d z e i u m Q K E K N T a d k Q 3 A X XV D L S g N B E O y N r x h f U Y 9 e B o M Q L 2 F X f B 2 D X j x G N A 9 I l j A 7 6 S R D Z m e X m V k h L P k E L x 4 U 8 e o X e f N v n C R 7 0 M S C h q K q m + 6 u I B Z c G 9 f 9 d n I r q 2 v r G / n N w t b 2 z u 5 e c f + g o a N E M a y z S E S q F V C N g k u s G 2 4 E t m K F N A w E N o P R 7 d R v P q H S P J K P Z h y j H 9 K B 5 H 3 O q L H S Q x l P u 8 W S W 3 F n I M v E y 0 g J M t S 6 x a 9 O L 2 J J i N I w Q b V u e 2 5 s / J Q q w 5 n A S a G T a I w p G 9 E B t i 2 V N E T t p 7 N T J + T E K j 3 S j 5 Q t a c h M / T 2 R 0 l D r c R j Y z p C a o V 7 0 p u J / X j s x / W s / 5 T J O D E o 2 X 9 R P B D E R m f 5 N e l w h M 2 J s C W W K 2 1 s J G 1 J F m b H p F G w I 3 u L L y 6 R x V v E u K x f 3 5 6 X q T R Z H H o 7 g G M r g w R V U 4 Q 5 q U A c G A 3 i G V 3 h z h P P i v D s f 8 9 a c k 8 0 c w h 8 4 n z + R / 4 1 Y < / l a t e x i t > (e)< l a t e x i t s h a 1 _ b a s e 6 4 = " 5 q D / s r l d S i a s j 6 P 3 a n P E w o s J H s E = " > A A A B 6 n i c b V D L S g N B E O y N r x h f U Y 9 e B o M Q L 2 F X f B 2 D X j x G N A 9 I Q p i d 9 C Z D Z m e X m V k h L P k E L x 4 U 8 e o X e f N v n C R 7 0 M S C h q K q m + 4 u P x Z c G 9 f 9 d n I r q 2 v r G / n N w t b 2 z u 5 e c f + g o a N E M a y z S E S q 5 V O N g k u s G 2 4 E t m K F N P Q F N v 3 R 7 d R v P q H S P J K P Z h x j N 6 Q D y Q P O q L H S Q z k 4 7 R V L b s W d g S w T L y M l y F D r F b 8 6 / Y g l I U r D B N W 6 7 b m x 6 a Z U G c 4 E T g q d R G N M 2 Y g O s G 2 p p C H q b j o 7 d U J O r N I n Q a R s S U N m 6 u + J l I Z a j 0 P f d o b U D P W i N x X / 8 9 q J C a 6 7 K Z d x Y l C y + a I g E c R E Z P o 3 6 X O F z I i x J Z Q p b m 8 l b E g V Z c a m U 7 A h e I s v L 5 P G W c W 7 r F z c n 5 e q N 1 k c e T i C Y y i D B 1 d Q h T u o Q R 0 Y D O A Z X u H N E c 6 L 8 + 5 8 z F t z T j Z z C H / g f P 4 A k 4 S N W Q = = < / l a t e x i t > (f ) < l a t e x i t s h a 1 _ b a s e 6 4 = " E e j U 6 b 1 G Z S A W z 4 M y N z W C i Q g 5 S l I = " > A A A B 6 n i c b V D L S g N B E O y J r x h f U Y 9 e B o M Q L 2 F X f B 2 D X j x G N A 9 I l j A 7 m d 0 M m Z 1 d Z m a F s O Q T v H h Q x K t f 5 M 2 / c Z L s Q R M L G o q q b r q 7 / E R w b R z n G x V W V t f W NX i Z u T C u R o 9 M t f v U F M 0 4 h J Q w X R u u s 6 i f E y o g y n g k 1 K v V S z h N A R C V n X U k k i p r 1 s d u o E n 1 h l g I N Y 2 Z I G z 9 T f E x m J t B 5 H v u 2 M i B n q R W 8 q / u d 1 U x N c e x m X S W q Y p P N F Q S q w i f H 0 b z z g i l E j x p Y Q q r i 9 F d M h U Y Q a m 0 7 J h u A u vU 4 V 0 1 g 7 O p X N w O 3 4 i a B f F 1 0 k = " > A A A B 6 n i c b V D L S g N B E O y N r x h f U Y 9 e B o M Q L 2 F X f B 2 D X j x G N A 9 I l j A 7 m S R D Z m e X m V 4 h L P k E L x 4 U 8 e o X e f N v n C R 7 0 M S C h q K q m + 6 u I J b C o O t + O 7 m V 1 b X 1 j f x m Y W t 7 Z 3 e v u H / Q M F G i G a + z S E a 6 F V D D p V C 8 j g I l b 8 W a 0 z C Q v B m M b q d + 8 4 l r I y L 1 i O O Y + y E d K N E X j K K V H s r i t F s s u R V 3 B r J M v I y U I E O t W / z q 9 C K W h F w h k 9 S Y t u f G 6 K d U o 2 C S T w q d x P C Y s h E d 8 L a l i o b c + O n s 1 A k 5 s U q P 9 C N t S y G Z q b 8 n U h o a M w 4 D 2 x l S H J p F b y r + 5 7 U T 7 F / 7 q V B x g l y x + a J + I g l G Z P o 3 6 Q n N G c q x J Z R p Y W 8 l b E g 1 Z W j T K d g Q v M W X l 0 n j r O J d V i 7 u z 0 v V m y y O P B z B M Z T B g y u o w h 3 U o A 4 M B v A M r / D m S Of F e X c + 5 q 0 5 J 5 s 5 h D 9 w P n 8 A m B O N X A = = < / l a t e x i t > (i) < l a t e x i t s h a 1 _ b a s e 6 4 = " x 9 j L s P d K d + i j 7 V j R r Z c t P c a Z W o k = " > A A A B 6 n i c b V D L S g N B E O y N r x h f U Y 9 e B o M Q L 2 F X f B 2 D X j x G N A 9 I l j A 7 6 S R j Z m e X m V k h L P k E L x 4 U 8 e o X e f N v n C R 7 0 G h B Q 1 H V T X d X E A u u j e t + O b m l 5 Z X V t f x 6 Y W N z a 3 u n u L v X 0 F G i G N Z Z J C L V C q h G w S X W D T c C W 7 F C G g Y C m 8 H o e u o 3 H 1 F p H s l 7 M 4 7 R D + l A 8 j 5 n 1 F j p r v x w 3 C 2 W 3 I o 7 A / l L v I y U I E O t W / z s 9 C K W h C g N E 1 T r t u f G x k + p M p w J n B Q 6 i c a Y s h E d Y N t S S U P U f j o 7 d U K O r N I j / U j Z k o b M 1 J 8 T K Q 2 1 H o e B 7 Q y p G e p F b y r + 5 7 U T 0 7 / 0 U y 7 j x K B k 8 0 X 9 R B A T k e n f p M c V M i P G l l C m u L 2 V s C F V l B m b T s G G 4 C 2 + / J c 0 T i r e e e X s 9 r R U v c r i y M M B H E I Z P L i A K t x A D e r A Y A B P 8 A K v j n C e n T f n f d 6 a c 7 K Z f f g F 5 + M b m Z i N X Q = = < / l a t e x i t > (j)Figure 13 depicts the average runtimes per evaluation step, accompanied by their corresponding standard deviations. We compared BAM with various unsupervised methods, noting that the evaluation 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 a fixed sample size M and graph dimension d.These empirical observations substantiate the computational efficiency of supervised approaches in the inference phase. The overall memory complexity of our proposed method isO(CM d + Cd 2 + M d 2 + M 2 d + 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(CM d + M d 2 ), attention-between-datapoints O(CM d + M 2 d), covariance matrix calculation O(C(M d + 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 . Additionally, C × C weight matrices are used, which require O(C 2 ) memory.The overall runtime complexity of our approach is O(C 2 M d + CM d 2 + CM 2 d + C 2 d 2 + Cd 3 ).We present a detailed breakdown of the single components:We first consider the computational complexity of the tensor multiplication defined in (1): a tensor A ∈ R M ×d×C is multiplied by a matrix B ∈ R C×C . This operation involves multiplying each of the M slices ∈ R d×C of A by B. The multiplication of a d × C matrix by a C × 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). Similarly, the parallel computation of the attention matrix K ⊙ Q for attention-between-attributes (defined in (4), matrix A in Figure 3 ). Computing the custom softmax can be expressed as d × 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 [20] . In summary, the overall time complexity of our approach is O(C 2 M d+CM d 2 +CM 2 d+C 2 d 2 +Cd 3 ).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-datapoints layer becomes the primary computational bottleneck, with memory complexity O(M 2 d) and runtime complexity O(CM 2 d).To address this limitation and enhance scalability, we propose a local attention strategy: dividing the sample axis M = ms into s subsamples of size m and applying attention within these segments. This modification reduces the runtime complexity to O(Cm 2 ds) = O( CM 2 d s ) and memory complexity to O(dm 2 s). By maintaining fixed subsample size m, the approach achieves linear scalability in M through linear increments in s. Similar principles can be applied to both the attention-betweendatapoints 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.I.1 Shape-agnostic architecture and the role of attention layers When employing a shape-agnostic architecture for matrices ∈ R M ×d , it is crucial to ensure that all elements within the M × d matrix can interact and influence one another. Consider a scenario where one axis of the matrix is expanded to the shape R M ×d×1 , followed by dense layers with 1 × C and C ×C weights. In this configuration, the dense layers carry out element-wise operations on the M ×d 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 × 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-trainable M × M and d × d attention matrices based on trainable C × C weights. This approach allows for a permutation-and shape-agnostic architecture, as the same set of trainable weights can be employed for any M × 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 we also experimented with Network-in-Network methods [38] , we found that the attention mechanism offers a more stable, efficient, and straightforward computation of non-trainable M × M and d × d attention matrices, using only trainable C × C weights.Our proposed bilinear attention mechanism is, to our knowledge, the first SPD layer to enable shapeagnostic computations. It uses trainable C × C weights to calculate non-trainable d × d attention matrices, allowing for adaptive weighting across different d × d SPD matrix sizes. This flexibility makes it a unique and essential component of our 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 × d input.Consider the setting as in Figure 3 While traditional self-attention computes scores to assess the importance of one data point to another, our bilinear attention mechanism extends this by exploring the interdependence of variable pairs. Specifically, for an output pair (i, j), its associated output value is determined not merely by a direct scalar relationship but by the bilinear form: k,l A i,k S k,l A l,j . Thus, instead of a singular focus on the relation "How does j 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 = 1, . . . , d} and {(l, j) | l = 1, . . . , d}. The "receptive field" adopts a cross-form within the scores A instead of being A i,j only, in the sense that relevant scores for the output at position (i, j) are not limited to A i,j but {A i,• ∪ A •,j }.Continuing with the single-channel assumption due to parallel channel processing, consider the quadratic form S d×d × S d×d → S d×d , (K, Q) → K T QK of the key-query interaction. The (i, j)-th entry of K T QK is K T i QK j for the columns K 1 , . . . , K d of K, which are often referred to as keys. Using the eigendecomposition of Q = U T DU one obtains for the (i, j)-th entry the bilinear form (U K i ) T D(U K j ). Note that U K i is a 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 1 2 -weighted bilinear-form (D1 2 U K i ) T (D 1 2 U K j) 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 interaction strengths of all pair sets {(i, k) | k = 1, . . . , d} and {(l, j) | l = 1, . . . , d} collectively influence the output. This is in contrast to standard attention, which uses the untransformed dot product (k, q) ∈ R d × R d → k T q ∈ R for the columns of key and query matrices K, Q.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-Eig layer performs efficiently within a moderate dimensional range, but it 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, can be computationally expensive for high-dimensional inputs (both in terms of 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 of 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. A possible extension to facilitate directional inference could involve using two separate embeddings for each variable, one for the variable as a 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 a 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 single-hidden-layer neural network can approximate any continuous function on a compact subset of the input space, given sufficient hidden units [27] , such that any mapping X → G 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 to be assured, the test data must be sufficiently close to the training data for out-of-sample approximation, which adds another layer of complexity to the problem.Answer: [Yes] Justification: The proof of Theorem 1 is 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 crossreferenced. • 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 to 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.The mapping A → AΣA T is quadratic and hence nonlinear in A. It is often referred to as bilinear as it is a special case of the mapping (A, B) → AΣB T . 


Decision: 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.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: 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.
questions: questions: 1. Provide a more detailed analysis of the theoretical computational complexity and practical scaling ability of the proposed method. 

2. Discuss potential strategies for optimizing the computational efficiency and handling larger datasets effectively.
limitations: limitations: N/A
paper_summary: paper_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.

In technical space:
architecture: The approach adopts alternating node-wise and sample-wise attention + covariance in embedded space + bilinear attention via SPD net
training data: ER sampling to generate graphs, Chebyshev polynomial functions as basis to generate mechanism and with Gaussian distribution for noise terms.
learning task: The main task is a three-class edge classification to output a moral graph. 

experiment shows good performance by comparing with other SOTA competitors, for both the moral graph prediction and CPDAG prediction.
strength_weakness: strength_weakness: 1.	Superiority 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.

2.	The 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.

3.	The 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.

4.	In 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.

5.	In addition, the experimental part is all tested on synthetic data, and should be tested on real benchmarks, such as Sachs, etc.
questions: questions: See my comments on the weakness part above.
limitations: limitations: N/A
paper_summary: paper_summary: This paper studies the problem of graph structure inference using a neural network with bilinear attention mechanism.
strength_weakness: strength_weakness: 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.

2. The advantages of BAM over the state-of-the-art approaches could be explained better from a conceptual perspective.
questions: questions: See Weaknesses.
limitations: limitations: Limitations have been adequately discussed.
paper_summary: paper_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.
strength_weakness: strength_weakness: - The method's complexity, particularly in terms of the bilinear attention mechanism and its application, may limit its accessibility to a broader audience.
- The training process, requiring extensive computational resources, may pose a challenge for researchers with limited access to high-performance computing facilities.
questions: questions: 1. How does the performance of the model scale with larger and more complex graph structures?
2. Please summarize the research motivation and major contributions of this work. 
3. Can the authors provide more details on the computational requirements and potential optimizations for the training process?
4. 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: 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.

Comments: comment: Thank you for your response and your positive assessment.
Thank you for your detailed response. It effectively addresses my questions. I will maintain my acceptance rating.
Thank you for considering our rebuttal and for raising your score.
I thank the authors for the rebuttal. After considering it and other reviews, I have raised my score.