Causality-Aware Transformer (CAT)

• CAT is a transformer-based model that first trains a multi-layer forecaster to capture long-range, nonlinear dependencies and then extracts directed causal graphs using gradient sensitivity analysis.
• It integrates user-specified excluded links through an attention-masking mechanism, effectively suppressing spurious causal relationships and enforcing prior knowledge constraints.
• Empirical evaluations demonstrate a 12.8% absolute improvement in F1-score and 98.9% accuracy in estimating causal lags compared to state-of-the-art methods.

Searching arXiv for the specified CAT paper and closely related transformer-based temporal causality work. Causality-Aware Transformer (CAT) most commonly denotes a transformer-based framework for temporal causal discovery and inference that couples a multi-layer time-series forecaster with post-training gradient analysis and an attention-masking mechanism for prior knowledge integration. In the formulation introduced in "Transforming Causality: Transformer-Based Temporal Causal Discovery with Prior Knowledge Integration," CAT is designed to address two stated difficulties in temporal causal analysis: complex nonlinear dependencies and spurious correlations. Its workflow is to train a forecaster on historical multivariate data, extract a directed causal graph and associated lags from the trained model through gradient-based sensitivity analysis, and optionally constrain attention with user-excluded links so that forbidden causal paths are suppressed throughout the network (Huang et al., 21 Aug 2025).

1. Conceptual definition and problem setting

CAT is a temporal causal discovery model built around a forecasting-first paradigm. Rather than estimating a causal graph directly, it first trains a multi-layer Transformer-based time-series forecaster to model long-range, nonlinear temporal relationships among variables, and then interprets the trained predictor by measuring how perturbations in lagged inputs affect target forecasts (Huang et al., 21 Aug 2025).

The underlying causal object is a directed graph with lag annotations. In CAT, a directed edge $X_i \to Y_j$ is declared when the normalized gradient sensitivity of the forecast for $Y_j$ to some lagged value of $X_i$ exceeds a user-defined threshold $\theta$. The corresponding causal lag is the $\arg\max$ over lagged gradient magnitudes. This makes CAT a representation-learning approach to temporal causality in which the causal graph is reconstructed from the behavior of a trained predictor rather than from explicit structural constraints in the training objective alone (Huang et al., 21 Aug 2025).

Within the broader transformer-for-causality literature, this places CAT alongside forecasting-based discovery models such as CausalFormer, which interprets a causality-aware transformer using regression relevance propagation (Kong et al., 2024), and sparse-attention Granger-style models such as the Sparse Attention Transformer, also referred to as CAT in the supplied material, which recovers causal influence by masking variables and comparing residual variances (Mahesh et al., 2024). A related but more general claim appears in "Transformer Is Inherently a Causal Learner," where decoder-only transformers trained autoregressively are shown to encode lagged causal structure in gradient sensitivities under assumptions A1-A4 (Wang et al., 9 Jan 2026).

2. Transformer architecture

CAT’s core forecaster tokenizes both static variables and temporal variables. Given $N$ source time-series variables $X=(X_1,\dots,X_N)$, each of length $S$, each series is normalized and split into $T$ patches of length $P$. Each patch is linearly projected into a $Y_j$0-dimensional embedding, while static variables are projected to a single $Y_j$1-dimensional vector. Standard sinusoidal positional encoding is then added so that temporal order is retained in the token sequence (Huang et al., 21 Aug 2025).

The encoder is organized as $Y_j$2 stacked Transformer encoder levels, each containing $Y_j$3 identical Transformer blocks. If $Y_j$4 denotes the input to level $Y_j$5, a block computes multi-head self-attention,

$Y_j$6

followed by residual connection, layer normalization, a position-wise feed-forward network,

$Y_j$7

and a second residual-plus-normalization step. For attention head $Y_j$8, the attention weight from token $Y_j$9 to token $X_i$0 is

$X_i$1

After $X_i$2 levels, each time-series variable is reduced to a single token by a temporal “strided convolution” procedure such that $X_i$3 (Huang et al., 21 Aug 2025).

