Neural Autoregressive Flows for Markov Boundary Learning

Khoa Nguyen; Bao Duong; Viet Huynh; Thin Nguyen

cs.LGMar 21, 2026

Neural Autoregressive Flows for Markov Boundary Learning

Authors: Khoa Nguyen, Bao Duong, Viet Huynh, Thin Nguyen

Organizations: Deakin Applied Artificial Intelligence Initiative, Geelong, Australia · Edith Cowan University, Perth, Australia

Abstract

Recovering Markov boundary -- the minimal set of variables that maximizes predictive performance for a response variable -- is crucial in many applications. While recent advances improve upon traditional constraint-based techniques by scoring local causal structures, they still rely on nonparametric estimators and heuristic searches, lacking theoretical guarantees for reliability. This paper investigates a framework for efficient Markov boundary discovery by integrating conditional entropy from information theory as a scoring criterion. We design a novel masked autoregressive network to capture complex dependencies. A parallelizable greedy search strategy in polynomial time is proposed, supported by analytical evidence. We also discuss how initializing a graph with learned Markov boundaries accelerates the convergence of causal discovery. Comprehensive evaluations on real-world and synthetic datasets demonstrate the scalability and superior performance of our method in both Markov boundary discovery and causal discovery tasks.

Figures & tables

Fig. 1: Illustration of weight masking mechanism in FANS , as described in Subsection III-B2 , for specific subset

\mathbf{S}=\{X_{2},X_{3},X_{5}\}

. This simplifies the original illustration of MADE architecture for better visualization, showing only a hidden layer with a block of nodes and output layer with a block of nodes. Full autoencoder may consist of multiple hidden layers, each with several blocks of hidden nodes, which can be permuted, as shown in Fig. 1 from [ 15 ] . Solid connections represent fully-connected weight regions in Step 1, while red connections indicate autoregressive masks in Step 2. Other subsets will have different masks. In compact version (right figure), assuming maximum size of observed subsets is 3, each block of hidden layer reduces to 2 nodes. Original element indices of

\mathbf{S}

are rescaled and colliders resolved:

[2,3,5]\rightarrow\left[\left\lceil\frac{2*2}{5-1}\right\rceil,\left\lceil\frac{3*2}{5-1}\right\rceil,\left\lceil\frac{5*2}{5-1}\right\rceil\right]\rightarrow[1,2,3]

Dataset	Nodes	Edges	Observations
Sachs [ 34 ]	11	17	853
SynTReN [ 40 ]	20	24	500
SERGIO [ 10 ]	100	137	9000
ARTH150 [ 28 ]	107	150	5000

TABLE I: Real datasets

Graph type	Metric	IAMB [ 38 ]	S2TMB [ 14 ]	EEMB [ 42 ]	TC [ 32 ]	KMB [ 43 ]	DAT-Moral [ 2 ]	FANSITEMB (ours)
Linear data
d30	nDCG ( $\uparrow$ )	47.18 $\pm$ 12.81	93.01 $\pm$ 3.74	90.78 $\pm$ 6.54	98.91 $\pm$ 1.25	70.58 $\pm$ 7.60	90.00 $\pm$ 3.84	99.34 $\pm$ 1.02
	AveP ( $\uparrow$ )	40.34 $\pm$ 15.41	90.56 $\pm$ 5.11	87.42 $\pm$ 8.93	98.42 $\pm$ 1.81	60.69 $\pm$ 9.57	86.14 $\pm$ 5.17	99.08 $\pm$ 1.39
	F1 ( $\uparrow$ )	48.25 $\pm$ 12.72	93.58 $\pm$ 3.79	91.39 $\pm$ 6.40	98.68 $\pm$ 1.72	65.36 $\pm$ 7.04	21.28 $\pm$ 3.22	99.44 $\pm$ 0.87
d100-1	nDCG ( $\uparrow$ )	41.71 $\pm$ 5.36	94.27 $\pm$ 1.90	96.01 $\pm$ 2.00	99.90 $\pm$ 0.22	$-$	97.42 $\pm$ 1.34	99.86 $\pm$ 0.22
	AveP ( $\uparrow$ )	34.54 $\pm$ 5.37	91.80 $\pm$ 2.74	94.29 $\pm$ 2.88	99.85 $\pm$ 0.34	$-$	96.33 $\pm$ 1.72	99.78 $\pm$ 0.33

TABLE II: Markov boundary discovery performance in % . The numbers are mean ± standard deviation over 5 independent simulations, except for real data. Missing entries (

-

) correspond to instances with runtime exceeding the limit of 10 hours. Bold : best performance, italic : second-best performance.

