MJP

    Markov jump process
DAG

    directed acyclic graph
MSA

    multi-sequence alignment
Glasso

    graphical lasso
SPD

    symmetric positive semi-definite
SEM

    structural equation model
NLP

    natural language processing
ER

    Erdős–Rényi
MLP

    multilayer perceptron
MRF

    Markov random field

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

Philipp Froehlich Heinz Koeppl  Department of Electrical Engineering and
Information TechnologyTechnische Universität Darmstadt, Germany

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

##  1 Introduction

The inference of causal relationships is central to various scientific fields
like biology [Bühlmann et al., 2014, Jones et al., 2012], climate science
[Nowack et al., 2020], economics [Barfuss et al., 2016], and social studies
[Gerstenberg et al., 2021]. These relationships are often represented as
directed edges within dag directed acyclic graph (DAG), a methodology
pioneered in agricultural research about a century ago by Wright [1921] and
widely adopted today. One key application is in estimating gene regulatory
networks from experimental data [Spirtes et al., 2000a, Lauritzen and
Spiegelhalter, 1988]. Graph structure inference, the process of deriving such
graphical representation from observational data, is crucial for gaining
insights into high-dynamical systems [Glymour et al., 2019].

Graph structure inference typically employs unsupervised learning methods to
estimate the underlying graph through either score-based approaches, which
rank graphs by predefined metrics, or constraint-based approaches that
determine edge existence between variable pairs using conditional independence
tests [Vowels et al., 2022]. 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 [Ke et al., 2022]. Constraint-based
methods usually require a large sample size [Vowels et al., 2022], rely on an
elusive optimal threshold hyperparameter111While a neural network approach
also involves hyperparameters, they are primarily involved in the training
process, such that they are less sensitive and less uncertain compared to
those in constraint-based methods, which directly influence the sparsity of
the matrices predicted by the model., and Shah and Peters [2020] 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.
[2015a, b, 2017], Li et al. [2020], Ke et al. [2022], Lorch et al. [2022], Dai
et al. [2023], have recently emerged as an appealing alternative to
unsupervised methods. In this rising approach, models are 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. It proves effective even with
small datasets and provides the option for fine-tuning when labeled data is
accessible.

Supervised causal learning techniques strongly rely on the ability of neural
networks to extract dependency information from observational data matrices.
Moving beyond traditional methods for predicting dependencies of variables
using neural networks, exemplified by the works of Ke et al. [2022], Lorch et
al. [2022], Rao et al. [2021], Kossen et al. [2021], Song et al. [2019], which
typically embed observational data into an expanded space to implicitly
capture dependency information within Euclidean space, we introduce a novel
observation-to-dependency framework that inputs a data matrix, processes it
through parallel channels, and computes transformed covariance matrices. This
process leverages a novel attention mechanism tailored to the geometry of
symmetric positive semi-definite (SPD) matrix space, enhancing graph structure
inference from observational data. Conditional independencies and the behavior
of covariance matrices are intrinsically linked, emphasizing the potential of
the SPD matrix space processing approach to efficiently decode essential
information for graph structure inference, obscured in observational data.
Additionally, utilizing covariance matrices offers a natural method for
deriving $d\times d$ adjacency-shaped matrices from $M\times d$ observational
data matrices in a permutation- and shape-invariant manner, thereby avoiding
the compression of observational data into embedding vectors.

Since covariance matrices are symmetric, our method is inherently suited for
undirected graph estimation and, with an additional step, can infer edge
directions. This is achieved through a two-step strategy: first estimating the
graph’s skeleton and immoralities, then testing these to generate a completed
partially directed acyclic graph (CPDAG) estimate, aligning with the PC-
algorithm but with reduced computational complexity and minimized errors. Our
method’s ability to infer symmetrical information simplifies directional
inference, offering a comprehensive solution for both undirected and directed
graph estimation from observational data.

Lastly, we employ randomly generated multivariate Chebyshev polynomials as
dependency functions within structural equation models for training data
generation in the supervised approach, using their ability to approximate
well-behaved functions with factorially decreasing coefficients.

#### Background.

Our neural network architecture incorporates the self-attention mechanism
[Vaswani et al., 2017, Bahdanau et al., 2015], a well-established method in
natural language processing (NLP) and computer vision with state-of-the-art
performance [Khan et al., 2022]. Kossen et al. [2021] adapted axial attention
[Ho et al., 2019] for non-parametric learning with matrix-shaped input,
introducing attention mechanisms both between datapoints and between
attributes. This axial attention serves as the foundational mechanism for
supervised causal discovery methods, as employed by Ke et al. [2022], Lorch et
al. [2022].

Incorporating manifold constraints into neural network architectures ensures
adherence to geometric constraints, drawing on the analysis of retractions, as
described by Absil et al. [2008], within the SPD matrix space.

