Moment Matters: Mean and Variance Causal Graph Discovery from Heteroscedastic Observational Data

This paper proposes a Bayesian, moment-driven framework that infers separate mean and variance causal graphs from heteroscedastic observational data, supported by new identification theory, variational inference with uncertainty quantification, and curvature-aware optimization to outperform existing baselines in recovering distinct causal structures.

The Gist

Imagine you are a detective trying to figure out how a complex machine works. You can't take it apart; you can only watch it run and record what happens. This is the challenge of Causal Discovery: figuring out what causes what just by watching the data.

Usually, when scientists look at this data, they build a simple map. They draw an arrow from "Rain" to "Wet Grass." This tells them that rain causes the grass to get wet. But it doesn't tell them how the rain changes the grass. Does it make the grass grow taller (change the average)? Does it make the grass wobble more wildly in the wind (change the variability)?

The Problem: The "One-Size-Fits-All" Map
In the real world, things are messy. Sometimes a cause changes the average outcome, and sometimes it changes the variability (how much things bounce around).

Think of a Drug Engineer trying to create a medicine.

• They want the drug to lower a patient's fever (changing the mean).
• But they also want the fever to stay stable, not spike and crash unpredictably (changing the variance).

If the engineer only has a standard map, they see that "Drug A" affects the fever. But they don't know if Drug A is the one that stabilizes the fever or just lowers it. They might try to tweak the wrong part of the machine, wasting time and money. They need a map that separates the "Average Effect" arrows from the "Stability Effect" arrows.

The Solution: A Two-Layer Map
This paper proposes a new way to look at data. Instead of drawing one blurry map, the authors' method draws two distinct maps at the same time:

1. The Mean Map: Shows what causes the average value to change.
2. The Variance Map: Shows what causes the wobble or uncertainty to change.

They call this "Moment Matters." In statistics, "moments" are just fancy words for the average (first moment) and the spread (second moment). The paper argues that to understand complex systems, we must care about both.

How Does It Work? (The Analogy of the Noisy Radio)
Imagine you are listening to a radio station.

• The Mean: This is the volume of the music. If you turn the volume knob, the music gets louder or softer.
• The Variance: This is the static or crackle. If you turn a different knob, the music stays the same volume, but the static gets louder or quieter.

In the past, scientists tried to figure out which knob did what by guessing. Sometimes they got it right; sometimes they got it wrong, especially if the radio was very noisy.

The authors built a smart, Bayesian detective (a computer program using probability).

1. It learns two maps at once: It doesn't guess one map and then try to split it. It learns the "Volume Map" and the "Static Map" simultaneously.
2. It admits uncertainty: Instead of saying "This is definitely the cause," it says, "There is an 80% chance this knob controls the volume, and a 20% chance it controls the static." This is crucial for real-world decisions (like medicine) where being wrong is dangerous.
3. It uses "Curvature-Aware" optimization: Imagine trying to find the bottom of a bumpy valley in the dark. Standard methods might get stuck on a small bump. This new method is like having a hiker who can feel the shape of the ground (curvature) and knows exactly which way to step to avoid getting stuck, making the search much faster and more accurate.

Why Does This Matter?
The paper tested this on fake data, semi-real gene data, and real protein data.

• In Biology: It helped identify which proteins control the average activity of a cell and which ones control the variability (noise) between cells. This is huge for understanding diseases.
• In Economics: It can help figure out what causes an economy to be stable versus what causes it to be volatile.
• In Fairness: It can detect if a decision-making algorithm is treating different groups fairly on average, or if it's creating huge, unpredictable swings in outcomes for specific groups (which is a form of hidden discrimination).

The Bottom Line
This paper gives us a new pair of glasses. Before, we saw the world in black and white (Cause $\to$ Effect). Now, we can see in color, distinguishing between what changes the status quo and what changes the stability. By separating these two forces, we can make better decisions in medicine, finance, and AI, especially when we don't have a lot of data to work with.

Technical Summary

1. Problem Statement

Standard causal discovery methods typically infer a single "moment-agnostic" causal graph, where an edge $X_i \to X_j$ implies that $X_i$ causes a change in $X_j$ without specifying which statistical moment (mean, variance, higher-order moments) is affected. However, real-world data often exhibits heteroscedasticity, where the variance of a variable changes depending on its causes.

In complex systems (e.g., drug discovery, systems biology, economics), different causes may influence the mean (expected value) and the variance (variability) of a target variable differently. A standard graph fails to distinguish these mechanisms, limiting interpretability and the design of targeted interventions. The paper addresses two key questions:

1. Can we separately identify the mean causal graph ($G_M$) and the variance causal graph ($G_V$) solely from observational data?
2. How can we perform this inference with principled uncertainty quantification, especially under data scarcity?

2. Methodology

The authors propose a Bayesian, moment-driven causal discovery framework based on a new class of Structural Causal Models (SCMs).

A. Model Formulation: Mean-Variance HNM

The authors define a Mean-Variance Heteroscedastic Noise Model (HNM). For each variable $X_j$, the structural equation is:
$$X_j = m_j(\mathbf{X}_{pa_M(j)}) + v_j(\mathbf{X}_{pa_V(j)}) E_j$$
Where:

• $\mathbf{X}_{pa_M(j)}$ and $\mathbf{X}_{pa_V(j)}$ are subsets of parents influencing the mean and variance, respectively.
• $m_j(\cdot)$ is the mean function.
• $v_j(\cdot)$ is the variance function (strictly positive).
• $E_j$ is a zero-mean noise term with constant variance.
• The Mean Causal Graph ($G_M$) contains edges from $\mathbf{X}_{pa_M(j)}$ to $X_j$.
• The Variance Causal Graph ($G_V$) contains edges from $\mathbf{X}_{pa_V(j)}$ to $X_j$.
• The union of these graphs forms the moment-agnostic graph $G$.

B. Theoretical Identifiability

The paper establishes sufficient conditions under which $G_M$ and $G_V$ are uniquely identifiable from the observational distribution $P(\mathbf{X})$:

1. Causal Sufficiency & Minimality: Standard assumptions (no unobserved confounders, no redundant edges).
2. DAG & Shared Permutation: Both $G_M$ and $G_V$ are Directed Acyclic Graphs (DAGs) and share the same topological ordering (permutation). This ensures their union $G$ is also a DAG.
3. Function Constraints:
   • $m_j$ must be nonlinear.
   • $v_j$ must be piecewise but non-constant.
   • Noise $E_j$ must be Gaussian.
     Under these conditions, the paper proves that the separation of mean and variance parents is theoretically possible, distinguishing this from standard HNMs where only the union graph is identifiable.

C. Bayesian Inference Framework

To infer the posterior distribution over $G_M$ and $G_V$ from finite data, the authors develop a Variational Inference (VI) approach:

• DAG Distribution: They model the posterior over adjacency matrices $A_M$ and $A_V$ using a shared permutation matrix $\Pi$ to enforce the DAG constraint and the shared ordering assumption.
  • $A_M = \Pi^\top U_M \Pi$ and $A_V = \Pi^\top U_V \Pi$, where $U_M, U_V$ are upper-triangular.
• Differentiable Sampling: To enable gradient-based optimization, they use Gumbel-Softmax for sampling binary upper-triangular matrices and SoftSort (from DDS) for sampling the continuous relaxation of the permutation matrix.
• Likelihood: The likelihood is modeled using Multi-Layer Perceptrons (MLPs) for $m_j$ and $v_j$, with masking applied based on the sampled adjacency matrices.
• Optimization Strategy:
  • Curvature-Aware Optimization: To address the ill-posedness of Maximum Likelihood Estimation (MLE) in heteroscedastic models (where gradients scale inversely with variance), they alternate between updating mean-related parameters (using MSE gradients, effectively scaling by variance) and variance-related parameters.
  • Prior Incorporation: To handle small sample sizes, they incorporate domain knowledge regarding node orderings (e.g., "Gene A regulates Gene B") by projecting the permutation parameters onto a feasible set defined by pairwise constraints using Quadratic Programming (QP).

3. Key Contributions

1. New SCM Class & Identifiability Theory: Proposes the Mean-Variance HNM and derives the first theoretical conditions for the separate identifiability of mean and variance causal graphs.
2. Bayesian Uncertainty Quantification: Develops a variational inference framework that outputs a posterior distribution over both graphs, allowing for the calculation of edge probabilities and structural uncertainty (crucial for decision-making).
3. Robust Optimization: Introduces a curvature-aware optimization scheme and a prior incorporation technique for node orderings to improve sample efficiency and stability in heteroscedastic settings.
4. Empirical Validation: Demonstrates state-of-the-art performance on synthetic, semi-synthetic (gene expression), and real-world (protein signaling) datasets.

4. Experimental Results

The method was evaluated against baselines including Bayesian ANM methods (MC3, DDS) and point-estimation HNM methods (ICDH, HOST).

• Synthetic & Semi-Synthetic Data:
  • The proposed method significantly outperformed baselines in recovering both $G_M$ and $G_V$ (measured by Structural Hamming Distance and F1 score).
  • Baselines like MC3 and DDS failed under heteroscedasticity due to model misspecification (assuming additive noise).
  • Point-estimation methods (ICDH, HOST) struggled to identify the variance graph, often relying on mean-structure cues.
  • The method remained robust even under non-Gaussian noise (Laplace, Student-t), outperforming Gaussian-based baselines.
• Real-World Data (Sachs Dataset):
  • On protein signaling data ($d=11$), the method achieved comparable performance to state-of-the-art moment-agnostic methods (HOST) for the union graph.
  • Case Study: Successfully identified the edge $MEK \to ERK$ as a variance-controlling relationship with high posterior probability, aligning with known biological literature regarding cell-to-cell variability.
  • Prior Knowledge: Incorporating 25-50% of known node orderings improved accuracy, particularly in low-sample regimes ($n=100$), with only a modest increase in runtime.

5. Significance

This work bridges a critical gap in causal discovery by moving beyond "moment-agnostic" reasoning. By explicitly modeling and separating the drivers of mean and variance:

• Interpretability: It reveals how causes affect outcomes (e.g., shifting the average vs. increasing volatility), which is vital for fields like pharmacology (stabilizing drug effects) and algorithmic fairness (identifying sources of outcome disparity).
• Decision Support: The Bayesian framework provides uncertainty estimates, enabling safer interventions in high-stakes domains like medicine where data is scarce.
• Theoretical Advancement: It extends the scope of identifiable SCMs to include higher-order moment structures, paving the way for future research into skewness and kurtosis causal graphs.