The output stage linearly projects the final token for each target variable $X_i$4 to the forecast window $X_i$5. Numerical targets use L1 loss, and categorical targets use cross-entropy. This choice keeps the causal extraction stage decoupled from target type: the transformer is optimized as a forecaster, and causal interpretation is imposed after training through sensitivity analysis rather than through a bespoke graph-structured decoder (Huang et al., 21 Aug 2025).

3. Gradient-based causal graph and lag recovery

After forecaster training, CAT extracts causal relations by finite-difference approximation of partial derivatives. Let $X_i$6 denote the forecast for target $X_i$7. For each source variable $X_i$8 and lag $X_i$9,

$\theta$0

These values are stacked into a gradient-score matrix $\theta$1 with entries $\theta$2 (Huang et al., 21 Aug 2025).

To make scores comparable across variables and lags, CAT normalizes by the global maximum: $\theta$3 A directed edge $\theta$4 is declared when

$\theta$5

and the inferred lag is the $\theta$6 of $\theta$7 (Huang et al., 21 Aug 2025).

This extraction mechanism is closely aligned with a broader methodological shift in transformer-based causal discovery from raw attention inspection to differential or attributional criteria. SAT interprets masked attention and prediction-error changes as a Granger-causality analogue (Mahesh et al., 2024). CausalFormer argues that analyzing attention weights alone is an incomplete mapping from model parameters to causality and instead propagates relevance through the whole network (Kong et al., 2024). "Transformer Is Inherently a Causal Learner" goes further by establishing a score-gradient energy criterion,

$\theta$8

with $\theta$9 under assumptions A1-A4, and then estimates a practical surrogate using aggregated relevance scores (Wang et al., 9 Jan 2026). Taken together, these works indicate that CAT’s gradient-based graph recovery belongs to a technically coherent line of research in which causal structure is inferred from input-output sensitivities of trained transformers.

4. Prior knowledge integration and spurious-correlation control

A distinctive feature of CAT is the integration of user-specified excluded links directly into multi-head attention. Let $\arg\max$0 be a binary adjacency-mask matrix over variables, where $\arg\max$1 means that the link “variable $\arg\max$2 variable $\arg\max$3” has been excluded by prior knowledge. CAT expands this variable-level matrix into a token-level mask $\arg\max$4 by replicating rows and columns according to the number of tokens per variable in layer $\arg\max$5 (Huang et al., 21 Aug 2025).

The standard attention score is then replaced by a masked score,

$\arg\max$6

with $\arg\max$7 for prohibited entries so that forbidden attention weights become zero. The mask is applied consistently across all heads and blocks in a level. Because source and target variables are kept in distinct token groups, the paper states that prohibited edges never slip through via multi-layer attention propagation (Huang et al., 21 Aug 2025).

The stated purpose of this mechanism is to mitigate spurious causal relationships. CAT couples hard attention exclusion with standard weight decay on all projection matrices and a relatively high gradient threshold $\arg\max$8, with $\arg\max$9 given as an example, to prune weak, noise-driven sensitivities (Huang et al., 21 Aug 2025). This architecture-level masking differs from related approaches that also encode causal constraints in attention but target different tasks. "CaTs and DAGs" masks cross-attention according to a known DAG so that predictions remain causally compliant at inference time (Vowels et al., 2024), while the DAG-aware Transformer for causal effect estimation masks self-attention using ancestor reachability in a known causal graph (Liu et al., 2024). CAT instead uses exclusion masks within a temporal discovery pipeline, where the graph is still being inferred from data but some edges may be ruled out a priori.

5. Empirical evaluation

CAT was benchmarked against six state-of-the-art baselines: cMLP, cLSTM, TCDF, DVGNN, CUTS, and CausalFormer. The evaluation used three dataset families: Basic Synthetic with “fork,” “V-structure,” “mediator,” and “diamond” motifs and 1,000-point series; Lorenz96 nonlinear climate with 10 variables and 1,000 points; and NetSim fMRI network simulations with node counts 5–50 and lengths 50–5,000. Performance was measured using F1-score on recovered edges and Precision of Delay (PoD), defined as the fraction of edges with correctly estimated lag within one time-step (Huang et al., 21 Aug 2025).

