Graph Structure Inference with BAM: Introducing the Bilinear Attention Mechanism

TL;DR

This paper introduces BAM, a neural network with a bilinear attention mechanism, for supervised graph structure learning that effectively detects dependencies in data, including linear and non-linear relationships, with robust generalization.

Contribution

The paper presents a novel bilinear attention mechanism (BAM) for explicit dependency processing and a training approach using simulated data with Chebyshev polynomials for robust graph inference.

Findings

Robust detection of various dependencies including non-linear relationships.

Effective in undirected graph estimation and competitive in partial DAG estimation.

Generalizes well across different data types and dependency structures.

Abstract

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.