Stabilizing Causal Structure Learning under Heteroscedasticity: Analysis and Mitigation of Optimization Failures

Abstract

This study focuses on learning causal directed acyclic graphs (DAGs) under heteroscedastic noise models (HNMs), where each effect is modeled as a function of its causes and a Gaussian noise term whose variance depends on the causes. While HNMs theoretically guarantee identifiability of causal structures, we show that gradient-based continuous DAG learning can fail in practice due to an adverse interaction between heteroscedastic likelihood optimization and the acyclicity constraint. Specifically, because the reconstruction gradient is scaled by the predicted variance, it can be heavily attenuated in early training; as a result, the DAG parameters may be updated primarily by the acyclicity constraint before the data reconstruction signal is sufficiently learned, hindering effective structure learning. We identify and formalize this failure mode. To mitigate it, we propose a graduated optimization strategy based on a surrogate loss that decouples the variance term from the reconstruction loss, thereby preventing early gradient attenuation. We further introduce a scheduling coefficient that initially assigns a high weight to the surrogate loss for stable mean learning, and then gradually transitions to the full heteroscedastic likelihood to refine variance estimates and strictly enforce acyclicity. This strategy avoids the identified failure mode and guides the learned DAG to better reflect the data. Experimental results on both synthetic and real-world datasets verify the effectiveness of our approach. Our Github repository including code and supplementary material is here: https://github.com/Sinegi/HNM.