The reported aggregate result is a 12.8% absolute improvement in F1 over the next best method, together with 98.9% accuracy in estimating causal lags. A specific example given in the paper is Lorenz96, where CAT attains 0.84 F1 versus 0.69 by CausalFormer. Ablation studies further report that removing the prior-masking mechanism drops F1 by 6–8 points, while replacing the multi-layer stack with a single layer, as in CausalFormer, loses an additional 4–5 points (Huang et al., 21 Aug 2025).

These results situate CAT within a competitive transformer-based discovery landscape. CausalFormer reports state-of-the-art performance on synthetic, simulated, and real datasets using a causality-aware transformer plus a decomposition-based causality detector (Kong et al., 2024). SAT reports AUC-ROC of 0.70–0.78 and F1-score of 0.63–0.72 across synthetic benchmark regimes, compared with 0.47–0.67 AUC and 0.51–0.66 F1 for VAR-based Granger causality (Mahesh et al., 2024). "Transformer Is Inherently a Causal Learner" reports that its transformer-based discovery method maintains $N$0 at $N$1, preserves $N$2 up to $N$3, and achieves AUROC/AUPRC of $N$4, $N$5, and $N$6 on AQI, Traffic, and Medical, ranking top-3 among 16 methods on the CausalTime benchmark (Wang et al., 9 Jan 2026). A plausible implication is that CAT’s strongest empirical distinction within this cluster is not merely the use of transformers for causal recovery, but the combination of gradient-based lag extraction with an explicit prior-knowledge masking interface.

6. Terminology, related usages, and interpretive issues

The acronym “CAT” is not unique in the literature. In addition to the temporal causal discovery framework described above (Huang et al., 21 Aug 2025), it has been used for Causal Attention Tuning in LLMs (Han et al., 1 Sep 2025), CAT-XPLAIN for inherently explainable vision transformers (Khanal et al., 2022), and Causality-Aware Transformer Networks for robotic navigation (Wang et al., 2024). In adjacent work, “CaT” also denotes Causal Transformers constrained by a predefined DAG (Vowels et al., 2024). For precise interpretation, the acronym therefore has to be resolved by task domain and paper context.

A recurring interpretive issue across these works is whether transformer attention itself should be read as causality. SAT is explicitly built on the intuition that self-attention scores can be re-interpreted as measures of causal influence (Mahesh et al., 2024). CausalFormer takes the opposite methodological stance that parameters such as attention weights and convolution weights alone are insufficient, because other components such as fully connected layers and activation functions also matter for causal discovery (Kong et al., 2024). The 2026 identifiability result for autoregressive transformers reframes the question again: under conditional exogeneity, no instantaneous effects, lag-window coverage, and faithfulness, nonzero score-gradient energy rather than raw attention is what exactly recovers lagged parents (Wang et al., 9 Jan 2026). CAT, as defined in (Huang et al., 21 Aug 2025), is methodologically closer to this latter view, since its graph extraction is based on gradient sensitivity and thresholding, while attention masking is used primarily for prior knowledge integration and spurious-link suppression.

This distribution of viewpoints suggests that “causality-aware” is best understood as a family resemblance rather than a single architectural recipe. In the supplied literature, the designation covers at least four distinct strategies: forecasting-plus-gradient extraction (Huang et al., 21 Aug 2025), sparse attention with Granger-style masking (Mahesh et al., 2024), whole-model relevance decomposition (Kong et al., 2024), and explicit attention constraints from known causal graphs (Vowels et al., 2024, Liu et al., 2024). What unifies them is the attempt to move beyond correlation-only sequence modeling by embedding temporal priority, structural constraints, or attributional identifiability into transformer computation.