Executing Causal Structure Learning with Linear-Attention Transformers

2026-10-07

AI Summary

This study addresses the challenge that standard neural models struggle to reliably reproduce the optimization trajectories of causal discovery algorithms. To this end, it proposes a fixed-weight linear attention Transformer that precisely executes continuous causal graph updates by retaining Lagrange multipliers. Theoretically, we prove that multiplier retention is essential for exact execution, establish conditions under which rounding errors remain bounded as network depth increases, and decouple algorithmic execution from causal recovery capabilities. Empirically, the proposed model achieves floating-point precision consistency on synthetic data and seven benchmark networks while inheriting the performance of reference solvers. These results validate the feasibility of employing Transformers to exactly execute causal discovery algorithms.

Abstract

Transformers can execute algorithms on data given in their input. We ask whether they can do the same for causal discovery. We study a standard continuous method that repeatedly updates a candidate causal graph while enforcing acyclicity. We explicitly construct a fixed-weight transformer whose forward pass exactly reproduces one update of this method, so repeated blocks reproduce its optimization trajectory. The transformer carries the current graph and the algorithm's multiplier between updates. We show that retaining the multiplier is essential for exact execution, since different multiplier values can lead to different next updates. We also give conditions under which, within a fixed stage, the number of updates needed to reach a target accuracy can be computed in advance and rounding errors stay bounded as depth grows. Experiments show that the constructed block agrees with a reference update to floating-point precision, while arithmetic replay on synthetic data and seven published benchmark network topologies inherits the reference solver's successes and failures. This separates accurate algorithm execution from accurate causal recovery. In contrast, the ordinary attention models tested under our training budgets do not reliably execute the update or transfer to larger graphs. Whether gradient training can learn an executor in the architecture class of the construction remains open.

Problem

• Causal Structure Learning
• Transformers
• Algorithm Execution
• Causal Discovery
• Linear Attention

Innovation

• Causal Structure Learning
• Linear-Attention Transformers
• Algorithm Execution
• Fixed-weight Construction
• Acyclicity Constraint