# Graph Structure Inference with BAM: Introducing the Bilinear Attention Mechanism

arXiv: https://arxiv.org/abs/2402.07735 · HTML: https://mlsanity.com/paper/2402.07735/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 32 of its 76 references that are in the map.
- First posted on arXiv 12 Feb 2024.

## Builds on (most cited first, 30 of 32)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Neural Machine Translation by Jointly Learning to Align and Translate](https://mlsanity.com/paper/1409.0473/) (2014)
- [Group Normalization](https://mlsanity.com/paper/1803.08494/) (2018)
- [Transformers in Vision: A Survey](https://mlsanity.com/paper/2101.01169/) (2021)
- [On Layer Normalization in the Transformer Architecture](https://mlsanity.com/paper/2002.04745/) (2020)
- [Going deeper with Image Transformers](https://mlsanity.com/paper/2103.17239/) (2021)
- [DAGs with NO TEARS: Continuous Optimization for Structure Learning](https://mlsanity.com/paper/1803.01422/) (2018)
- [Network In Network](https://mlsanity.com/paper/1312.4400/) (2013)
- [Axial Attention in Multidimensional Transformers](https://mlsanity.com/paper/1912.12180/) (2019)
- [DAG-GNN: DAG Structure Learning with Graph Neural Networks](https://mlsanity.com/paper/1904.10098/) (2019)
- [AutoInt: Automatic Feature Interaction Learning via Self-Attentive Neural Networks](https://mlsanity.com/paper/1810.11921/) (2018)
- [ReZero is All You Need: Fast Convergence at Large Depth](https://mlsanity.com/paper/2003.04887/) (2020)
- [Differentiable Causal Discovery from Interventional Data](https://mlsanity.com/paper/2007.01754/) (2020)
- [D'ya like DAGs? A Survey on Structure Learning and Causal Discovery](https://mlsanity.com/paper/2103.02582/) (2021)
- [CAM: Causal additive models, high-dimensional order search and penalized regression](https://mlsanity.com/paper/1310.1533/) (2013)
- [Causal Inference and Causal Explanation with Background Knowledge](https://mlsanity.com/paper/1302.4972/) (2013)
- [A Riemannian Network for SPD Matrix Learning](https://mlsanity.com/paper/1608.04233/) (2016)
- [Self-Attention Between Datapoints: Going Beyond Individual Input-Output Pairs in Deep Learning](https://mlsanity.com/paper/2106.02584/) (2021)
- [Amortized Inference for Causal Structure Learning](https://mlsanity.com/paper/2205.12934/) (2022)
- [A survey on modern trainable activation functions](https://mlsanity.com/paper/2005.00817/) (2020)
- [Discovering Causal Signals in Images](https://mlsanity.com/paper/1605.08179/) (2016)
- [Approximate Kernel-based Conditional Independence Tests for Fast Non-Parametric Causal Discovery](https://mlsanity.com/paper/1702.03877/) (2017)
- [Learning to Induce Causal Structure](https://mlsanity.com/paper/2204.04875/) (2022)
- [Structure Learning in Graphical Modeling](https://mlsanity.com/paper/1606.02359/) (2016)
- [Structural Agnostic Modeling: Adversarial Learning of Causal Graphs](https://mlsanity.com/paper/1803.04929/) (2018)
- [Concave Penalized Estimation of Sparse Gaussian Bayesian Networks](https://mlsanity.com/paper/1401.0852/) (2014)
- [Supervised Whole DAG Causal Discovery](https://mlsanity.com/paper/2006.04697/) (2020)
- [Discovering Cyclic Causal Models with Latent Variables: A General SAT-Based Procedure](https://mlsanity.com/paper/1309.6836/) (2013)
- [Parsimonious modeling with Information Filtering Networks](https://mlsanity.com/paper/1602.07349/) (2016)

## Abstract (arXiv)

In statistics and machine learning, detecting dependencies in datasets is a central challenge. We propose a novel neural network model for supervised graph structure learning, i.e., the process of learning a mapping between observational data and their underlying dependence structure. 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. By leveraging structural equation models and employing randomly generated multivariate Chebyshev polynomials for the simulation of training data, our method demonstrates robust generalizability across both linear and various types of non-linear dependencies. We introduce a novel bilinear attention mechanism (BAM) for explicit processing of dependency information, which operates on the level of covariance matrices of transformed data and respects the geometry of the manifold of symmetric positive definite matrices. Empirical evaluation demonstrates the robustness of our method in detecting a wide range of dependencies, excelling in undirected graph estimation and proving competitive in completed partially directed acyclic graph estimation through a novel two-step approach.