{"id":"b786d8e2-5a48-41e4-9055-ae13a7b773b3","arxiv_id":"2506.18285","paper_version":1,"verdict":"CONDITIONAL","confidence":"MODERATE","novelty_score":6.0,"correctness_risk":"medium","formal_verification":"none","parameter_count":3,"one_line_summary":"ADAG pre-trains a linear transformer to map observed data from many related tasks directly to DAG adjacency matrices, enabling fast zero-shot causal discovery on new order-consistent or heterogeneous datasets.","lead":"This paper trains a transformer-based model to predict causal graph structures from many small related datasets, then applies it instantly to new datasets. If it works, this could make causal discovery much faster and more reliable when data per task is scarce.","discovery_kind":"new_application","skeptic_critique":{"model":"deepseek-v4-flash","headline":"The d=5 zero-shot experiments may be confounded by structure memorization: with a fixed order there are only 2^10 possible DAGs, and 70k training domains cover essentially all of them.","rationale":"The paper has genuine merits: a novel attention-based formulation for jointly estimating multiple linear SEMs, self-supervised training without ground-truth graphs, and a clear statement of the scoped settings (Section 3.1, footnote 1). The central claim is that a single pre-trained model zero-shot infers DAGs for new order-consistent or heterogeneous domains. The strongest quantitative support is at d=5 (SHD 0.027-0.038 in Table 1, and the small-sample results in Table 2), but the d=5 evaluation is confounded by the finiteness of the structure space. With a fixed order, only 1024 DAG structures exist; 70,000 training domains essentially exhaust this set. Consequently, the test domains are not structurally unseen, so the d=5 numbers can be explained by memorization of the structure set rather than by a generalizable kernel map. The heterogeneous setting is even more restrictive: the test DAG is identical to the training DAG by design, so structural generalization is not tested at all. The d=10/20 results are not vulnerable to this objection because the structure space is 2^45, making train/test structural overlap negligible; those results do support a softer version of the claim, namely zero-shot inference within the same topological order and same ER1 generative protocol. The proposed split-by-structure test would distinguish memorization from genuine generalization. If the d=5 performance degrades under a structure-disjoint split, the headline and small-sample claims must be re-scoped; if it does not, the concern is resolved. This does not invalidate the paper, but it does mean the reader's CONDITIONAL verdict is the right call: the central claim should be accepted only with the explicit requirement of a structure-disjoint evaluation and a discussion of the finite-structure confound at d=5.","tokens_in":17606,"tokens_out":14605,"duration_ms":156753,"concrete_test":"Re-run the order-consistent d=5 experiment with a strict structure split: enumerate all 1024 DAGs consistent with the training order, hold out 20% of these structures for testing, generate the 70,000 training domains only from the remaining 80%, and evaluate on held-out structures. If the d=5 SHD rises substantially (e.g., from 0.038 toward the NOTEARS or MultiDAG level), the current zero-shot claim is an artifact of structure memorization. For the heterogeneous claim, additionally test on a DAG structure that was never used during training; since the current protocol uses one shared structure, this directly tests whether the model can infer a new graph rather than recognize a memorized one.","verdict_should_be":"UNCHANGED","load_bearing_attack":"The order-consistent d=5 evaluation in Table 1 and Table 2 is presented as zero-shot generalization to unseen DAGs, but the structure space consistent with a fixed order on 5 nodes has only 2^(d(d-1)/2)=1024 members. The training protocol uses 70,000 domains (Section 4.2 and Appendix C.2), so every one of the 1024 structures appears in training many times. The 'held-out' test domains therefore introduce new data and new edge weights, but almost never a new DAG structure; the model can memorize the mapping from data statistics to each of the 1024 structures and only estimate weights for the new domain. The heterogeneous setting is even more explicit: by the footnote in Section 2 and the dataset description in Section 4, test domains are generated from the same DAG as training, so no new structure is ever introduced. Thus the d=5 results (SHD 0.027-0.271) do not establish zero-shot inference on genuinely new DAGs. The d=10/20 results (SHD 0.075-0.217) are less affected because 2^45 possible structures make train/test structural overlap negligible, but the paper's flagship small-sample experiments are all d=5.","agreement_with_reader":"partial"},"referee_report":{"model":"deepseek-v4-flash","summary":"The manuscript proposes Attention-DAG (ADAG), a linear-transformer-based model that learns a nonlinear kernel mapping from observational data matrices to weighted adjacency matrices for linear structural equation models. ADAG is pre-trained across many synthetic domains by minimizing a reconstruction loss subject to an acyclicity constraint, then applied to held-out domains in a zero-shot fashion. The paper targets two multi-domain settings: heterogeneous data, where all domains share the same DAG structure with varying edge weights, and order-consistent data, where domains have different DAGs that share a common topological order. The experiments report strong accuracy and very fast inference on synthetic ER1 graphs with d=5,10,20, additional low-sample robustness experiments, an ablation on the number of training domains, non-Gaussian noise ablations, and a Sachs-dataset transfer experiment. The central claim is that a single pre-trained model can zero-shot infer DAGs for new order-consistent or heterogeneous domains and that the pre-trained kernel captures a shared low-dimensional prior that mitigates small-sample ill-posedness.","tokens_in":17875,"tokens_out":4861,"duration_ms":51933,"significance":"If the zero-shot generalization claim is established, ADAG would be a genuinely useful contribution: it would be among the first practical foundation-model-style approaches to DAG learning, with the attractive property that inference at test time is a single forward pass rather than iterative optimization. The paper also has concrete strengths: the training objective is unsupervised and, in principle, does not require ground-truth graphs; the authors provide extensive implementation details, ablation studies on domain count and noise type, and low-sample experiments; and the reported inference speed is orders of magnitude faster than optimization-based baselines. However, the evidence as presented does not fully support the central generalization claim because the experimental design allows structure-level overlap between training and test domains in the d=5 order-consistent setting, and by construction in the heterogeneous setting. The comparison with baselines also conflates the pre-training information advantage with algorithmic superiority.","major_comments":[{"comment":"The zero-shot comparison does not test genuinely unseen DAG structures. With a fixed topological order on d=5 nodes, the number of possible DAGs is 2^(d(d-1)/2)=1024, while the training protocol uses M=70,000 domains (Appendix C.2). Since all order-consistent domains share the same order (footnote in §2 and dataset settings in §4), every one of the 1024 possible structures appears many times during pre-training. The held-out test domains therefore introduce new data and new edge weights, but almost never a new DAG structure. The SHD of 0.038 in Table 1 and the low-sample results in Table 2 consequently measure within-distribution structure recognition (or memorization) rather than zero-shot generalization to unseen DAGs. This is load-bearing for the paper's central claim, and the d=10/20 results are less affected but do not repair the flagship d=5 experiments. To support the claim, report evaluations on structures excluded from training, with a train/test split performed at the level of graph structures.","section":"§4.2, Table 1 (d=5, order-consistent)"},{"comment":"In the heterogeneous experiments, the test domains are generated from the same DAG as the training domains, as stated in the footnote in §2 and in §4. Thus the heterogeneous setting never presents a new graph structure to the model; it only varies the edge weights. The SHD values of 0.027-0.271 in Table 1 show that the learned kernel can adapt to new causal mechanisms for a fixed structure, but they do not demonstrate structural generalization. The conclusion's statement that ADAG 'generalize[s] effectively to test domains with unseen DAGs and mechanisms' is therefore not supported by the heterogeneous experiments, and the abstract's 'unseen tasks' wording should be qualified accordingly.","section":"§4.1-4.2, heterogeneous setting"},{"comment":"The comparison against baselines conflates the pre-training information advantage with algorithmic superiority. ADAG is pre-trained on 70,000 domains generated from the same generative protocol as the test domains, while NOTEARS and DAGMA are applied independently to each test domain with no access to training domains, and MetaDAG and MultiDAG are run only on the 1000 test domains. The large apparent gains in SHD, relative error, and runtime therefore include the benefit of having seen 70,000 domains from the target distribution. This is the intended advantage of a foundation model, but the paper's phrasing 'consistently outperforming all state-of-the-art baselines' should be contextualized, and the claim of a fair comparison in §4.2 does not explain why the baselines were not given an equivalent multi-domain training stage.","section":"§4.2, baseline comparison"},{"comment":"The paper's claim that pre-training on multiple domains reduces the ill-posedness of small-sample DAG learning is not yet backed by a proof for the regime considered. Appendix D.1 explicitly states that the identifiability conditions from Chen et al. (2021) may be violated when n is comparable to d, and that a relaxation is left to future work. Similarly, Lemma D.1 provides an RKHS identifiability statement for the parameter matrix, but it is not shown how the attention-based architecture implements or approximates this RKHS. The empirical low-sample results in Table 2 are valuable, but without a proof or a precise statement of the conditions under which ADAG improves identifiability, the theoretical claim in the abstract and in §3.3 should either be removed or clearly labeled as a conjecture.","section":"§3.3 and Appendix D"}],"minor_comments":[{"comment":"The typo 'ADGA' should be corrected to 'ADAG'.","section":"§4.3, second paragraph"},{"comment":"The caption says there are 'five dominant eigenvalues consistent with the structure of an ER1 graph with d=5', but the plotted eigenvalues appear to be from the weighted adjacency matrices, not from the data covariance. Please clarify in the caption and the text what object is being decomposed.","section":"§4.1, Figure 2"},{"comment":"The runtime of MultiDAG differs by two orders of magnitude between the heterogeneous setting (0.0151 s) and the order-consistent setting (1.9911 s), which is a surprising discrepancy given the same problem size; please explain this difference.","section":"Table 1, d=5 row"},{"comment":"A fixed threshold of 0.3 is applied to the inferred weighted adjacency matrices for all methods and all settings. Since the scale of ADAG's kernel output may vary with architecture and data, a sensitivity analysis for this threshold would strengthen the empirical evaluation.","section":"§4.2, threshold choice"},{"comment":"The Sachs evaluation uses a pre-training set generated from an assumed ordering of the Sachs variables plus synthetic DAGs with known ground truth. This is an oracle-assisted transfer setting, and the main text should state that the ordering assumption is used in the pre-training stage; otherwise readers may over-interpret Table 3 as a fully unsupervised real-data result.","section":"Appendix C.1, Table 3"}],"recommendation":"major_revision","confidential_remarks":"The main issue is structural leakage in the d=5 and heterogeneous experiments, which directly affects the zero-shot claims. If the authors can add structure-level held-out experiments for order-consistent data and reframe the heterogeneous setting as weight transfer rather than structure transfer, the paper could become publishable. The paper is otherwise well organized and the experiments are reasonably detailed for a first submission."},"author_rebuttal":null,"desk_editor":{"model":"deepseek-v4-flash","letter":"Colleague,\n\nThe one thing to know: ADAG is a genuinely new way to do multi-task DAG learning—an attention-based map from data matrices to weighted adjacency matrices trained with reconstruction and acyclicity constraints—and it is fast. But the headline d=5 zero-shot results are probably inflated by structure memorization. With a fixed topological order on 5 nodes there are only 2^10=1024 possible DAGs, and the paper trains on 70,000 domains drawn from the same ER1 protocol. Almost every test structure is certainly in training. So the d=5 'zero-shot' numbers mostly measure the model's ability to fit edge weights for a known structure, not to generalize to genuinely new graphs. The heterogeneous setting is worse: by the paper's own definition, test domains share the same DAG as training, so zero-shot generalization there means new mechanisms, not new structure. The d=10 and d=20 results are less affected because the structure space is enormous, so those are more informative.\n\nWhat deserves credit: the architecture is new, the training objective is sensible, and the empirical work is fairly extensive for a preprint. The model hits SHD close to zero and runs in milliseconds. The low-sample robustness story is plausible, and the non-Gaussian noise ablation adds scope. The paper is honest about its limited setting and openly leaves the small-sample identifiability theory to future work. That is refreshing.\n\nThe soft spots beyond memorization: the baseline comparison is not apples-to-apples. The baselines (NOTEARS, DAGMA, MetaDAG, MultiDAG) receive no equivalent pre-training data, so part of the reported gain is just having seen the distribution. The paper mentions three repeats but only gives standard deviations in one appendix table. No code is released, which matters for a paper whose contribution is a trained model. The theoretical appendix leans on prior results and the small-sample case—the one that motivates the paper—is explicitly left open.\n\nMy take: the central idea is worth taking seriously, but the 'foundation model' claim needs a narrower, more honest statement. This paper deserves a serious referee, but it needs major revision before acceptance: fairer baselines, error bars, code, and either new experiments on genuinely unseen structures or a clear reframing of what zero-shot means here. I'd probably cite it for the architecture, not for the benchmark claims.\n\nRecommendation: send to peer review, with the caveat that the memorization issue must be addressed.","headline":"Genuinely new architecture and fast inference, but the d=5 zero-shot numbers look like structure memorization; worth a serious referee, with major revisions.","tokens_in":18388,"tokens_out":2851,"would_cite":true,"duration_ms":27173,"reading_group":"maybe","serious_thinker":"yes","would_accept_peer_review":true},"rs_alignment":null,"lean_confirmation":null,"pith_extraction":{"msc":["68T07","62H22"],"pacs":[],"model":"deepseek-v4-flash","headline":"The paper claims that ADAG, an attention network pre-trained on 70,000 synthetic domains, learns a nonlinear map from observed data to weighted causal graphs and infers new DAGs zero-shot with near-zero structural error.","keywords":["causal discovery","DAG learning","foundation model","attention mechanism","linear structural equation model","zero-shot inference","multi-task learning","small-sample regime"],"falsifier":"Hold the training protocol fixed (ER1, $d=5$, equal-variance Gaussian noise, edge weights in $[-2,-0.5]\\cup[0.5,2]$) but reverse the topological order used to generate test-domain DAGs, or switch test graphs to ER2 with two edges per node, then measure zero-shot SHD of the pre-trained ADAG. If SHD jumps from the reported $\\sim0.03$ toward the $\\sim1$–$3$ range of per-task baselines, the generalization claim is confined to the training distribution rather than being a general causal prior.","tokens_in":17368,"feed_emoji":"🔗","tokens_out":12305,"duration_ms":101859,"temperature":0.7,"pith_summary":"This paper seeks to establish that causal DAG discovery can be pre-trained once and then applied to new problems without per-task optimization. The proposed model, ADAG, uses stacked attention blocks to learn a nonlinear kernel that maps the observed data matrix of a linear structural equation model directly to the weighted adjacency matrix encoding both graph structure and edge coefficients. In two settings that share structure across tasks—heterogeneous data (same DAG, varying edge weights) and order-consistent data (different DAGs, one shared causal order)—a model pre-trained on 70,000 synthetic domains recovers unseen test graphs with near-zero structural error, e.g., SHD 0.027 for five variables with heterogeneous data, in about 0.0003 seconds per domain. The larger claim is that multi-task pre-training discovers a shared low-dimensional generative prior that mitigates the ill-posedness of small-sample causal discovery. If correct, this points toward reusable causal-discovery models for domains where data are scarce.","feed_headline":"Pre-trained attention model infers causal DAGs in a millisecond","feed_subtitle":"One forward pass recovers unseen graphs with near-zero structural error, even from only 25 samples per task.","key_machinery":"The central object is the attention-based kernel map. Data from one domain are tokenized as $H^{(0)} = X \\in \\mathbb{R}^{d\\times n}$, one token per variable; $L$ linear-attention layers compute $$$H^{{(\\ell)}}$ = \\operatorname{Attn}[$H^{{(\\ell-1)}}$]$H^{{(\\ell-1)}}$ + $H^{{(\\ell-1)}}$,\\quad \\operatorname{Attn}[H] = \\$\\sigma$\\bigl(\\tfrac{1}{\\sqrt{d}} H W_Q W_K^\\top H^\\top\\bigr)\\in\\mathbb{R}^{d\\times d},$$ and the final layer outputs $A = W_{\\mathrm{out}}^{P,x}\\,\\sigma\\bigl(\\tfrac{1}{\\sqrt{d}} H^{(L)} W_Q^{\\mathrm{out}}(W_K^{\\mathrm{out}})^\\top (H^{(L)})^\\top\\bigr)$. The attention matrix is exactly the weighted adjacency matrix: entry $(i,j)$ is the strength of the proposed edge $X_i\\to X_j$. Augmented-Lagrangian training with the trace-exponential acyclicity constraint is what makes the unsupervised output a DAG rather than an arbitrary matrix.","core_discovery":"ADAG claims that a nonlinear attention-based kernel map $A[X(1:n); \\Theta]$ can take the $n$ observed samples of $d$ variables and output the $d \\times d$ weighted adjacency matrix of the underlying linear SEM, $X = A^\\top X + E$. The map is trained by minimizing the reconstruction loss $\\sum_\\eta \\|X^\\eta - (A^\\eta)^\\top X^\\eta\\|_F^2$ across $M$ domains subject to the acyclicity constraint $h(A^\\eta)=\\operatorname{tr}(e^{A^\\eta\\circ A^\\eta})-d=0$, using augmented Lagrangian iterations. After pre-training on $M=70{,}000$ ER1 domains for $d\\in\\{5,10,20\\}$ with equal-variance Gaussian noise, the learned kernel performs zero-shot inference on 1,000 held-out domains: SHD 0.027 (heterogeneous, $d=5$) and 0.038 (order-consistent, $d=5$), with relative adjacency error around 0.12 and inference runtime 0.0003 seconds, outperforming both single-task optimizers and multi-task learners in accuracy and speed. The paper further claims that the kernel captures the common low-dimensional structure of the training distribution and that this prior keeps performance relatively stable when each domain has only $n=25$ or $n=50$ observations.","pith_inferences":["Editorial: The benchmark only tests same-protocol generalization. A decisive next experiment is to train on ER1 and zero-shot test on ER2 graphs or on unequal-variance Gaussian noise; low SHD there would support the shared-prior story, while collapse to baseline SHD would bound the claim to the training generator.","Editorial: Because each variable is one token and the output is $d\\times d$, the architecture does not obviously scale beyond the tested $d\\le 20$ regime; adapting it to hundreds of variables would likely require latent grouping or hierarchical attention.","Editorial: ADAG's training loss is unsupervised, so labeled or interventional samples, when available, could be added as a fine-tuning signal; the paper does not explore this, but it is a natural route to relax the equal-variance Gaussian identifiability dependence.","Editorial: The theoretical section relies on prior identifiability results and explicitly leaves the small-sample joint identifiability proof to future work; if the empirical prior effect is real, it suggests the sample-complexity conditions of the order-recovery result can be relaxed."],"forward_implications":["A model pre-trained under ADAG's protocol can replace per-task DAG optimization: a new domain in the same structural family is solved by one forward pass, cutting inference time from seconds (single-task baselines) or minutes (multi-task baselines) to fractions of a millisecond.","Small-sample performance shifts: at $n=25$ or $n=50$, ADAG keeps SHD roughly 0.05–0.27 for $d=5$, while single-task baselines degrade to SHD 0.54–1.56, so the shared prior behaves like regularization.","A model trained on $n=100$ transfers to test domains with only 25 or 50 observations when samples are duplicated to the training length, supporting the claim that high-resource pre-training can serve low-resource downstream tasks.","The learned kernel's covariance matrix and dominant eigenvectors match the ground-truth adjacency structure, making the shared prior visible rather than a black-box effect.","The same architecture extends in principle to nonlinear SEMs by feeding the last hidden representation $H^{(L)}$ into the reconstruction loss instead of raw $X$, a direction the paper states as a direct extension."],"supporting_citations":[{"why":"Supplies the continuous acyclicity constraint $h(A)=\\operatorname{tr}(e^{A\\circ A})-d$ and the optimization framing that ADAG's reconstruction loss and augmented-Lagrangian training build on; also provides the NOTEARS baseline.","marker":"Zheng et al. [2018]"},{"why":"Defines the order-consistent multi-task setting and its data-generation protocol, provides the MultiDAG baseline, and contributes the shared-order identifiability argument the paper adapts.","marker":"Chen et al. [2021]"},{"why":"Defines the heterogeneous multi-task setting (shared graph, varying mechanisms) and its data-generation protocol, and provides the MetaDAG bilevel-optimization baseline.","marker":"Lu and Gao [2023]"},{"why":"Gives the equal-variance Gaussian identifiability result the paper invokes to claim each domain's DAG is uniquely recoverable under the benchmark noise model.","marker":"Peters and Bühlmann [2014]"},{"why":"Provides the non-Gaussian additive-noise identifiability result (LiNGAM) the paper cites when extending its identifiability argument beyond Gaussian noise.","marker":"Shimizu et al. [2006]"},{"why":"Provides the mechanism the paper relies on for its central conceptual claim: linear transformers discover low-dimensional cross-task priors that regularize ill-posed inverse problems.","marker":"Lu and Yu [2025]"},{"why":"Supplies the real-world protein-signaling dataset used in the paper's real-data evaluation protocol.","marker":"Sachs et al. [2005]"}],"fun_headline_variants":["ADAG zero-shot causal DAG inference in 0.3 ms","Pretrained attention model recovers unseen DAGs","From 25 samples to accurate causal graphs via ADAG","Foundation model speeds causal discovery to milliseconds"],"cache_read_input_tokens":3200,"weakest_assumption_plain":"The load-bearing premise is that any new task comes from the same generative recipe used for the 70,000 training domains: the same number of variables, the same sparse graph density, the same edge-weight range, equal-variance Gaussian noise, and either the same DAG (heterogeneous) or the same topological order (order-consistent). Outside that recipe the trained kernel has no stated adaptation mechanism.","fun_headline_variants_meta":{"raw":{"variants":["ADAG zero-shot causal DAG inference in 0.3 ms","Pretrained attention model recovers unseen DAGs","From 25 samples to accurate causal graphs via ADAG","Foundation model speeds causal discovery to milliseconds"]},"model":"deepseek-v4-flash","effort":"low","cost_usd":0.000246,"raw_usage":{"total_tokens":1614,"prompt_tokens":1094,"completion_tokens":520,"prompt_tokens_details":{"cached_tokens":384},"prompt_cache_hit_tokens":384,"prompt_cache_miss_tokens":710,"completion_tokens_details":{"reasoning_tokens":455}},"tokens_in":710,"tokens_out":520,"duration_ms":5413,"temperature":1.0,"reasoning_tokens":455,"cache_read_input_tokens":384,"cache_creation_input_tokens":0},"cache_creation_input_tokens":0},"created_at":"2026-08-15T18:52:33.930610+00:00","model_set":{"reader":"deepseek-v4-flash"},"falsifier":"Hold the training protocol fixed (ER1, $d=5$, equal-variance Gaussian noise, edge weights in $[-2,-0.5]\\cup[0.5,2]$) but reverse the topological order used to generate test-domain DAGs, or switch test graphs to ER2 with two edges per node, then measure zero-shot SHD of the pre-trained ADAG. If SHD jumps from the reported $\\sim0.03$ toward the $\\sim1$–$3$ range of per-task baselines, the generalization claim is confined to the training distribution rather than being a general causal prior.","supporting_citations":[{"cited_title":"Meta-dag: Meta causal discovery via bilevel optimization","cited_arxiv_id":null,"evidence_quote":"Defines the heterogeneous multi-task setting (shared graph, varying mechanisms) and its data-generation protocol, and provides the MetaDAG bilevel-optimization baseline."}],"review_version":2}