cs.LGOct 7, 2026
SaveExecuting Causal Structure Learning with Linear-Attention Transformers
Organizations: School of Interdisciplinary Research Indian Institute of Technology Delhi · University of Florida
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.
Figures & tables
Figure 1: One update, run two ways. The solver carries the edge matrix and multiplier ; the fixed-weight transformer carries the same state in . With identical controls, the two transitions agree exactly (Theorem 2 ). The surfaces are schematic.
Figure 2: One update, two implementations. Left: one Poly-prox iteration. Right: the same iteration as the fixed-weight transformer carries it out on the encoded state of ( 5 ), line for line; blue is linear attention, orange bilinear/ReLU, green persistent state. Theorem 2 shows both return the same .
Figure 3: Benchmark topologies (top: ASIA, CANCER, EARTHQUAKE, SURVEY; bottom: SACHS, CHILD, ALARM). Median-SHD draw at , SHD beside each name; red dots are true edges; the column labelled “transformer” shows the direct arithmetic replay of the block. Last column, thresholded at : green kept by both, red missed by both, gold added by both; the black disagreement category is empty on every row.
Figure 4: What the trained models produce. Red dots mark true edges. (1) A softmax model trained at and applied at after updates. (2) At , a collapsed-input linear model withholds , while a softmax model receives them; this illustration changes both architecture and state access. Table 4 separately tests write-back with identical inputs and architectures. (3) Reference-solver updates , and . The right column is the solver’s returned estimate. Colour shows , scaled within each panel; the rows use two illustrative problem instances.
Figure 5: How close small trained transformers get to the exact update. One-step error is normalized so copying scores (dashed line). Circles: trained and tested at ; squares: transfer to . Top: smaller- and larger-corpus one-step errors, then smaller-corpus rollout distances. Bottom: identical state inputs with residual write-back or prediction from scratch; the latter stays above copying. Five training runs are attempted per condition; means and standard errors use finite run summaries (failure counts in Tables 3 and 4 ).
Appendix figures & tables7 assets
Supplementary material from the paper’s appendix.
Appendix
| Protocol | Radius attempts | Passing snapshots | Saved inside | |||
|---|---|---|---|---|---|---|
| Original centres, recorded | 270 | 0/90 | 0 | 0 | 0 | 0 |
| Polished centres, recorded | 725 | 46/90 | 21 | 12 | 13 | 25 |
| Polished centres, new | 636 | 67/90 | 30 | 20 | 17 | 27 |
Table 1: Verified interval fixed-stage tests. Each dimension has snapshots from ten SEM instances; snapshots and radius attempts are not independent experimental units. The polished policies retain one passing ball per successful snapshot. The last column counts saved matching-control endpoints in that policy’s passing ball, not automatically successful continuations of the original adaptive controller. Only and original snapshot states, respectively, lie in the polished passing balls.
| Starting policy | Seed | Snapshot | Certified updates | Final verified distance bound | |
|---|---|---|---|---|---|
| Seeded half-radius perturbation | 3 | 5 | 0 | 134 | |
| Seeded half-radius perturbation | 5 | 0 | 0 | 43 | |
| Seeded half-radius perturbation | 10 | 3 | 0 | 978 | |
| Saved matching-control endpoint | 3 | 8 | 0 | 186 | |
| Saved matching-control endpoint | 5 | 9 | 0 | 426 | |
| Saved matching-control endpoint | 10 | 3 | 0 | 546 |
Table 2: Repeated execution of the constructed block in serial float64 arithmetic. All use the recorded step size and , with controls fixed throughout. Within each starting policy and dimension, the selector prefers the recorded-step policy and selects the smallest positive predicted update count, up to . All six trajectories remain finite, stay inside their verified balls, retain exactly clean scratch/immutable coordinates, and meet the fixed-point distance target by the predicted update count. Final bounds use exact rational a posteriori residuals divided by , or the verified triangle bound through the centre, whichever is smaller.
Figure 6: Accuracy of graph recovery. Lower structural Hamming distance (SHD) means fewer edge mistakes. Each bar averages simulations with observations; error bars show standard errors. On raw data, the continuous solvers and sort-regress outperform PC and GES. Rescaling each variable to unit variance makes sort-regress the worst method at variables. PC and GES are unchanged by this rescaling.
| Model | error | error (zero-shot) | trajectory error, 12 updates |
|---|---|---|---|
| paper-size corpus | |||
| softmax, with memory | |||
| linear, with memory | |||
| linear, no memory (collapsed input) | |||
| learned-threshold unroll | |||
| 10 corpus, 2.5 epochs | |||
Table 3: Trained update executors, mean standard error over five attempted training runs. Step error is normalized so that copying the input scores ; the median over test cases is taken within a seed. A bracketed count gives the seeds whose value was non-finite, which are excluded from that mean; trajectory means are dominated by a few divergent seeds.
| Attention | Residual write-back | error | error | trajectory error |
|---|---|---|---|---|
| softmax (paper-size) | persistent ( ) | |||
| softmax (paper-size) | absent (output from scratch) | |||
| linear (paper-size) | persistent ( ) | |||
| linear (paper-size) | absent (output from scratch) | |||
| softmax (large corpus) | persistent ( ) | |||
| softmax (large corpus) | absent (output from scratch) |
Table 4: Residual write-back with identical inputs, five training runs. “Persistent” writes the update back into the state registers, ; “absent” gives the model the same inputs and the same parameter count but produces and from scratch. One-step errors are normalized so that copying scores ; the trajectory column is the unnormalized after twelve unrolled updates. A bracketed count gives the seeds whose value was non-finite, excluded from that mean. These are the numbers quoted in Section 8 .
Figure 7: Arithmetic replay and conditioning checks. (1) One-update replay error across eight states per dimension. (2) Effective-multiplier growth in the two-node stationary family. (3) Relative gradient perturbations remain below the bound of Corollary 5.4 ; fixed absolute perturbations are amplified as grows. (4) Undamped multiplier error accumulates as while .
Figure 8: Stationarity along saved accepted primal outputs, using each update’s pre-controller controls. At each of normalized progress fractions , trajectory contributes its saved index . Lines are medians and bands are minima/maxima over ten input-seed trajectories per dimension, not confidence intervals. The horizontal axes do not align outer-stage boundaries. The plotted first point is the first primal output; at the separately recorded initial , and . Finite control changes can increase the augmented residual as the constraint and its gradient decrease.