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

Dear Reviewer wDfW,

Thank 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.

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.

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.

Computational complexity and resource requirements We 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.

Nevertheless, 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. Although 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.

Memory complexity: The 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.

Computational time complexity: We 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:

The 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)$.

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.

Handling larger datasets: We 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.

Consider 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.

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

Best,

Authors of Submission15606

Dear Reviewer 2mp1,

Thank 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.

Significance and advantages of SPD layers: We 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.

Theoretical guarantees: Thank 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.

Moreover, 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.

That 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.

We 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.

Best,

Authors of Submission15606

Dear Reviewer rbW5,

Thank you for the insightful and constructive feedback. We sincerely appreciate your recognition of the potential impact of our approach.

Significance of SPD layers: We 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.

Evidence for the effectiveness of moralized edge estimation: Thank 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.

We 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.

Multidimensional couplings: We 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.

For 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.

Regarding 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. Additionally 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.

Rationale behind using Chebyshev polynomials:

Chebyshev 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.

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

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

Our 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.

Evaluation on real benchmarks: We 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.

For 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.

We will incorporate all points discussed above in the final version of our paper.

Best,

Authors of Submission15606

Dear Reviewer wqVr,

Thank you for the insightful and constructive feedback. We greatly appreciate your recognition of the potential of our approach.

Efficiency: We 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).

Memory complexity: The 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.

Computational time complexity: We 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:

We 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)$. Similarly, 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)$.

Calculating $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. In 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)$.

Handling larger datasets: We 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.

Consider 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.

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

Best,

Authors of Submission15606