Figure 1: Neural network architecture: An input of arbitrarily shape $(M,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 symmetric positive semi-definite (SPD) matrices. The matrices are
then transformed into Euclidean space using the $\operatorname{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 along the channel axis and applying a softmax layer on the channel
axis.

##  2 Method

###  2.1 Supervised approach for graph learning

#### Problem formulation.

We consider an underlying DAG $\mathcal{G}=(\mathcal{V},\mathcal{E})$ with
nodes $\mathcal{V}=\\{v_{1},\dots,v_{d}\\}$ and edges $\mathcal{E}$. The
binary adjacency matrix $\boldsymbol{A}\in\\{0,1\\}^{d\times d}$ denotes the
edge presence. For nodes $v\in\mathcal{V}$, we model the data as a structural
equation model (SEM) as
$X_{v}=f_{v}(X_{\text{pa}_{\mathcal{G}}(v)},\epsilon_{v})$, where $f_{v}$ is a
measurable function of parent nodes and zero-mean error $\epsilon_{v}$. Each
$X_{v}$ follows a distinct probability distribution determined by the SEM.
Here, $X_{\text{pa}_{\mathcal{G}}(v)}$ includes random variables in the parent
set of a node $v$. The data matrix $\boldsymbol{X}\in\mathbb{R}^{M\times d}$
consists of $M$ independent samples from this model. The aim of the graph
inference task is to reconstruct the structure of $\mathcal{G}$, typically
estimating the adjacency matrix $\boldsymbol{A}\in\mathbb{R}^{d\times d}$ from
data $\boldsymbol{X}$.

#### Three-class edge classification.

Our graph structure learning method employs a three-class classification
paradigm for the first-step edge inference. We utilize the assumptions of
faithfulness and the Markov condition to be able to identify the Markov
equivalence class. This enables us to classify pairs of variables distinctly
into:

  * •

Skeleton edges: Representing 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: Signifying variables that remain conditionally independent
considering all other variables.

Our approach distinguishes between skeleton and moralized edges, modeling them
as undirected, which translates to symmetric adjacency matrices. This sets our
method apart from some existing algorithms that primarily focus on the moral
graph [Friedman et al., 2008, Hsieh et al., 2014, Shalom et al., 2022].

Given the assumptions stated earlier, the three-class classification problem
can be uniquely solved by applying specific independence tests to the data
distribution. This can be inferred from the CPDAG’s identifiability under
these assumptions [Spirtes et al., 2000b, Peters et al., 2017], and the
subsequent identifiability of the three-class problem from the CPDAG. In
Appendix B, we provide an explicit proof, outlining the conditional
independence relations that lead to a unique solution.

The targets for prediction can be considered as an extension of the binary
adjacency matrix $\boldsymbol{A}\in\\{0,1\\}^{d\times d}$ to a set of one-hot
encoded adjacency matrices, denoted as
$\widetilde{\boldsymbol{A}}\in\\{0,1\\}^{d\times d\times 3}$. Here, for each
$i,j$, the vector $\widetilde{\boldsymbol{A}}_{i,j,\cdot}\in\\{0,1\\}^{3}$
represents a one-hot encoded classification among the three classes: _skeleton
edge_ , _moralized edge_ , and _no-edge_.

#### Simulation of training data.

Figure 2: Scatterplots illustrating example non-linear dependencies governed
by the structural equation model employed in this study.

To develop a model capable of generalizing across a broad spectrum of
functional dependencies and graph structures, we simulate random graphs,
denoted $\mathcal{G}_{i}$, and generate corresponding random data matrix
realizations, $\boldsymbol{X}_{i}$ for $i=1,\dots$. These realizations are
created from a structural equation model, reflecting the underlying graph
structure and incorporating randomly selected dependency functions. This
results in an input/label pair $(\boldsymbol{X}_{i},\mathcal{G}_{i})$ used for
supervised learning. To enhance the efficiency of neural network training and
prevent overfitting, we employ on-the-fly training, whereby each input/label
pair $(\boldsymbol{X}_{i},\mathcal{G}_{i}),i=1,\dots$ is generated just before
training and discarded once the neural network weights are updated.

Our approach utilizes random Erdős–Rényi graphs, denoted $ER(d,q)$, where the
number of nodes $d$ and the expected degree $q$ are sampled from discrete
uniform distributions $d\sim\mathcal{U}(\\{10,\dots,100\\})$ and $q\sim
U(\\{1,\dots,\min(\frac{d}{3},5)\\})$, respectively. This configuration allows
for denser graphs than those explored in other studies [Dai et al., 2023, Ke
et al., 2022, Yu et al., 2019]. Our model’s versatility enables it to train
without restriction to a fixed $(d,q)$ pair, which proves advantageous when
the graph density is unknown. For each graph $\mathcal{G}_{i}$, we generate a
data matrix $\boldsymbol{X}_{i}\in\mathbb{R}^{M\times d}$ using sem SEM, where
the sample size $M$ is drawn from a discrete uniform distribution,
$M\sim\mathcal{U}(\\{50,\dots,1000\\})$.

We employ random multivariate Chebyshev polynomial functions to generate
diverse continuous training data. Chebyshev polynomials effectively
approximate real-world functions, showing factorial decay in their
coefficients [Xiang and Liu, 2020, Trefethen, 2008], i.e., for the $n$-th
coefficient $c_{n}$ it holds $\|c_{n}\|\leq\frac{C}{n!}$ for a constant $C$,
making higher-degree terms negligible. In our model, error terms use Gaussian
mixture models for a wide range of error distributions [Reynolds, 2009].
Details about the parameterization of the SEM can be found in Appendix D.

###  2.2 Shape-agnostic neural network for graph structure learning

Figure 3: Bilinear self-attention layer. Gray indicates non-trainable tensors,
and red trainable weights. Matrix multiplication is performed after necessary
transposition to match axis dimensions. The double arrow signifies the use of
the matrix as a bilinear operator. $\widetilde{\sigma}$ denotes the custom
softmax, defined in 1.

For simplicity, we denote each layer’s input by $\boldsymbol{X}$ or, if in the
SPD manifold, by $\boldsymbol{S}$. The output is denoted by $\boldsymbol{H}$,
so as not to introduce new notation for every layer. We use the multi-
dimensional analogue of matrix multiplication: For a tensor
$\boldsymbol{A}\in\mathbb{R}^{I\times J\times K}$ and a matrix
$\boldsymbol{B}\in\mathbb{R}^{K\times L}$ we denote

$\boldsymbol{A}\boldsymbol{B}=\boldsymbol{C}\in\mathbb{R}^{I\times J\times
L}\quad\text{with}\quad\boldsymbol{C}_{ijl}=\sum_{k}\boldsymbol{A}_{ijk}\boldsymbol{B}_{kl}$

The network architecture is depicted in Figure 1.

#### Channel embedding.

We perform an embedding of the input $\boldsymbol{X}\in\mathbb{R}^{M\times d}$
to obtain a hidden representation with $C$ channels. For this, one axis for
$\boldsymbol{X}$ is extended to
$\widetilde{\boldsymbol{X}}\in\mathbb{R}^{M\times d\times 1}$ and then
trainable weights $\boldsymbol{W}_{1}\in\mathbb{R}^{1\times C}$,
$\boldsymbol{W}_{2}\in\mathbb{R}^{C\times C}$ are used to obtain

$\boldsymbol{H}=\widetilde{\boldsymbol{X}}+\text{relu}\left(\widetilde{\boldsymbol{X}}\boldsymbol{W}_{1}\right)\boldsymbol{W}_{2}\in\mathbb{R}^{M\times
d\times C}$

where broadcasting is used for the addition.

#### Observational data self-attention.

A description of the observational attention layers is provided in Appendix
A.2.

#### Bilinear attention layer.

After several iterations of data-attention layers, we derive a transformed
data matrix $\boldsymbol{X}\in\mathbb{R}^{M\times d\times C}$. This matrix is
used to compute covariance matrices for each channel, resulting in a tensor of
covariance matrices across $C$ channels:

$\small\boldsymbol{\Sigma}=\frac{1}{M-1}\left(\left(\boldsymbol{X}-\boldsymbol{\mu_{X}}\right)^{T(3,2,1)}\odot\left(\boldsymbol{X}-\boldsymbol{\mu_{X}}\right)^{T(3,1,2)}\right)^{T(2,3,1)},$

where we define for tensors $\boldsymbol{A}\in\mathbb{R}^{I\times J\times K}$,
$\boldsymbol{B}\in\mathbb{R}^{I\times K\times L}$ the $I$-parallel matrix
multiplication

$\boldsymbol{A}\odot\boldsymbol{B}:=\mathopen{}\mathclose{{}\left(\boldsymbol{A}_{i,\cdot,\cdot}\boldsymbol{B}_{i,\cdot,\cdot}}\right)_{i=1,\dots,I}\in\mathbb{R}^{I\times
J\times L},$

and
$\boldsymbol{\mu_{X}}=\frac{1}{M}\boldsymbol{1}_{M}^{T}\boldsymbol{X}\in\mathbb{R}^{1\times
d\times C}$ is a tensor of sample means.

We denote by $\mathcal{S}^{d\times d}_{\succeq}$ the cone of $d\times d$ SPD
matrices, and by $\mathcal{S}^{d\times d\times
C}_{\succeq}=\mathcal{S}^{d\times
d}_{\succeq}\times\dots\times\mathcal{S}^{d\times d}_{\succeq}$ we denote the
$C$-ary Cartesian power of $\mathcal{S}^{d\times d}_{\succeq}$. It holds
$\boldsymbol{\Sigma}\in\mathcal{S}^{d\times d\times C}_{\succeq}$.

Given the nonpositive curvature of the SPD matrices’ Riemannian manifold
[Bhatia, 2009], utilizing SPD matrices in traditional neural networks that
operate in Euclidean space poses challenges [Pennec et al., 2006, Huang and
Gool, 2017]. Thus, specialized SPD networks that respect this manifold’s
geometry are developed.

Bilinear222The mapping
$\boldsymbol{A}\mapsto\boldsymbol{A}\boldsymbol{\Sigma}\boldsymbol{A}^{T}$ is
quadratic, not linear, in $\boldsymbol{A}$. However, it’s often referred to as
’bilinear’ as a specific case of the map
$(\boldsymbol{A},\boldsymbol{B})\mapsto\boldsymbol{A}\boldsymbol{\Sigma}\boldsymbol{B}^{T}$.
matrix multiplication, analogous to a dense layer in Euclidean space, serves
as a primary tool for SPD-nets [Wang et al., 2022]. This is because the
mapping
$\boldsymbol{\Sigma}\rightarrow\boldsymbol{A}^{T}\boldsymbol{\Sigma}\boldsymbol{A}$,
$\mathcal{S}^{d\times d}_{\succeq}\rightarrow\mathcal{S}^{d\times
d}_{\succeq}$ is an endomorphism, i.e., it preserves the space of symmetric,
positive semi-definite matrices $\mathcal{S}^{d\times d}_{\succeq}$ of
dimension $d$.

In contrast to existing SPD layers, which are typically used for image
processing tasks where only a single SPD matrix serves as the covariance
descriptor of an image, we do not directly parameterize weights
$\mathbf{W}\in\mathbb{R}^{d\times d_{\text{out}}}$ to be applied as
$\mathbf{W}^{T}\mathbf{\Sigma}\mathbf{W}$ to a matrix
$\mathbf{\Sigma}\in\mathbb{R}^{d\times d}$. Instead, we parameterize weights
$\boldsymbol{W}\in\mathbb{R}^{C\times C}$ to act as linear combinations
$\boldsymbol{S}\boldsymbol{W}$ on a set of covariance matrices
$\boldsymbol{S}\in\mathcal{S}^{d\times d\times C}_{\succeq}$, leveraging an
attention mechanism to create inner weights adaptable to different matrix
dimensions $d$ for each SPD matrix in $\boldsymbol{S}$.

Our bilinear attention mechanism preserves the positive definiteness of
matrices, requiring only a single matrix logarithm computation for the
transition to the associated Euclidean space. In this space, the softmax
operation over the three classes becomes applicable, enabling the computation
of predictions as output. The architecture, illustrated in Figure 3, imposes
non-negativity constraints on weights
$\boldsymbol{W}_{K}\in\mathbb{R}^{C\times C}_{+}$,
$\boldsymbol{W}_{Q}\in\mathbb{R}^{C\times C}_{+}$ such that
$\boldsymbol{S}\boldsymbol{W}_{K},\boldsymbol{S}\boldsymbol{W}_{Q}\in\mathcal{S}^{d\times
d\times C}_{\succeq}$, leveraging cone properties. We suggest initializing
positive weight matrices $\boldsymbol{W}^{+}$ with samples from
$U\mathopen{}\mathclose{{}\left(0,\frac{2}{n_{\text{in}}}}\right)$, where
$n_{\text{in}}$ is the number of input neurons. In expectation, the diagonal
entries in each matrix within $\boldsymbol{S}\boldsymbol{W}^{+}$ should stay
within a comparable range, while off-diagonal entries are expected to be drawn
to zero due to the symmetric distribution of positive and negative values.
Consequently, keys and queries have essentially initiated as perturbed
identity mappings.

For input $\boldsymbol{S}\in\mathcal{S}^{d\times d\times C}_{\succeq}$, we
obtain keys
$\boldsymbol{K}=\boldsymbol{S}\boldsymbol{W}_{K}\in\mathcal{S}^{d\times
d\times C}_{\succeq}$ and queries
$\boldsymbol{Q}=\boldsymbol{S}\boldsymbol{W}_{Q}\in\mathcal{S}^{d\times
d\times C}_{\succeq}$, which are combined in a bilinear fashion, parallel over
the $C$ channels by calculating

$\boldsymbol{K}\otimes\boldsymbol{Q}:=\left(\boldsymbol{K}_{\cdot,\cdot,c}\boldsymbol{Q}_{\cdot,\cdot,c}\boldsymbol{K}_{\cdot,\cdot,c}\right)_{c=1,\dots,C}\in\mathcal{S}^{d\times
d\times C}_{\succeq}$

for each channel $c$, where indexing corresponds to the channel axis, results
in a tensor of $C$ SPD matrices.

To preserve positive definiteness, which is in general compromised by the non-
symmetric output of standard softmax application on matrix rows, we propose a
custom softmax function. We first define a diagonal scaling matrix as

$\boldsymbol{\Lambda}(\boldsymbol{S}):=\operatorname{diag}\mathopen{}\mathclose{{}\left(\frac{1}{{\exp[\boldsymbol{S}]\boldsymbol{1}_{d}}}}\right),$

where $\exp[\cdot]$ denotes the elementwise application of the exponential
function, $\boldsymbol{1}_{d}$ is a vector of length $d$ with all entries
being $1$, and $\operatorname{diag}$ transforms a vector of length $d$ into a
$d\times d$ diagonal matrix. The quotient is also taken elementwise. With
this, we propose the custom softmax function $\widetilde{\boldsymbol{\sigma}}$
as follows:

$\displaystyle\widetilde{\boldsymbol{\sigma}}:\mathcal{S}^{d\times
d}_{\succeq}\rightarrow\mathcal{S}^{d\times
d}_{\succeq}\quad\quad\widetilde{\boldsymbol{\sigma}}(\boldsymbol{S}):=\sqrt{\boldsymbol{\Lambda}(\boldsymbol{S})}\exp[\boldsymbol{S}]\sqrt{\boldsymbol{\Lambda}(\boldsymbol{S})},$
(1)

where the square root is elementwise.

Note that the elementwise application of the exponential function preserves
positive definiteness. This is because SPD matrices are closed under addition
and the Hadamard product [Bhatia, 2009], and it holds that
$\exp[\boldsymbol{S}]=\sum_{n=0}^{\infty}\frac{1}{n!}[\boldsymbol{S}]^{n}$,
with elementwise exponentiation $[\cdot]^{n}$. Similar to the standard softmax
applied over the rows, using this softmax for a $\mathcal{S}^{d\times
d}_{\succeq}$ matrix returns positive values that altogether sum to $d$.
However, in contrast to the standard softmax, the rows do not sum up to $1$.
We demonstrate that our modified softmax function
$\widetilde{\boldsymbol{\sigma}}$ additionally regularizes the eigenvalues:

######  Theorem 1.

For any $\boldsymbol{S}\in\mathcal{S}^{d\times d}_{\succeq}$, the largest
eigenvalue of $\widetilde{\boldsymbol{\sigma}}(\boldsymbol{S})$ is $1$.

###### Proof.

Let $\widetilde{\boldsymbol{S}}:=\exp[\boldsymbol{S}]$. By similarity
transformation, the eigenvalues of
$\widetilde{\boldsymbol{\sigma}}(\boldsymbol{S})$ are equal to the eigenvalues
of $\widetilde{\boldsymbol{S}}\boldsymbol{\Lambda}(\boldsymbol{S})$. It holds
that
$\widetilde{\boldsymbol{S}}\boldsymbol{\Lambda}(\boldsymbol{S})\widetilde{\boldsymbol{S}}\boldsymbol{1}=\widetilde{\boldsymbol{S}}\boldsymbol{1}$,
which demonstrates that $\widetilde{\boldsymbol{S}}\boldsymbol{1}$ is the
Perron eigenvector corresponding to the eigenvalue $1$. The assertion now
follows from the Perron-Frobenius theorem. ∎

######  Proposition 2.

The custom softmax $\widetilde{\boldsymbol{\sigma}}$ is invariant to additive
shifting, i.e.,
$\widetilde{\boldsymbol{\sigma}}(\boldsymbol{S}+\alpha)=\widetilde{\boldsymbol{\sigma}}(\boldsymbol{S})$
for each $\alpha\in\mathbb{R}$.

Proposition 2 shows that $\widetilde{\boldsymbol{\sigma}}$, unlike standard
softmax, does not need a scaling constant, and it can easily manage exploding
$\exp$ values via maximum-value scaling.

We obtain the attention matrix $\boldsymbol{A}\in\mathcal{S}^{d\times d\times
C}_{\succeq}$ by applying $\widetilde{\boldsymbol{\sigma}}$ channelwise:

$\boldsymbol{A}:=\left(\widetilde{\boldsymbol{\sigma}}\mathopen{}\mathclose{{}\left((\boldsymbol{K}\otimes\boldsymbol{Q}}\right)_{\cdot,\cdot,c})\right)_{c=1,\dots,C}\in\mathcal{S}^{d\times
d\times C}_{\succeq}$

Finally, the output of the bilinear layer is computed for each channel
separately as

$\boldsymbol{H}=\boldsymbol{A}\otimes\boldsymbol{S}\in\mathcal{S}^{d\times
d\times C}_{\succeq}.$

####  SPD activation function.

We use polynomials of degree $3$ with learnable weights $w_{k}\geq 0$,
$\sum_{k}w_{k}\leq 1$ as activation function in the SPD space. For a
discussion and a motivation, see Appendix A.1.

#### Log-Eig layer and output softmax.

Data representation transitions from the SPD space to Euclidean space through
the Log-Eig layer, as Huang and Gool [2017] proposed. This transformation,
given an input matrix
$\boldsymbol{S}=\boldsymbol{U}\boldsymbol{D}\boldsymbol{U}^{T}$ via
eigendecomposition can be expressed as:

$l:\mathcal{S}^{d\times d}_{\succeq}\rightarrow\mathbb{R}^{d\times d},\quad
l(\boldsymbol{S}):=\boldsymbol{\log}\mathopen{}\mathclose{{}\left(\boldsymbol{S}}\right):=\boldsymbol{U}\boldsymbol{\log}\mathopen{}\mathclose{{}\left(\boldsymbol{D}}\right)\boldsymbol{U}^{T}$

The final layer, equipped with a softmax activation function, consists of $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 G. Unlike
traditional attention where the $(i,j)$-th score in the attention matrix
$\boldsymbol{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 $\sum_{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
${\boldsymbol{A}_{i,\cdot}\cup\boldsymbol{A}_{\cdot,j}}$ rather than just
$A_{i,j}$, leading to attention scores forming a cross shape within the matrix
$\boldsymbol{A}$. This approach allows columns of the matrix $\mathbf{S}$ to
attend to each other, highlighting a more complex interaction pattern.

#### Implementation details.

Additional details regarding residual connections, normalization techniques,
and the use of multiple heads in attention layers are elaborated in Appendix
A.3 and A.4.

#### 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 [Meek, 1995]. Details about the CPDAG estimation step can be found in
Appendix C.

##  3 Related Work

Causal discovery largely relies on unsupervised learning methods, namely
constraint-based and score-based approaches [Vowels et al., 2022]. The former
infer conditional independencies [Spirtes et al., 2000b, Hyttinen et al.,
2013, Drton and Maathuis, 2017], while the latter optimize a score function
under acyclicity constraints [Chickering, 2002, Goudet et al., 2018].
Continuous optimization techniques within score-based methods have gained
attention, utilizing various strategies [Zheng et al., 2018, Yu et al., 2019,
Brouillard et al., 2020].

Some studies, inspired by the protein contact prediction task that utilizes
multi-sequence alignment (MSA), have aimed to estimate relations from MSA data
matrices. For instance, Rao et al. [2021] employed axial attention [Ho et al.,
2019] on MSA matrices. Similarly, Li et al. [2019] estimated precision
matrices of MSAs and applied a convolutional neural network to predict
specific protein contacts. While the latter method shares our intuition of
operating on the inverse covariance matrix, our approach significantly differs
by transforming data by axial attention prior to covariance matrix calculation
and implementing tailored layers for SPD matrix processing. Moreover, because
of the known fixed data dimension of $20$ naturally-occurring residue types,
MSA models are specifically designed for the unique conditions of the MSA
prediction task.

Our work is most closely aligned with the studies by Ke et al. [2022] and
Lorch et al. [2022], which similarly leverage attention mechanisms between
samples and attributes to estimate adjacency matrices. However, the
methodologies differ in their strategies for deriving an adjacency matrix from
observational data representations. Ke et al. [2022] employ an autoregressive
transformer approach, whereas Lorch et al. [2022] utilize the dot product of
embedding vectors derived from max-pooling across the sample axis of the
observational data matrices. We demonstrate that the derivation of output
adjacency matrices can be naturally and efficiently achieved through
covariance matrices, which allows for the direct processing of dependence
information and enhances learning efficiency through the application of
geometric learning on the SPD manifold.

Also, while the model of Ke et al. [2022] trains different models for varying
data dimension, our model can be trained across several different shapes of
sample numbers and attribute numbers, thus eliminating the need for re-
training on different datasets. Although the model developed by Lorch et al.
[2022] is capable of evaluation on data with varying dimensions, their
implementation does not support practical training across variable sample
sizes due to memory allocation constraints, and it is restricted to handling a
limited range of dimensions for training simultaneously.

Additionally, our model addresses identifiability challenges by first
estimating an undirected graph, then proceeding to CPDAG estimation employing
Meek’s rules. This methodology ensures that edge directionality is deduced
only when it can be reliably determined, indicated by the identification of
v-structures. In contrast, the networks developed by Ke et al. [2022] and
Lorch et al. [2022] deduce a DAG utilizing both observational and
interventional data, dependent on the availability of interventional data or
making random guesses for edges that cannot be directly inferred.

While Li et al. [2020] also used permutation-invariant models for causal
inference, their method directly computes correlation matrices, which may
limit its efficacy when covariance matrices are not sufficient. They apply
layers of the form
$\boldsymbol{Y}=\omega_{1}\boldsymbol{S}+\omega_{2}\boldsymbol{1}\boldsymbol{1}^{T}\boldsymbol{S}+\omega_{3}\boldsymbol{S}\boldsymbol{1}\boldsymbol{1}^{T}+\omega_{4}\boldsymbol{1}\boldsymbol{1}^{T}\boldsymbol{S}\boldsymbol{1}\boldsymbol{1}^{T}+b$,
with scalar weights $\omega_{i},i=1,\dots,4$, and bias $b$. Despite preserving
invariances, the limited free parameters could restrict its representational
power.

Lopez-Paz et al. [2015b, a] proposed a supervised learning framework that
leverages kernel mean embeddings. Ma et al. [2022] address identifiability
with independence tests and cascade classifiers for supervised skeleton
learning, trained on vicinal graphs specific to observational data. Our
method, however, adopts a more general approach, training a network that does
not require re-training for new evaluation samples. Dai et al. [2023]’s work
on immoralities aligns with our second CPDAG estimation step.

##  4 Experiments

Figure 4: Undirected graph estimation results, arranged in each case from
worst (left) to best (right). AUC values. (a) and (b) AUC values for different
dependencies, with (a) $d=50$, $M=200$ and (b) $d=100$, $M=50$. (c) shows
accuracy values for the same dependencies at $d=100$, $M=50$. (d) and (e)
present AUC values for different sample sizes with Chebyshev (d) and cosine
(e) dependencies at $d=100$. (f) displays structural Hamming distance for
varying sample sizes in a high-dimensional setting ($d=100$) for Chebyshev
dependency. Figure 5: CPDAG estimation results ordered from worst (left) to
best (right). (a)-(c) SHD for various dependencies at (a) $d=20$, $M=200$, (b)
$d=50$, $M=200$, and (c) $d=100$, $M=500$. (d) AUC for Chebyshev dependency
across $M=50,100,200,500,1000$ at $d=100$. (e)+(f) SHD at $d=100$ for
Chebyshev and sine dependencies, respectively, over the same $M$ values.

#### Baseline algorithms.

We evaluate BAM’s performance against other algorithms: PC [Spirtes et al.,
2000b], PC-HSIC [Zhang et al., 2012], rcot and rcit [Strobl et al., 2019],
ccdr [Aragam and Zhou, 2015], GES [Chickering, 2002], GIES [Hauser and
Bühlmann, 2012], LiNGAM [Shimizu et al., 2006], MMPC [Tsamardinos et al.,
2006], CAM [Bühlmann et al., 2014], SAM [Kalainathan et al., 2022], all
implemented in the causal discovery toolbox [Kalainathan et al., 2020], glasso
[Friedman et al., 2008], as implemented by Pedregosa et al. [2011], along with
DAG-GNN [Yu et al., 2019]. For a supervised causal discovery baseline, we use
the AVICI model from Lorch et al. [2022], where we use a pre-trained version
(Avici SCM-v0), as well as a version trained from scratch on the same
Chebyshev data as our method. For this, we used default hyperparameters,
particularly, we used data dimensions $d\in\\{2,5,10,20,30,40,60,80,100\\}$.
For the number of samples $M$, which can only be chosen as a single value due
to memory allocation constraints, we selected $M$ = 150, as larger sizes
exceeded our server’s graphical memory limits (approximately 80 GB). Default
hyperparameters from Kalainathan et al. [2020] are used. For glasso, cross-
validation for the sparsity parameter is performed as proposed by [Pedregosa
et al., 2011]. Glasso and MMPC are excluded from CPDAG estimation as they
focus on undirected graphs. For algorithms which estimate directed edges, the
skeleton is computed for undirected prediction. Addtionally, we compare SHD
and accuracy results against a zero graph without any edges as naive baseline.

#### Performance indicators.

We employ the Area Under the Precision-Recall Curve (AUC) for evaluating both
undirected and CPDAG graph estimations, recognizing its appropriateness for
imbalanced binary classification tasks such as sparse graph detection [He and
Garcia, 2009]. For CPDAG estimations, we further utilize the Structural
Hamming Distance (SHD), a benchmark metric in structure learning [Yu et al.,
2019, Ke et al., 2022]. For the undirected graph estimation task, SHD is
equivalent to accuracy, defined as the percentage of correctly inferred edges,
which we also consider in this study.

#### Results.

Figure 4 showcases BAM’s efficacy in undirected graph prediction. Trained on
synthetic Chebyshev polynomial data, BAM was evaluated across various
dependency relations, as shown in Figure 4 (a), (b), (c). It consistently
outperforms other methods in both moderate and high dimensions. Remarkably,
BAM excels in capturing intricate non-monotonic dependencies such as cosine
and $x^{2}$ dependencies, illustrated in Figure 4 (a), (b), (e). When applied
to mainly monotonic dependencies like sine in Figure 4 (c), BAM demonstrates
superior performance across multiple sample sizes, affirming its
generalizability.

To contrast with Avici’s distinct model architecture, plots (d) and (f)
present the performance for varying sample sizes $M$ on Chebyshev dependency
data, utilized in training both BAM and the re-trained Avici version. These
results demonstrate BAM’s superior performance over Avici on the training set.

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 (a graph devoid of edges) in SHD, as shown in Appendix F.3. Therefore,
our SHD analysis focuses on the low-dimensional context shown in (a) and the
moderate-dimensional context in (b) and in in (c) another regime where $d<M$,
specifically $d=100$ and $M=500$, 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 from data.
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.

Table 1: Ablation study results featuring loss values for evaluation. Data generated under Chebyshev dependencies. ”-” indicates the removal of a corresponding layer. $\Delta\text{Param}$ quantifies the difference in the number of parameters to the full model. model | loss ↓ | $\Delta\text{Param}$
---|---|---
FULL | $0.173\pm 0.007$ |
$-$ bilinear | $0.202\pm 0.005$ | 80 K
$-$ bilinear $-$ LogEig | $0.271\pm 0.012$ | 100 K
$-$ obs. att. | $0.189\pm 0.006$ | 120 K
$-$ obs. att. $-$ Dense | $0.205\pm 0.006$ | 160 K

#### Ablation Study.

Table 1 presents the results of our ablation study, highlighting the critical
role of the bilinear layer. Its removal leads to a significant increase in
loss metrics. The omission of the LogEig layer results in a pronounced
deterioration of the loss, showing that direct predictions from the SPD-space
are leading to problems. Interestingly, bilinear data processing alone, even
in the absence of additional data processing, yields relatively good results.
This may be attributable to the ability of the embedding layer’s nonlinearity
to decode data non-monotonicities, which could be obscured when solely relying
on covariance matrices. Further details regarding the ablation study can be
found in the Appendix.

#### Efficiency.

Training of the neural network was executed in approximately 6 hours on an
A-100 GPU with 81,920 MiB of graphical memory. Inference typically requires
less a few seconds. In contrast many unsupervised approaches incur a
significantly higher runtime. A time comparison is provided in Appendix F.4.

##  5 Conclusion

In this study, we introduced a novel neural network model for supervised graph
structure learning, addressing the identifiability issue in observational data
and modeling dependence relations through random Chebyshev polynomial
dependencies. We introduced an observational-to-dependency processing
approach, operating in both the Euclidean observational space and the SPD
covariance manifold. Our model incorporates a novel bilinear attention
mechanism and a permutation- and shape-agnostic architecture. We also
establish theoretical properties confirming the model’s robustness.
Simulations demonstrate that BAM outperforms existing models in various
scenarios.

BAM operates on transformed data’s covariance matrices, enabling explicit
processing of dependency information in the SPD matrix manifold, suggesting
its potential applicability to other optimization challenges on the SPD
manifold. Also, the observation-to-dependency processing model opens promising
avenues for research, especially in applications requiring an understanding of
variable dependencies.

## References

  * Abadi et al. [2015] Martín Abadi, Ashish Agarwal, Paul Barham, and et al.  TensorFlow: Large-scale machine learning on heterogeneous systems, 2015\.  URL https://www.tensorflow.org/.  Software available from tensorflow.org.
  * Absil et al. [2008] P-A Absil, Robert Mahony, and Rodolphe Sepulchre.  _Optimization algorithms on matrix manifolds_.  Princeton University Press, 2008.
  * Apicella et al. [2021] Andrea Apicella, Francesco Donnarumma, Francesco Isgrò, and Roberto Prevete.  A survey on modern trainable activation functions.  _Neural Networks_ , 138:14–32, 2021.
  * Aragam and Zhou [2015] Bryon Aragam and Qing Zhou.  Concave penalized estimation of sparse gaussian bayesian networks.  _The Journal of Machine Learning Research_ , 16(1):2273–2328, 2015.
  * Ba et al. [2016] Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton.  Layer Normalization.  _arXiv:1607.06450_ , 2016.
  * Bachlechner et al. [2021] Thomas Bachlechner, Bodhisattwa Prasad Majumder, Henry Mao, Gary Cottrell, and Julian McAuley.  ReZero is All You Need: Fast Convergence at Large Depth.  In _UAI_ , pages 1352–1361, 2021.
  * Bahdanau et al. [2015] Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio.  Neural Machine Translation by Jointly Learning to Align and Translate.  In _ICLR_ , 2015.
  * Barfuss et al. [2016] Wolfram Barfuss, Guido Previde Massara, T. Di Matteo, and Tomaso Aste.  Parsimonious modeling with information filtering networks.  _Phys. Rev. E_ , 94, 2016.
  * Bhatia [2009] Rajendra Bhatia.  Positive Definite Matrices.  In _Positive Definite Matrices_. Princeton University Press, 2009\.
  * Brouillard et al. [2020] Philippe Brouillard, Sébastien Lachapelle, Alexandre Lacoste, Simon Lacoste-Julien, and Alexandre Drouin.  Differentiable Causal Discovery from Interventional Data.  In _NeurIPS_ , volume 33, pages 21865–21877, 2020.
  * Bühlmann et al. [2014] Peter Bühlmann, Markus Kalisch, and Lukas Meier.  High-Dimensional Statistics with a View Toward Applications in Biology.  _Annu. Rev. Stat. Appl._ , 1(1):255–278, 2014\.
  * Bühlmann et al. [2014] Peter Bühlmann, Jonas Peters, and Jan Ernest.  Cam: Causal additive models, high-dimensional order search and penalized regression.  _The Annals of Statistics_ , 2014.
  * Chickering [2002] David Maxwell Chickering.  Optimal structure identification with greedy search.  _Journal of machine learning research_ , 3(Nov):507–554, 2002.
  * Chung et al. [2016] Hoon Chung, Sung Joo Lee, and Jeon Gue Park.  Deep Neural Network Using Trainable Activation Functions.  In _IJCNN_ , pages 348–352, 2016.
  * Dai et al. [2023] Haoyue Dai, Rui Ding, Yuanyuan Jiang, Shi Han, and Dongmei Zhang.  ML4C: Seeing Causality Through Latent Vicinity.  In _Proceedings of the 2023 SIAM International Conference on Data Mining (SDM)_ , pages 226–234, 2023.
  * Drton and Maathuis [2017] Mathias Drton and Marloes H. Maathuis.  Structure learning in graphical modeling.  _Annual Review of Statistics and Its Application_ , 4(1):365–393, 2017.
  * Friedman et al. [2008] Jerome Friedman, Trevor Hastie, and Robert Tibshirani.  Sparse inverse covariance estimation with the graphical lasso.  _Biostatistics_ , 9(3):432–441, 2008.
  * Gerstenberg et al. [2021] Tobias Gerstenberg, Noah D Goodman, David A Lagnado, and Joshua B Tenenbaum.  A Counterfactual Simulation Model of Causal Judgments for Physical Events.  _Psychol. Rev._ , 128(5):936, 2021.
  * Glymour et al. [2019] Clark Glymour, Kun Zhang, and Peter Spirtes.  Review of causal discovery methods based on graphical models.  _Frontiers in Genetics_ , 10:524, 2019.
  * Goudet et al. [2018] Olivier Goudet, Diviyan Kalainathan, Philippe Caillou, Isabelle Guyon, et al.  Learning Functional Causal Models with Generative Neural Networks.  In _Explainable and Interpretable Models in Computer Vision and Machine Learning_ , pages 39–80, 2018.
  * Guillot and Rajaratnam [2015] D. Guillot and B. Rajaratnam.  Functions Preserving Positive Definiteness for Sparse Matrices.  _Proc. Am. Math. Soc._ , 367:627–649, 2015.
  * Hauser and Bühlmann [2012] Alain Hauser and Peter Bühlmann.  Characterization and greedy learning of interventional markov equivalence classes of directed acyclic graphs.  _The Journal of Machine Learning Research_ , 13(1):2409–2464, 2012.
  * He and Garcia [2009] Haibo He and Edwardo A. Garcia.  Learning from imbalanced data.  _IEEE Transactions on Knowledge and Data Engineering_ , 21(9):1263–1284, 2009.  10.1109/TKDE.2008.239.
  * He et al. [2016] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.  Deep Residual Learning for Image Recognition.  In _CVPR_ , pages 770–778, 2016.
  * Ho et al. [2019] Jonathan Ho, Nal Kalchbrenner, Dirk Weissenborn, and Tim Salimans.  Axial Attention in Multidimensional Transformers.  _arXiv:1912.12180_ , 2019.
  * Hoffman and Withers [1988] Michael E. Hoffman and William Douglas Withers.  Generalized chebyshev polynomials associated with affine weyl groups.  _Transactions of the American Mathematical Society_ , 308(1):91–104, 1988.
  * Hsieh et al. [2014] Cho-Jui Hsieh, Mátyás A. Sustik, Inderjit S. Dhillon, and Pradeep Ravikumar.  QUIC: Quadratic Approximation for Sparse Inverse Covariance Estimation.  _J. Mach. Learn. Res._ , 15(83):2911–2947, 2014\.
  * Huang and Gool [2017] Zhiwu Huang and Luc Van Gool.  A Riemannian Network for SPD Matrix Learning.  In _AAAI_ , page 2036–2042, 2017.
  * Hyttinen et al. [2013] A. Hyttinen, F. Eberhardt, O. Hoyer, and M. Jarvisalo.  Discovering cyclic causal models with latent variables: a general sat-based procedure.  In _Proceedings of the 29th Conference on Uncertainty in Artificial Intelligence_ , 2013.
  * Jones et al. [2012] David T Jones, Daniel W A Buchan, Domenico Cozzetto, and Massimiliano Pontil.  PSICOV: precise structural contact prediction using sparse inverse covariance estimation on large multiple sequence alignments.  _Bioinformatics_ , 28(2):184–190, 2012.
  * Kalainathan et al. [2020] Diviyan Kalainathan, Olivier Goudet, and Ritik Dutta.  Causal discovery toolbox: Uncovering causal relationships in python.  _The Journal of Machine Learning Research_ , 21(1):1406–1410, 2020.
  * Kalainathan et al. [2022] Diviyan Kalainathan, Olivier Goudet, Isabelle Guyon, David Lopez-Paz, and Michèle Sebag.  Structural Agnostic Modeling: Adversarial Learning of Causal Graphs.  _J. Mach. Learn. Res._ , 23(219):1–62, 2022\.
  * Ke et al. [2022] Nan Rosemary Ke, Silvia Chiappa, Jane X Wang, Jorg Bornschein, Anirudh Goyal, Melanie Rey, Matthew Botvinick, Theophane Weber, et al.  Learning to induce causal structure.  In _ICML: Workshop on Spurious Correlations, Invariance and Stability_ , 2022.
  * Khan et al. [2022] Salman Khan, Muzammal Naseer, Munawar Hayat, Syed Waqas Zamir, Fahad Shahbaz Khan, and Mubarak Shah.  Transformers in Vision: A Survey.  _ACM Comput. Surv._ , 54(10s), 2022.
  * Kingma and Ba [2015] Diederik P. Kingma and Jimmy Ba.  Adam: A method for stochastic optimization.  In Yoshua Bengio and Yann LeCun, editors, _3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings_ , 2015.
  * Koller and Friedman [2009] Daphne Koller and Nir Friedman.  _Probabilistic Graphical Models: Principles and Techniques_.  MIT press, 2009.
  * Kossen et al. [2021] Jannik Kossen, Neil Band, Clare Lyle, Aidan N Gomez, Thomas Rainforth, and Yarin Gal.  Self-Attention Between Datapoints: Going Beyond Individual Input-Output Pairs in Deep Learning.  In _NeurIPS_ , 2021.
  * Lauritzen and Spiegelhalter [1988] S. L. Lauritzen and D. J. Spiegelhalter.  Local computations with probabilities on graphical structures and their application to expert systems.  _Journal of the Royal Statistical Society. Series B (Methodological)_ , 50(2):157–224, 1988.  ISSN 00359246.
  * Li et al. [2020] Hebi Li, Qi Xiao, and Jin Tian.  Supervised Whole DAG Causal Discovery.  _arXiv:2006.04697_ , 2020.
  * Li et al. [2019] Yang Li, Jun Hu, Chengxin Zhang, Dong-Jun Yu, and Yang Zhang.  ResPRE: high-accuracy protein contact prediction by coupling precision matrix with deep residual neural networks.  _Bioinformatics_ , 35(22):4647–4655, 2019.
  * Lin et al. [2013] Min Lin, Qiang Chen, and Shuicheng Yan.  Network in network.  _arXiv preprint arXiv:1312.4400_ , 2013.
  * Lopez-Paz et al. [2015a] David Lopez-Paz, Krikamol Muandet, and Benjamin Recht.  The randomized causation coefficient.  _JMLR_ , 16:2901–2907, 2015a.
  * Lopez-Paz et al. [2015b] David Lopez-Paz, Krikamol Muandet, Bernhard Schölkopf, and Iliya Tolstikhin.  Towards a learning theory of cause-effect inference.  In _ICML_ , pages 1452–1461, 2015b.
  * Lopez-Paz et al. [2017] David Lopez-Paz, Robert Nishihara, Soumith Chintala, Bernhard Schölkopf, and Leon Bottou.  Discovering causal signals in images.  In _Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR)_ , July 2017.
  * Lorch et al. [2022] Lars Lorch, Scott Sussex, Jonas Rothfuss, Andreas Krause, and Bernhard Schölkopf.  Amortized inference for causal structure learning.  In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, editors, _Advances in Neural Information Processing Systems_ , volume 35, pages 13104–13118. Curran Associates, Inc., 2022.
  * Ma et al. [2022] Pingchuan Ma, Rui Ding, Haoyue Dai, Yuanyuan Jiang, Shuai Wang, Shi Han, and Dongmei Zhang.  ML4S: Learning Causal Skeleton from Vicinal Graphs.  In _SIGKDD_ , page 1213–1223, 2022.
  * Meek [1995] Christopher Meek.  Causal Inference and Causal Explanation with Background Knowledge.  In _UAI_ , pages 403–410, 1995.
  * Nowack et al. [2020] Peer Nowack, Jakob Runge, Veronika Eyring, and Joanna D Haigh.  Causal networks for climate model evaluation and constrained projections.  _Nature communications_ , 11(1):1415, 2020.
  * Pedregosa et al. [2011] F. Pedregosa, G. Varoquaux, A. Gramfort, V. Michel, B. Thirion, O. Grisel, M. Blondel, P. Prettenhofer, R. Weiss, V. Dubourg, J. Vanderplas, A. Passos, D. Cournapeau, M. Brucher, M. Perrot, and E. Duchesnay.  Scikit-learn: Machine learning in Python.  _Journal of Machine Learning Research_ , 12(85):2825–2830, 2011.
  * Pennec et al. [2006] Xavier Pennec, Pierre Fillard, and Nicholas Ayache.  A Riemannian Framework for Tensor Computing.  _Int. J. Comput. Vis._ , 66(1):41–66, 2006\.
  * Peters et al. [2017] Jonas Peters, Dominik Janzing, and Bernhard Schölkopf.  _Elements of Causal Inference_.  MIT Press, Cambridge, Massachusetts, 2017.
  * Puschel and Rotteler [2007] Markus Puschel and Martin Rotteler.  Algebraic signal processing theory: 2-d spatial hexagonal lattice.  _IEEE Transactions on Image Processing_ , 16(6):1506–1521, 2007.
  * Rao et al. [2021] Roshan M Rao, Jason Liu, Robert Verkuil, Joshua Meier, John Canny, Pieter Abbeel, Tom Sercu, and Alexander Rives.  MSA Transformer.  In _ICML_ , volume 139, pages 8844–8856. PMLR, 2021.
  * Reynolds [2009] Douglas A Reynolds.  Gaussian mixture models.  _Encyclopedia of biometrics_ , 741(659-663), 2009.
  * Schoenberg [1942] I. J. Schoenberg.  Positive definite functions on spheres.  _Duke Mathematical Journal_ , 9:96–108, 1942.  MR0005922 (3,232c).
  * Shah and Peters [2020] R. D. Shah and J. Peters.  The hardness of conditional independence testing and the generalised covariance measure.  _The Annals of Statistics_ , 48(3), 2020.
  * Shalom et al. [2022] Gal Shalom, Eran Treister, and Irad Yavneh.  pista: preconditioned iterative soft thresholding algorithm for graphical lasso.  _arXiv:2205.10027_ , 2022.
  * Shimizu et al. [2006] Shohei Shimizu, Patrik O Hoyer, Aapo Hyvärinen, Antti Kerminen, and Michael Jordan.  A linear non-gaussian acyclic model for causal discovery.  _Journal of Machine Learning Research_ , 7(10), 2006.
  * Song et al. [2019] Weiping Song, Chence Shi, Zhiping Xiao, Zhijian Duan, Yewen Xu, Ming Zhang, and Jian Tang.  AutoInt: Automatic Feature Interaction Learning via Self-Attentive Neural Networks.  In _CIKM_ , pages 1161–1170, 2019.
  * Spirtes et al. [2000a] Pater Spirtes, Clark Glymour, Richard Scheines, Stuart Kauffman, Valerio Aimale, and Frank Wimberly.  Constructing Bayesian network models of gene expression networks from microarray data.  In _Proceedings of the Atlantic Symposium on Computational Biology_. Carnegie Mellon University, 2000a.
  * Spirtes et al. [2000b] Peter Spirtes, Clark N Glymour, Richard Scheines, and David Heckerman.  _Causation, Prediction, and Search_.  MIT Press, 2000b.
  * Strobl et al. [2019] Eric V Strobl, Kun Zhang, and Shyam Visweswaran.  Approximate kernel-based conditional independence tests for fast non-parametric causal discovery.  _Journal of Causal Inference_ , 7(1):20180017, 2019.
  * Touvron et al. [2021] Hugo Touvron, Matthieu Cord, Alexandre Sablayrolles, Gabriel Synnaeve, and Hervé Jégou.  Going deeper with Image Transformers.  In _ICCV_ , pages 32–42, 2021.
  * Trefethen [2008] Lloyd N. Trefethen.  Is Gauss Quadrature Better than Clenshaw–Curtis?  _SIAM Review_ , 50(1):67–87, 2008.
  * Tsamardinos et al. [2006] Ioannis Tsamardinos, Laura E Brown, and Constantin F Aliferis.  The max-min hill-climbing bayesian network structure learning algorithm.  _Machine learning_ , 65:31–78, 2006.
  * Vaswani et al. [2017] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Ł ukasz Kaiser, and Illia Polosukhin.  Attention is All you Need.  In _NeurIPS_ , volume 30, 2017.
  * Vowels et al. [2022] Matthew J. Vowels, Necati Cihan Camgoz, and Richard Bowden.  D’Ya Like DAGs? A Survey on Structure Learning and Causal Discovery.  _ACM Comput. Surv._ , 55(4):82:1–82:36, nov 2022.  ISSN 0360-0300.
  * Wang et al. [2022] Rui Wang, Xiao-Jun Wu, Tianyang Xu, Cong Hu, and Josef Kittler.  Deep Metric Learning on the SPD Manifold for Image Set Classification.  _IEEE Trans. Circuits Syst. Video. Technol._ , 2022.
  * Wright [1921] Sewall Wright.  Correlation and causation.  _J. Agric. Res._ , 20(7):557–585, 1921.
  * Wu and He [2018] Yuxin Wu and Kaiming He.  Group Normalization.  In _ECCV_ , pages 3–19, 2018.
  * Xiang and Liu [2020] Shuhuang Xiang and Guidong Liu.  Optimal decay rates on the asymptotics of orthogonal polynomial expansions for functions of limited regularities.  _Numerische Mathematik_ , 145(1):117–148, 2020\.
  * Xiong et al. [2020] Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, et al.  On Layer Normalization in the Transformer Architecture.  In _ICML_ , volume 119, pages 10524–10533, 2020.
  * Yu et al. [2019] Yue Yu, Jie Chen, Tian Gao, and Mo Yu.  DAG-GNN: DAG structure learning with graph neural networks.  In _ICML_ , pages 7154–7163, 2019.
  * Zhang et al. [2012] Kun Zhang, Jonas Peters, Dominik Janzing, and Bernhard Schölkopf.  Kernel-based conditional independence test and application in causal discovery.  _arXiv preprint arXiv:1202.3775_ , 2012.
  * Zheng et al. [2018] Xun Zheng, Bryon Aragam, Pradeep K Ravikumar, and Eric P Xing.  DAGs with NO TEARS: Continuous Optimization for Structure Learning.  In _NeurIPS_ , volume 31, 2018.

Graph Structure Inference with BAM:
Introducing the Bilinear Attention Mechanism
(Appendix)

##  Appendix A Architectural Details

###  A.1 SPD Activation Function

To obtain an activation function for the SPD net, we leverage the following
theorem, a direct consequence of (Theorem 4.11 Guillot and Rajaratnam, 2015),
which is based on the work of Schoenberg [1942]:

######  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]\to[-1,1]$ with

$\displaystyle f(x)=\sum_{k=1}^{\infty}x^{k}w_{k},$
$\displaystyle\quad\text{\rm{subject to}}\quad\sum_{k=1}^{\infty}w_{k}\leq
1\quad\text{\rm{with}}\quad w_{k}\geq 0,\quad\forall k\in\mathbb{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_{\text{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 [Chung et al., 2016] for
general neural networks, not focusing on SPD data. Additionally, [Apicella et
al., 2021] provides various types of trainable activation functions for
Euclidean neural networks.

###  A.2 Observational Attention

This layer accepts an input $\boldsymbol{X}\in\mathbb{R}^{M\times d\times C}$,
from which it generates keys
$\boldsymbol{K}=\boldsymbol{X}\boldsymbol{W}_{K}\in\mathbb{R}^{M\times d\times
c}$, queries
$\boldsymbol{Q}=\boldsymbol{X}\boldsymbol{W}_{Q}\in\mathbb{R}^{M\times d\times
c}$, and values
$\boldsymbol{V}=\boldsymbol{X}\boldsymbol{W}_{V}\in\mathbb{R}^{M\times d\times
C}$, with $\boldsymbol{W}_{K},\boldsymbol{W}_{Q}\in\mathbb{R}^{C\times c}$,
and $\boldsymbol{W}_{V}\in\mathbb{R}^{C\times C}$.

For each $m=1,\dots,M$, keys and queries are combined in parallel along the
inner axis, leading to

$\boldsymbol{K}\odot\boldsymbol{Q}:=\mathopen{}\mathclose{{}\left(\boldsymbol{K}_{m,\cdot,\cdot}\boldsymbol{Q}_{m,\cdot,\cdot}}\right)_{m=1,\dots,M}\in\mathbb{R}^{M\times
d\times d}.$

This results in the attention weights

$\boldsymbol{A}=\sigma\mathopen{}\mathclose{{}\left(\frac{\boldsymbol{K}\odot\boldsymbol{Q}^{T(1,3,2)}}{\sqrt{c}}}\right)\in\mathbb{R}^{M\times
d\times d},$

where $T(\operatorname{perm})$ denotes a permutation of the axes according to
a permutation $\operatorname{perm}$ and $\sigma$ is the softmax operation
along the last axis. For example, if $\text{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 as

$\boldsymbol{H}=\boldsymbol{A}\odot\boldsymbol{V}\in\mathbb{R}^{M\times
d\times C}.$

Figure 6: Left: Observational data self-attention layer across attributes.
Gray denotes non-trainable tensors, and red represents trainable weights. The
model-channel dimension, $c$, is chosen to be smaller than $C$ for efficiency.
Matrix multiplication is performed after necessary transposition to match axis
dimensions. Right: Bilinear self-attention layer. Gray indicates non-trainable
tensors, and red denotes trainable weights. The double arrow signifies the use
of the matrix as a bilinear operator.

###  A.3 Residual Connections and Normalization.

Normalization is critical in attention networks [Xiong et al., 2020], with
various methodologies available [Wu and He, 2018, Ba et al., 2016]. We adopt
normalization and a residual connection [He et al., 2016] on the input of all
attention layers, expressed as
$\boldsymbol{S}+\operatorname{Attention}(\operatorname{LayerNorm}(\boldsymbol{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 methods
from Touvron et al. [2021], Bachlechner et al. [2021].

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.

###  A.4 Multiple Heads

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.

##  Appendix B Theoretical Foundation for the Three-Class Edge Classification
Problem

Building upon the foundational principles of Markov and faithfulness [Koller
and Friedman, 2009], 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\rightarrow Y\implies X\mathchoice{\mathrel{\hbox to0.0pt{ \kern 7.63889pt
\kern-5.27776pt$\displaystyle\not$
\hss}{\perp\mkern-9.5mu\perp}}}{\mathrel{\hbox to0.0pt{ \kern 7.63889pt
\kern-5.27776pt$\textstyle\not$ \hss}{\perp\mkern-9.5mu\perp}}}{\mathrel{\hbox
to0.0pt{ \kern 2.61118pt \kern-4.11108pt$\scriptstyle\not$
\hss}{\perp\mkern-9.5mu\perp}}}{\mathrel{\hbox to0.0pt{ \kern 1.11118pt
\kern-3.3333pt$\scriptscriptstyle\not$
\hss}{\perp\mkern-9.5mu\perp}}}Y\mid\boldsymbol{C}\quad\forall\boldsymbol{C}\in\mathcal{P}\mathopen{}\mathclose{{}\left(\mathcal{V}\setminus\\{X,Y\\}}\right).$

For an immorality between $X$ and $Y$, there exists a set of nodes
$\boldsymbol{C}\in\mathcal{P}\mathopen{}\mathclose{{}\left(\mathcal{V}\setminus\\{X,Y\\}}\right)$
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.:

$\displaystyle X\rightarrow Z\leftarrow
Y,X\mathchoice{\mathrel{\ooalign{$\displaystyle\leftarrow$\cr$\rotatebox[origin={c}]{60.0}{$\displaystyle-$}\mkern-2.0mu$\cr}}}{\mathrel{\ooalign{$\textstyle\leftarrow$\cr$\rotatebox[origin={c}]{60.0}{$\textstyle-$}\mkern-2.0mu$\cr}}}{\mathrel{\ooalign{$\scriptstyle\leftarrow$\cr$\rotatebox[origin={c}]{60.0}{$\scriptstyle-$}\mkern-2.0mu$\cr}}}{\mathrel{\ooalign{$\scriptscriptstyle\leftarrow$\cr$\rotatebox[origin={c}]{60.0}{$\scriptscriptstyle-$}\mkern-2.0mu$\cr}}}Y,X\mathchoice{\mathrel{\ooalign{$\displaystyle\rightarrow$\cr$\rotatebox[origin={c}]{60.0}{$\displaystyle-$}\mkern
2.0mu$\cr}}}{\mathrel{\ooalign{$\textstyle\rightarrow$\cr$\rotatebox[origin={c}]{60.0}{$\textstyle-$}\mkern
2.0mu$\cr}}}{\mathrel{\ooalign{$\scriptstyle\rightarrow$\cr$\rotatebox[origin={c}]{60.0}{$\scriptstyle-$}\mkern
2.0mu$\cr}}}{\mathrel{\ooalign{$\scriptscriptstyle\rightarrow$\cr$\rotatebox[origin={c}]{60.0}{$\scriptscriptstyle-$}\mkern
2.0mu$\cr}}}Y\implies\exists\boldsymbol{C}\in\mathcal{P}\mathopen{}\mathclose{{}\left(\mathcal{V}\setminus\\{X,Y\\}}\right):X\perp\mkern-9.5mu\perp
Y\mid\boldsymbol{C},X\mathchoice{\mathrel{\hbox to0.0pt{ \kern 7.63889pt
\kern-5.27776pt$\displaystyle\not$
\hss}{\perp\mkern-9.5mu\perp}}}{\mathrel{\hbox to0.0pt{ \kern 7.63889pt
\kern-5.27776pt$\textstyle\not$ \hss}{\perp\mkern-9.5mu\perp}}}{\mathrel{\hbox
to0.0pt{ \kern 2.61118pt \kern-4.11108pt$\scriptstyle\not$
\hss}{\perp\mkern-9.5mu\perp}}}{\mathrel{\hbox to0.0pt{ \kern 1.11118pt
\kern-3.3333pt$\scriptscriptstyle\not$
\hss}{\perp\mkern-9.5mu\perp}}}Y\mid\mathcal{V}\setminus\\{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.,

$\displaystyle
X\mathchoice{\mathrel{\ooalign{$\displaystyle\leftarrow$\cr$\rotatebox[origin={c}]{60.0}{$\displaystyle-$}\mkern-2.0mu$\cr}}}{\mathrel{\ooalign{$\textstyle\leftarrow$\cr$\rotatebox[origin={c}]{60.0}{$\textstyle-$}\mkern-2.0mu$\cr}}}{\mathrel{\ooalign{$\scriptstyle\leftarrow$\cr$\rotatebox[origin={c}]{60.0}{$\scriptstyle-$}\mkern-2.0mu$\cr}}}{\mathrel{\ooalign{$\scriptscriptstyle\leftarrow$\cr$\rotatebox[origin={c}]{60.0}{$\scriptscriptstyle-$}\mkern-2.0mu$\cr}}}Y,X\mathchoice{\mathrel{\ooalign{$\displaystyle\rightarrow$\cr$\rotatebox[origin={c}]{60.0}{$\displaystyle-$}\mkern
2.0mu$\cr}}}{\mathrel{\ooalign{$\textstyle\rightarrow$\cr$\rotatebox[origin={c}]{60.0}{$\textstyle-$}\mkern
2.0mu$\cr}}}{\mathrel{\ooalign{$\scriptstyle\rightarrow$\cr$\rotatebox[origin={c}]{60.0}{$\scriptstyle-$}\mkern
2.0mu$\cr}}}{\mathrel{\ooalign{$\scriptscriptstyle\rightarrow$\cr$\rotatebox[origin={c}]{60.0}{$\scriptscriptstyle-$}\mkern
2.0mu$\cr}}}Y,\quad\nexists Z:X\rightarrow Z\leftarrow Y\quad\implies
X\perp\mkern-9.5mu\perp Y\mid\mathcal{V}\setminus\\{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 any
$Z\in\mathcal{V}\setminus\\{X,Y\\}$ if there is a v-structure $X\rightarrow
Z\leftarrow 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.

##  Appendix C Details of the CPDAG Estimation Model

Figure 7: CPDAG Estimation Architecture: In the graph, black edges represent
undirected connections, whereas the red edge signifies an immorality. Columns
in the data matrix corresponding to parent nodes, potential common children,
and neighboring nodes are taken as input for the neural network and are
processed through three distinct feed-forward networks. The embeddings are
concatenated along the variable axis and processed through a network with data
matrix attention, bilinear attention, and a LogEig layer. Mean pooling and
sigmoid activation are applied to output probabilities for potential common
children.

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 $\boldsymbol{x}_{\text{pa}}\in\mathbb{R}^{M\times
2},\boldsymbol{x}_{\text{cc}}\in\mathbb{R}^{M\times|\text{cc}|},\boldsymbol{x}_{\text{ne}}\in\mathbb{R}^{M\times|\text{ne}|}$
of the data matrix $\boldsymbol{X}$ corresponding to
$\text{pa},\text{cc},\text{ne}$ as inputs for the neural network. Each of
these submatrices undergoes a dimensionality expansion to
$\widetilde{\boldsymbol{x}}_{\text{pa}}\in\mathbb{R}^{M\times 2\times
1},\widetilde{\boldsymbol{x}}_{\text{cc}}\in\mathbb{R}^{M\times|\text{cc}|\times
1},\widetilde{\boldsymbol{x}}_{\text{ne}}\in\mathbb{R}^{M\times|\text{ne}|\times
1}$. Next, three separate feed-forward subnetworks
$l_{\text{pa}},l_{\text{cc}},l_{\text{ne}}$ are applied to the three inputs
$\boldsymbol{x}_{\text{pa}},\boldsymbol{x}_{\text{cc}},\boldsymbol{x}_{\text{ne}}$.
Each of these layers has the same architecture: For an input
$\widetilde{\boldsymbol{x}}\in\mathbb{R}^{M\times d\times 1}$, weight matrices
$\boldsymbol{W}_{1}\in\mathbb{R}^{1\times C}$ and
$\widetilde{\boldsymbol{W}}_{1}\in\mathbb{R}^{C\times C}$, together with a
bias vector $\boldsymbol{b}_{1}\in\mathbb{R}^{C}$ are used to embed
$\widetilde{\boldsymbol{x}}$ to

$\boldsymbol{h}_{1}=\tanh\mathopen{}\mathclose{{}\left(\widetilde{\boldsymbol{x}}\boldsymbol{W}_{1}+\boldsymbol{b}_{1}}\right)\widetilde{\boldsymbol{W}}_{1}\in\mathbb{R}^{M\times
d\times C}$

Then, two residual layers of the form

$\boldsymbol{h}_{i+1}=\boldsymbol{h}_{i}+\tanh\mathopen{}\mathclose{{}\left(\boldsymbol{h}_{i}\boldsymbol{W}_{i}+\boldsymbol{b}_{i}}\right)\widetilde{\boldsymbol{W}}_{i}$

with $\boldsymbol{W}_{i}\in\mathbb{R}^{C\times C}$,
$\widetilde{\boldsymbol{W}}_{i}\in{R}^{C\times C}$,
$\boldsymbol{b}_{i}\in\mathbb{R}^{C}$, $i=1,2$ are applied to obtain
representations $\boldsymbol{h}_{\text{pa}}\in\mathbb{R}^{M\times 2\times
C},\boldsymbol{h}_{\text{cc}}\in\mathbb{R}^{M\times|\text{cc}|\times
C},\boldsymbol{h}_{\text{ne}}\in\mathbb{R}^{M\times|\text{ne}|\times C}$. For
the addition of the bias terms, broadcasting is used, i.e.,

$\boldsymbol{b}\in\mathbb{R}^{C}\mapsto\widetilde{\boldsymbol{b}}\in\mathbb{R}^{M\times
d\times C}\text{ with
}\widetilde{\boldsymbol{b}}_{m,l,c}=\boldsymbol{b}_{c}\quad m=1,\dots,M,\text{
}l=1\dots,d,\text{ }c=1,\dots,C.$

Now, the representations are concatenated along the dimension axis to obtain a
$M\times(2+|\text{cc}|+|\text{ne}|)\times C$ tensor. We use
$d=2+|\text{cc}|+|\text{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-pooling

$\boldsymbol{X}\mapsto\left(\frac{1}{d}\boldsymbol{X}^{T(3,2,1)}\boldsymbol{1}_{d}\right)^{T}$

to inflate one of the variable axes of the $\mathbb{R}^{d\times d\times C}$
output of LogEig to obtain a $\mathbb{R}^{d\times C}$ batch of $C$ vectors of
dimension $d$. After applying a dense $C\times 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.

##  Appendix D Parameterization of the SEM

###  D.1 Chebyshev Polynomials for Training

For the Chebyshev polynomial, we utilize the following parameterization:

$\displaystyle f_{v}(\boldsymbol{x}_{\text{pa}},\epsilon_{1},\epsilon_{2})$
$\displaystyle=\sum_{w\in\text{pa}_{\mathcal{G}}\\!(v)}\beta_{w}\sum_{n=1}^{r}\alpha_{n}T_{n}(x_{w})+\epsilon_{1}$
$\displaystyle+\alpha_{\text{m}}\left(\sum_{s,t\in\text{pa}(\mathcal{G}),s<t}\delta_{s,t}T_{\text{m}}(x_{s},x_{t})+\sum_{w\in\text{pa}(\mathcal{G})}T_{\text{m}}(x_{w},\epsilon_{2})\right)\quad\forall
v\in\mathcal{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:

$T_{\text{m}}(x,y):=\frac{(x-\mu_{x})(y-\mu_{y})}{(1+\left|\mu_{x}\right|)(1+\left|\mu_{y}\right|)}$
(2)

where m in the index stands for ”multidimensional”.

Here, $\mu_{x}\sim U[-1,1]$, $\mu_{y}\sim U[-1,1]$, and the coefficients
$\alpha_{1},\dots,\alpha_{r},\alpha_{\text{m}}$ are calculated from
$\alpha_{i}=\frac{\widetilde{\alpha}_{i}}{\sum_{n}\widetilde{\alpha}_{n}+\widetilde{\alpha}_{\text{m}}}$
for $i=1,\dots,r$, with $\widetilde{\alpha}_{i}=\frac{\gamma_{i}}{i!}$,
$\gamma_{i}\sim U[-1,1]$, and $\widetilde{\alpha}_{\text{m}}\sim U[-1,1]$.
$\beta_{w}=\frac{\widetilde{\beta}_{w}}{|\text{pa}|}$ with
$\widetilde{\beta}_{w}\sim U[0.7,1.3]$. $\delta_{s,t}$ are random weights with
$\delta_{s,t}=\frac{\widetilde{\delta}_{s,t}}{\sum_{s,t\in\text{pa}(\mathcal{G}),s<t}|\widetilde{\delta}_{s,t}|}$,
$\widetilde{\delta}_{s,t}\sim U[-1,1]$

This parameterization is motivated by the observation that for a smooth
function—where higher-order derivatives are not significantly larger than the
lower-order ones—the coefficients of the Chebyshev approximation decrease in a
factorial manner, as noted by Xiang and Liu [2020]. Rapidly decreasing
Chebyshev coefficients were also observed by Trefethen [2008]. 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.

While multivariate Chebyshev polynomials constructed using Weyl-Groups were
considered in [Hoffman and Withers, 1988, Puschel and Rotteler, 2007], we
argue that our multivariate terms behave more nicely since they are bounded
between $[-1,1]$, akin to univariate Chebychev polynomials. Their construction
also intuitively incorporates multiplicative effects of two variables in a
randomly shifted way.

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.

###  D.2 Gaussian Mixture Error Terms

The additive error term $\epsilon_{1}$ follows a Gaussian mixture
distribution. We randomly determine the number of components $L\sim
U\\{1,\dots,5\\}$ from a discrete uniform distribution. Each component has
randomly assigned parameters for the means $\mu_{l}\sim U[-1,1]$, standard
deviations $\sigma_{l}\sim U[0.05,1]$, and weights $\widetilde{w}_{l}\sim
U[0.3,1]$, $w_{l}=\frac{\widetilde{w}_{l}}{\sum_{l}\widetilde{w}_{l}}$ such
that $\epsilon_{1}\sim\sum_{l=1}^{L}w_{l}N(\mu_{l},\sigma_{l}^{2})$. The
multiplicative error term $\epsilon_{2}$ is uniformly distributed, with
$\epsilon_{2}\sim U[-1,1]$.

###  D.3 Testing Dependencies

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)$ in

$f_{\text{test}}(\boldsymbol{x}_{\text{pa}},\varepsilon)=\sum_{w\in\text{pa}(\mathcal{G})}\alpha_{w}g(x_{w})+\epsilon.$

For the multi-dimensional multiplicative test dependency, we used

$f_{v}(\boldsymbol{x}_{\text{pa}},\epsilon_{1},\epsilon_{2})=\alpha_{\text{m}}\left(\sum_{s,t\in\text{pa}(\mathcal{G}),s<t}\delta_{s,t}T_{\text{m}}(x_{s},x_{t})+\sum_{w\in\text{pa}(\mathcal{G})}T_{\text{m}}(x_{w},\epsilon_{2})\right)\quad\forall
v\in\mathcal{V},$

with

$T_{\text{m}}(x,y):=\frac{(x-\mu_{x})(y-\mu_{y})}{(1+\left|\mu_{x}\right|)(1+\left|\mu_{y}\right|)}$

with $\mu_{x}\sim U[-1,1]$, $\mu_{y}\sim U[-1,1]$,
$\delta_{s,t}=\frac{\widetilde{\delta}_{s,t}}{\sum_{s,t\in\text{pa}(\mathcal{G}),s<t}|\widetilde{\delta}_{s,t}|}$,
$\widetilde{\delta}_{s,t}\sim U[-1,1]$.

##  Appendix E Training Details

###  E.1 Model Hyperparameters

The implementation was performed using TensorFlow [Abadi et al., 2015]. We
employed the ADAM optimizer by Kingma and Ba [2015]. For training, we used the
hyperparameters stated in Table 3:

Table 2: Hyperparameters for the undirected graph estimation Hyperparameter | Value
---|---
Layer Parameters |
Number of channels $C$ | $100$
Number of inner channels $c$ | $100$
Maximal degree activation function | $3$
Attention heads | 5
Number of layers |
Attention between attributes | 10
Attention between samples | 10
$C\times C$ dense observational layers | 10
bilinear attention + SPD activation | 10
Training Schedule |
epochs | $1000$
samples per epoch | $128$
Initial learning rate | $0.0005$
Learning rate decrease factor | $\left(\frac{1}{10}\right)^{1/500}$
Minibatchsize | 1

Table 3: Hyperparameters of the CPDAG estimation model Hyperparameter | Value
---|---
Layer Parameters |
Number of channels $C$ | $100$
Number of inner channels $c$ | $100$
Maximal degree activation function | $3$
Attention heads | 5
Number of layers |
Attention between attributes | 10
Attention between samples | 10
$C\times C$ dense observational layers | 10
bilinear attention + SPD activation | 10
Training Schedule |
epochs | $1000$
matrices per epoch | $1$
Initial learning rate | $0.0005$
Learning rate decrease factor | $\left(\frac{1}{10}\right)^{1/1000}$
Minibatchsize | 1

Additionally, we generated data with a random number of samples $M\sim
U\\{50,51,\dots,1000\\}$ and a random variable dimension $d\sim
U\\{10,11,\dots,100\\}$.

#### Ablation Study.

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\sim
U\\{50,51,\dots,1000\\}$ and a random variable dimension $d\sim
U\\{10,11,\dots,100\\}$.

###  E.2 Loss Function for the Three-Class Edge Classification Problem.

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,
$\mathcal{L}_{\text{p}}$.

The overall loss function is defined as
$\mathcal{L}_{\text{b}}+\mathcal{L}_{\text{c}}+\mathcal{L}_{\text{p}}$. Here,
$\mathcal{L}_{\text{c}}$ is the categorical crossentropy of the three
categories given by

$\mathcal{L}_{\text{c}}(\boldsymbol{A},\widehat{\boldsymbol{A}}):=H(\boldsymbol{A},\widehat{\boldsymbol{A}})=-\sum_{i=2}^{d}\sum_{j=1}^{i-1}\sum_{c=1}^{3}A_{i,j,c}\log(\widehat{A}_{i,j,c})$

denotes the categorical crossentropy of the three categories _no-edge_ ,
_skeleton edge_ , and _moralized edge_ with

$\boldsymbol{A}\in\mathbb{R}^{d\times d\times 3}\quad\text{with}\quad
A_{i,j,c}=\begin{cases}1&\text{if $(i,j)$ is in category $c$ in the ground-
truth DAG}\\\ 0&\text{else}\end{cases}$

and $\widehat{A}_{i,j,c}$ denotes the estimation by the algorithm on it.
$\boldsymbol{A}$ and $\widehat{\boldsymbol{A}}$ are symmetric along its first
two axes, i.e., $A_{i,j,c}=A_{j,i,c},\quad i,j=1,\dots,d,\quad c=1,\dots,3$.

$\mathcal{L}_{\text{b}}$ denotes the binary loss of no-edge vs. any edge
present (present edges $=$ skeleton edges $\cup$ moralized edges):

$\mathcal{L}_{\text{b}}(\boldsymbol{A}^{(b)},\widehat{\boldsymbol{A}}^{(b)}):=H(\boldsymbol{A}^{(b)},\widehat{\boldsymbol{A}}^{(b)})=-\sum_{i=2}^{d}\sum_{j=1}^{i}\left[A_{i,j}^{(b)}\log(\widehat{A}_{i,j}^{(b)})+(1-A_{i,j}^{(b)})\log(1-\widehat{A}_{i,j}^{(b)})\right]$

with

$\boldsymbol{A}^{(b)}\in\mathbb{R}^{d\times d}\quad\text{with}\quad
A_{i,j}^{(b)}=\begin{cases}1&\text{if an edge is estimated between $i$ and
$j$}\\\ 0&\text{else}\end{cases}$

being the adjacency matrix of no-edge vs. (direct edge $\cup$ moralized edge).

The penalty term, $\mathcal{L}_{\text{p}}$ is defined as:

$\mathcal{L}_{\text{p}}(\widehat{\boldsymbol{A}}):=\max\left(\widehat{\boldsymbol{A}}_{3}-[\widehat{\boldsymbol{A}}_{2}\widehat{\boldsymbol{A}}_{2}]^{0.5},0\right).$

$\widehat{\boldsymbol{A}}_{1},\widehat{\boldsymbol{A}}_{2},\widehat{\boldsymbol{A}}_{3}\in\mathbb{R}^{d\times
d}$ are estimates of $\boldsymbol{A}_{\cdot,\cdot,1}$,
$\boldsymbol{A}_{\cdot,\cdot,2}$, and $\boldsymbol{A}_{\cdot,\cdot,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.

##  Appendix F Further Experiments

###  F.1 Error Bars

In our experiments, each algorithm was evaluated on five distinct trials, each
involving a unique data matrix and corresponding ground-truth graph. The bars
in the figures represent the mean performance values across these trials,
while the error bars indicate the standard deviation. It is important to 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.

###  F.2 Additional Results on Undirected Graph Estimation

Figure 8: AUC values for undirected graph estimation in low-dimensional
regimes: (a) $d=10$, $M=200$ (b)$d=20$, $M=500$. Figure 9: AUC values for
undirected graph estimation for random MLP dependency: (a): $d=20$, (b):
$d=50$, (c): $d=100$

Figures 8 and 9 provide supplemental data on the task of undirected graph
estimation. Figures 8 showcases performance in low-dimensional settings
characterized by $d=10$, $M=200$ and $d=20$, $M=500$. In this scenario, our
method (BAM) also outperforms competing graph inference algorithms. To further
assess its capability to recognize multidimensional dependencies, we extended
our tests to cases where the dependency function within the SEM is modeled via
a randomly initialized multilayer perceptron (MLP) with a random number of
layers $\sim\mathcal{U}\\{1,\dots,5\\}$, a random number of hidden layers
$\sim\mathcal{U}\\{4,64\\}$, and $\operatorname{relu}$ or
$\operatorname{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.

###  F.3 Additional Results on CPDAG Estimation

Figure 10: (a) SHD values for the high-dimensional CPDAG estimation $d=100$,
$M=50$. (b) AUC values for the high-dimensional CPDAG estimation $d=100$,
$M=50$.

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
10 (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 10(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.

###  F.4 Time Comparison

Figure 11 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 $\pm$ 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.

Figure 11: Algorithm runtime for (a) $d=10$, (b) $d=20$, (c) $d=50$, (d)
$d=100$ in seconds per $M\times d$ data matrix inference.

##  Appendix G Interpretation

###  G.1 Shape-Agnostic Architecture and the Role of Attention Layers

When employing a shape-agnostic architecture for matrices
$\in\mathbb{R}^{M\times d}$, it is crucial to ensure that all elements within
the $M\times d$ matrix can interact and influence one another. Consider a
scenario where one axis of the matrix is expanded to the shape
$\mathbb{R}^{M\times d\times 1}$, followed by dense layers with $1\times C$
and $C\times C$ weights. In this configuration, the dense layers carry out
element-wise operations on the $M\times 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\times 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\times M$ and
$d\times d$ attention matrices based on trainable $C\times 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\times 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 [Lin et al., 2013], we
found that the attention mechanism offers a more stable, efficient, and
straightforward computation of non-trainable $M\times M$ and $d\times d$
attention matrices, using only trainable $C\times C$ weights.

Our proposed bilinear attention mechanism is, to our knowledge, the first SPD
layer to enable shape-agnostic computations. It uses trainable $C\times C$
weights to calculate non-trainable $d\times d$ attention matrices, allowing
for adaptive weighting across different $d\times 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\times d$ input.

###  G.2 Attention scores in the BAM layer

Consider the setting as in Figure 3 (right) and the computation of the output
$\boldsymbol{H}$ by $\boldsymbol{A}\otimes\boldsymbol{S}$, where
$\boldsymbol{A}\in\mathcal{S}^{d\times d\times C}_{\succeq}$ are the attention
scores and $\boldsymbol{S}\in\mathcal{S}^{d\times d\times C}_{\succeq}$ are
the input matrices into the BAM layer. Since
$\boldsymbol{A}\otimes\boldsymbol{S}$ is processed parallel across the
channels, we consider for simplicity the output of a single channel here and
assume $\boldsymbol{A}\in\mathcal{S}^{d\times d}_{\succeq}$ and
$\boldsymbol{S}\in\mathcal{S}^{d\times d}_{\succeq}$ to be quadratic, positive
definite $d\times d$ matrices.

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:
$\sum_{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
$\boldsymbol{A}$ in classical attention, the score matrix in bilinear
attention shows the interaction strengths of pair sets $\\{(i,k)\mid
k=1,\dots,d\\}$ and $\\{(l,j)\mid l=1,\dots,d\\}$. The ”receptive field”
adopts a cross-form within the scores $\boldsymbol{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
$\\{\boldsymbol{A}_{i,\cdot}\cup\boldsymbol{A}_{\cdot,j}\\}$.

###  G.3 Keys and Queries

Continuing with the single-channel assumption due to parallel channel
processing, consider the quadratic form $\mathcal{S}^{d\times
d}\times\mathcal{S}^{d\times d}\mapsto\mathcal{S}^{d\times d}$,
$(\boldsymbol{K},\boldsymbol{Q})\mapsto\boldsymbol{K}^{T}\boldsymbol{Q}\boldsymbol{K}$
of the key-query interaction. The $(i,j)$-th entry of
$\boldsymbol{K}^{T}\boldsymbol{Q}\boldsymbol{K}$ is
$\boldsymbol{K}^{T}_{i}\boldsymbol{Q}\boldsymbol{K}_{j}$ for the columns
$\boldsymbol{K}_{1},\dots,\boldsymbol{K}_{d}$ of $\boldsymbol{K}$, which are
often referred to as keys. Using the eigendecomposition of
$\boldsymbol{Q}=\boldsymbol{U}^{T}\boldsymbol{D}\boldsymbol{U}$ one obtains
for the $(i,j)$-th entry the bilinear form
$(\boldsymbol{U}\boldsymbol{K}_{i})^{T}\boldsymbol{D}(\boldsymbol{U}\boldsymbol{K}_{j})$.
Note that $\boldsymbol{U}\boldsymbol{K}_{i}$ is a similarity measure between
$\boldsymbol{K}$ and $\boldsymbol{Q}$ analogous to standard attention. So, for
bilinear attention, similarity scores are calculated between the keys
$\boldsymbol{K}$ and the eigenvectors of the queries $\boldsymbol{Q}$.
Afterwards, the $\boldsymbol{D}^{\frac{1}{2}}$-weighted bilinear-form
$(\boldsymbol{D}^{\frac{1}{2}}\boldsymbol{U}\boldsymbol{K}_{i})^{T}(\boldsymbol{D}^{\frac{1}{2}}\boldsymbol{U}\boldsymbol{K}_{j})$
is used to create covariance matrices by combining the similarity scores
between $\boldsymbol{U}$ and $\boldsymbol{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)\mid k=1,\dots,d\\}$ and
$\\{(l,j)\mid l=1,\dots,d\\}$ collectively influence the output.

This is in contrast to standard attention, which uses the untransformed dot
product
$(\boldsymbol{k},\boldsymbol{q})\in\mathbb{R}^{d}\times\mathbb{R}^{d}\mapsto\boldsymbol{k}^{T}\boldsymbol{q}\in\mathbb{R}$
for the columns of key and query matrices $\boldsymbol{K}$, $\boldsymbol{Q}$.

##  Appendix H Limitations

The model effectively captures smooth dependence relations using Chebyshev
polynomials. Although this approach excels across various types of
dependencies, it might have limitations for data structures that deviate
significantly from the generated synthetic data. However, adapting the
synthetic data generation to accommodate these structures is straightforward.

As for the Log-Eig layer, it performs efficiently within a moderate
dimensional range but may face computational challenges when scaling to higher
dimensions. Training our model with parameters similar to those used in this
study demands substantial memory resources; in our experiments, around 80 GB
of GPU memory was required.

The attention mechanism, while effective, can be costly for high-dimensional
(both, in $M$ and $d$) inputs. This can be ameliorated with local attention,
although this approach may introduce its own set of challenges. While the
model is effective for its intended applications, its architecture allows for
easy extensions. For instance, a separate embedding layer could be added for
both observational and interventional data to make use of interventional data.

An end-to-end approach for CPDAG estimation might offer further benefits. The
current model loses directional information in the covariance computation,
making an end-to-end approach for CPDAG estimation unfeasible with the
existing architecture. However, a simple extension could involve using two
separate embeddings for each variable, one for the variable as parent, one for
the variable as child, effectively doubling the dimensionality to $2d$, to
potentially facilitate directional inference. Lastly, like other neural
network-based approaches, there is a potential risk of overfitting (here on
Chebyshev polynomial dependencies), necessitating hyperparameter tuning.

  *[DAG]: directed acyclic graph
  *[SPD]: symmetric positive semi-definite
  *[NLP]: natural language processing
  *[SEM]: structural equation model
  *[MSA]: multi-sequence alignment
  *[MSAs]: multi-sequence alignment
  *[MLP]: multilayer perceptron