Nonlinear data
Graph type	Metric	DAGMA [ 3 ]	COSMO [ 24 ]	SDCD [ 26 ]	DAT [ 2 ]	MARVEL [ 25 ]	FANSITE-DCD (ours)
d30-G	SHD ( $\downarrow$ )	15.4 $\pm$ 8.85	25 $\pm$ 5.66	13 $\pm$ 5.43	420.2 $\pm$ 7.50	24.6 $\pm$ 6.23	12.8 $\pm$ 6.53
	AUC-ROC ( $\uparrow$ )	0.752 $\pm$ 0.095	0.822 $\pm$ 0.084	0.858 $\pm$ 0.074	0.502 $\pm$ 0.068	0.603 $\pm$ 0.039	0.902 $\pm$ 0.045
	AUC-PR ( $\uparrow$ )	0.488 $\pm$ 0.172	0.375 $\pm$ 0.077	0.570 $\pm$ 0.144	0.036 $\pm$ 0.013	0.144 $\pm$ 0.061	0.592 $\pm$ 0.098
d30-MN	SHD ( $\downarrow$ )	21.8 $\pm$ 7.82	48.4 $\pm$ 12.46	30.6 $\pm$ 9.81	419.8 $\pm$ 6.10	26.6 $\pm$ 8.91	18.2 $\pm$ 8.14
	AUC-ROC ( $\uparrow$ )	0.654 $\pm$ 0.066	0.777 $\pm$ 0.058	0.766 $\pm$ 0.037	0.488 $\pm$ 0.061	0.578 $\pm$ 0.041	0.866 $\pm$ 0.044

TABLE III: Causal discovery performance . The numbers are mean ± standard deviation over 5 independent simulations, except for real data. Bold : best performance, italic : second-best performance.

Fig. 2: Different model options for estimating our score criterion—conditional entropy. Experiments conducted on synthetic nonlinear graphs with

d=100

. FANS-CPT, FANS-FN, and FANS-FN-NUL represent three variants in our proposed class of masked autoregressive flows (MAFs) tailored for MB discovery, corresponding to: a compact version with leaf masks, a non-compact version with leaf-masks, and a non-compact version without leaf masks. NAF and BNAF are typical MAF classes, while ACE [ 37 ] represents a class of energy-based estimators. The left figure illustrates the reduction across training epochs of

\sum H(X\mid\mathcal{MB}_{X})

, which is the total conditional entropy as each variable conditions on its MB)—smaller values indicate better performance. The right figure shows the number of trainable parameters; fewer parameters imply greater efficiency. When FANS is fully parameterized and trained with leaf-masking (FANS-FN), it learns even better than the compact counterpart. However, the compact FANS , with half the parameters of NAF, is expressive enough to surpass NAF in training and estimation, demonstrating the efficacy of our FANS class, as discussed in Subsection III-B2 .

Causal Structure Learning

Causal structure learning infers directed relationships among variables from observational, interventional, or time-series data.

See all 123 papers

Mask2Cause: Causal Discovery via Adjacency Constrained Causal Attention
May 8, 2026Omar Muhammad, Pasupuleti Dhruv Shivkant, Deepak N. SubramaniTime-Series Causal Discovery Causal Structure Learning
1Indian Institute of Science, Bengaluru, India
Decoupled Causal Discovery
Sep 20, 2026Zhengkang Guan, Fei Wu, Kun KuangCausal Discovery Causal Structure Learning
College of Computer Science and Technology Zhejiang University
Estimating Causal Bounds from Observational Data With Neural Generative Models
Oct 18, 2026Max Zhu, Srivathsan Amruth, Sina Akbari +1Causal Inference Causal Effect Estimation
Department of Pure Mathematics and Mathematical Statistics University of Cambridge

May 8, 2026J/K move · Enter open · S save

Mask2Cause: Causal Discovery via Adjacency Constrained Causal Attention

Omar Muhammad, Pasupuleti Dhruv Shivkant, Deepak N. Subramani

1Indian Institute of Science, Bengaluru, India

Time-Series Causal Discovery Causal Structure Learning Causal Attention

Open paper AskPDF arXiv

Leveraging deep learning for causal discovery in time series remains challenging because existing neural methods predominantly rely on component-wise architectures that fail to capture shared system dynamics or employ decoupled post-hoc graph extraction that risks overfitting to spurious correlations. We propose

\textbf{Mask2Cause}

, an end-to-end framework that recovers the underlying causal graph directly during the forecasting forward pass. Our approach introduces an Inverted Variable Embedding and an Adjacency-Constrained Masked Attention mechanism, trained with homoscedastic or heteroscedastic objectives to capture causal influences in both mean and variance. Empirical results on diverse benchmarks, from synthetic chaotic dynamics to realistic biological simulations, demonstrate state-of-the-art causal discovery with significantly reduced parameter complexity compared to standard baselines. We further show that inferred causal structures can be used to reduce parameter count of forecasting models by more than 70% on average while maintaining predictive accuracy.

Abstract

Causal Structure Learning

Explore similar work

Mask2Cause: Causal Discovery via Adjacency Constrained Causal Attention

Decoupled Causal Discovery

Estimating Causal Bounds from Observational Data With Neural Generative Models