LEMD: Latent Environment Extrapolation and Message Disentanglement for Dynamic Graph Under Distribution Shift

1 Introduction

By coupling time, interaction topology, and features, dynamic graphs can naturally represent many systems in the real world, such as social networks [Berger-Wolf and Saia, 2006], commercial networks [Cheng et al., 2022], and collaboration networks [Hughes et al., 2010]. Dynamic graph neural networks (DyGNNs) have become an advanced paradigm in dynamic graph representation learning by combining temporal modeling with spatial message passing [Kumar et al., 2019; Barros et al., 2021]. Most DyGNNs rely on the implicit assumption that training and testing data are identically distributed [Zhang et al., 2022]. But in practice, the process of data generation is affected by unobserved environments that vary across domains and evolve over time [Creager et al., 2021]. Such environmental variations may affect both the topology and features of the dynamic graph [Yuan et al., 2023]. Many studies on generalization have recognized the challenge of data out of distribution (OOD) and attempted to solve it [Ye et al., 2025; Arjovsky et al., 2020]. However, most existing methods are mainly developed for Euclidean data [Li et al., 2021]. Due to the evolving of non-Euclidean topology, extending these techniques to dynamic graphs is non-trivial [Tian et al., 2024; Zhang et al., 2022].

Although some work has made important progress in learning robust representations of dynamic graphs [Zhang et al., 2024], there are still two open issues that have not been solved. Firstly, limited and unreliable environment modeling.

Considering the domain shift, the available training distribution may only cover a narrow range of conditions, while the target distribution may lie outside this support [Rosenfeld et al., 2021; Lu et al., 2024]. While unconstrained internal mixup of data or inferring the environment through pseudolabeling may be effective in some cases [Zhang et al., 2024; Yuan et al., 2023], the invariance of learning only from a limited source may fail when the target distribution exceeds the training support [Christiansen et al., 2022; Vuong et al., 2025]. It is still challenging to perform environment extrapolation with effective constraints. Secondly, insufficient disentanglement across temporal and spatial dimensions. As the toy dataset shown in Figure 1, the distribution shift of dynamic graph can jointly affect features and topology across temporal and spatial dimensions, while the potential relations remain invariant. Static disentanglement only for observable factors (e.g., edges) is incomplete and leads to shortcuts leakage. A proxy object that can contain real domain information and invariant relations should be designed and allow the process of disentanglement to evolve over time.

To address these limitations, a novel approach named LEMD (Latent environment Extrapolation and Message Disentanglement) is developed for dynamic graph under distribution shift. We start from the message-sufficiency view of DyGNNs: A multiset of messages received by a node v at time t can fully mediate its features and topology while containing sufficient potential relationships between nodes. Accordingly, the distribution shift of dynamic graphs can be studied as the shift in dynamic message distributions. Thereby, we first design an uncertainty-guided dynamic message extrapolation mechanism. By focusing on the sufficient proxy of distribution shift, this mechanism produces a diverse extrapolation set of counterfactual messages under uncertainty constraints, thereby improving the model's ability to generalize to unseen domains without requiring any explicit domain labels. Fine-grained disentanglement of causal and spurious information is proposed at the message level. For each node, each message from neighbor is disentangled, taking into account both topology and features. An evolving gating network is introduced to allow the disentanglement mechanism to vary over time. Finally, a unified training objective is proposed with theoretical guarantees. Extensive experiments on six real-world datasets show that our method achieves state of the art performance in two downstream tasks under distribution shift. Our contributions are as follows:

• A novel framework for dynamic graph OOD generalization. We propose LEMD to adapt to dynamic graph data shifts that are widely observed in the real world.
• Message extrapolation constrained by uncertainty. We extrapolate counterfactual message distributions to construct diverse reliable environments without requiring explicit domain labels.
• Evolving message disentanglement. We disentangle neighbor messages into causal and spurious components identifiably, and employ a time-evolving gating mechanism to allow the temporal shift of disentanglement.
• Theoretical guarantees and empirical superiority. We provide theory for the proposed method and objective

and achieve state-of-the-art on six real-world datasets across two downstream tasks under distribution shift.

2 RelatedWork

Dynamic graph representation learning. DyGNNs have developed into a class of temporal representation learners that unify message passing with temporal modeling [Ma et al., 2026]. In snapshot settings, a prevalent paradigm is to encode each snapshot with a shared GNN and then capture temporal dependence of node states via recurrent units or temporal attention [Sankar et al., 2020], offering a practical balance among expressivity, efficiency, and long range modeling. Recently, work such as EvolveGCN [Pareja et al., 2020] and SEIGN [Qin et al., 2023] have developed temporal modeling from the evolution of node representations to the evolution of both parameters and representations. Most dynamic graph representation learning methods follow the identical distribution assumption, which results in their limited generalization ability under the data distribution shift.

Causal learning for OOD generalization. A significant insight for OOD generalization is that robustness under distribution shift requires disentangling and d-separating the spurious factors, which means predictive shortcuts that hold in training but break in testing [Komanduri et al., 2024]. Causal intervention methods generate counterfactual variations to discourage dependence on shortcuts [Sun et al., 2021]. Invariant learning enforces stability of representations across environments, whether provided by domain labels [Ganin et al., 2016] or inferred as latent partitions [Creager et al., 2021]. Despite success on Euclidean data, the distribution shift of dynamic graphs affects evolving topology and features, making intervention and disentanglement challenging.

OOD generalization on graphs. Building on the intervention and invariance principles, recent work has begun to study OOD generalization for dynamic graphs [Tian et al., 2026; Tian et al., 2025; Ma et al., 2025]. DIDA [Zhang et al., 2022] disentangles spatial-temporal patterns via flipping the attention mechanism and constructs interventions by sampling and reassembling neighbors across spatial and temporal. SILD [Zhang et al., 2023] learns disentangled spectral masks with invariant filtering to suppress domain sensitive frequency components. DyAug [Zhang et al., 2024] extends graph rationalization to dynamic settings by separating rationale subgraphs over time and performing temporally consistent regularization. EAGLE [Yuan et al., 2023] focuses on environments inference, combining disentanglement with a diverse instantiation mechanism. Despite this progress, existing methods often struggle to produce faithful, extrapolative counterfactuals and to disentangle causal and spurious factors sufficiently. Accordingly, we treat messages as the sufficient mediator and combine reasonable environment extrapolation with an evolving and identifiable message disentanglement to achieve OOD generalization of DyGNNs.

3 Preliminaries

3.1 Dynamic graph domain generalization

Dynamic graphs. Adynamic graph is a sequence of graphs G1:T = {Gt}Tt=1 = {(Vt,At,Xt)}Tt=1 . Each snapshot Gt contains a node set Vt, an adjacency matrix At, and a node features matrix Xt ∈ R|Vt|×dx . The adjacency matrix is defined as: (At)uv = I[ u→v at time t ], for u, v ∈ Vt. For a node v ∈ Vt, we denote its incoming neighbor set at t as N t v = {u ∈ Vτ : (Aτ)uv = 1 | ∃τ ≤ t}.

Domain generalization on dynamic graphs. Dynamic graph domain generalization posits latent environments e ∈ E that govern the data generation process and induce distribution shifts in both topology and features. Conditioned on e, samples follow the joint distribution P(G1:T , y | e). Unobserved environment distribution differs: e ∼ Ptr(e) at training and e ∼ Pte(e) at testing [Wu et al., 2024].

We learn fθ that minimizes the expected testing loss:

where Pte is unknown during training and the environment index e is unobserved. Hence, the key challenge is to build robustness to unseen distribution shifts without relying on explicit environment labels.

3.2 Amessage passing pipeline of DyGNNs

Most DyGNNs implement a spatiotemporal message passing pipeline: exchange messages along edges and update node representations by aggregating messages [Gilmer et al., 2017]. Let htv ∈ Rd be the hidden state of node v at time t. We apply an attention-style DyGNN layer, where the message is defined as:

where qtv = Wqhtv , kτu = Wkhτu and vτu = Wvhτu . And W(·) denotes learnable projections, dk is the dimension of the latent variable kτu . Collecting all directed messages received by v at time t gives a message multiset Mtv := {mτ,t uv | u ∈ N t v }. Aggregation and update are then applied:

where AGG(Mtv ) ∈ Rd and UPD(·) : Rd → Rd .

4 Method

We propose LEMD as shown in Figure 2. Following the message fully mediated perspective in Sec. 4.1, the dynamic graph is induced into message multisets. Uncertain extrapolation (Sec. 4.2) and evolving causal disentanglement (Sec. 4.3) are applied for message multisets to generate diverse distributional shifts and suppress spurious shortcuts. Rigorous theories and proofs are given to support method validity.

4.1 Amessage-mediated causal view

Message sufficiency as explicitmediation. Given themessage passing pipeline in Sec. 3.2, a DyGNN layer constructs the received message multiset Mtv and then updates node states by ˜htv = UPD(AGG(Mtv )). Let Pred(·) be an arbitrary task head applied on ˜htv . Then for each prediction ˜ ytv :

˜ ytv = gθ(Mtv ), gθ ≜ Pred ◦ UPD ◦ AGG, (4) which formalizes a mediation principle: for fixed model parameters, the dependence on the dynamic graph history (A1:t, X1:t) is fully mediated by the message multiset Mtv .

Induced message distributions under latent environments. Recall that environments e ∈ E govern the datagenerating process (Sec. 3.1). SinceMtv is constructed from (A1:t,X1:t) given fixed model, each environment e induces a corresponding joint distribution in message space, denoted by P(Mtv , y | e). Accordingly, for node v at time t, the test message distribution is a mixture over environments:

and similarly PMtr is defined by Ptr(e). Therefore, Eq. (1) admits an equivalent message-space view:

Eq. (6) illustrates that domain shifts on (A1:t,X1:t) are equivalently captured as shifts in the distributions P(Mtv | e).

A structural causal model with label invariance. A structural view is adopted in which a stable semantic factor and an environment-specific bias factor jointly shape the observed dynamic graph and thus the induced messages. For each t, let Ct denote stable semantics and St denote environmentdependent biases. The key objects in our setting admit the following factorization (withMtv as an explicit mediator):

We assume counterfactual label invariance with respect to environment shifts:

i.e., changing environment can affect interactions and features, while the causal mechanism remains preserved.

4.2 Uncertainty-guided environment extrapolation

Counterfactuals in message space. Under the messagemediated view in Eq. (6), unseen environments perform as shifts in the induced message distributions. Since the environment is latent, we proxy such shifts by constructing counterfactuals in message space around the observed messages.

For node v at time t, we extrapolate from its received message multiset Mtv . Let S tv : Mtv → Rq be a parameterfree summary operator and S(M) be the collection of messages over all (v, t). We estimate an uncertainty descriptor ρ = U(S(M)) from observed summaries, where ρ ∈ Rq+ and U quantifies the feasible variation scale supported by the observed messages. This yields an uncertainty guided soft region of counterfactual message distributions:

where Q is a density distribution centered at the factual summary with scale ρ. We sample K counterfactual message multisets { ˆM(k) }Kk=1 from Ω(ρ) via the reparameterized trick with moment matching.

Reparameterized moment matching construction. According to Theorem 1, the smaller dimension q of ρ, the lower the upper bound of expectation. The low order moments of a feature can effectively summarize its information. Studies have shown that low order moments of messages can reflect environmental information to a certain extent [Li et al., 2020]. The low order moments of Mtv are used as its summary:

where µtv and σtv are the mean and standard deviation within Mtv respectively. Then, the scale of feasible extrapolation ρ is estimated by analyzing the dispersion of S tv (Mtv ):

where the δµ and δσ are the standard deviation of µ and σ across all time and nodes. ρ measures how heterogeneous all the incoming messages are.

According to the central limit theorem, Q is naturally realized as a bivariate diagonal Gaussian distribution in the summary space:

K counterfactual summaries are sampled from Q, i.e. (˜µt,(k) v , ˜σt,(k) v ) ∼ Q. In actual implementation, to ensure gradient propagation, the reparameterization trick is applied:

where ϵt,(k) µ,v , ϵt,(k) σ,v ∼ N(0, 1). σ is cut off by a certain minimum positive value to avoid σ < 0.

Given the sampled k-th counterfactual summary (˜µt,(k) v , ˜σt,(k) v ), the counterfactual message multiset is generated by an affine moment transform:

where i is the index of message withinMtv . This counterfactual generation method extrapolates messages under control while preserving their normalized structure, thereby providing a proxy for shifts induced by unseen environments.

Counterfactual invariance regularization. Let m(0) i denote a factual directed message and ˆm(k) i (k ≥ 1) its counterfactual variants generated by Eq. (9)-(14). Let mc,(k) i be the corresponding causal component produced by the disentangler in Sec. 4.3. We enforce counterfactual invariance by penalizing the variance across sampled environments:

where Vark(·) denotes the mean squared deviation over k ∈ {0, . . . , K} and M is the collection of all messages.

4.3 Evolving message disentanglement

Disentangling the messages with evolving parameters. Our causal formulation assumes a stable semantic factor Ct that fully determines the target mechanism as shown in Eq. (8), while the latent environment e affects topology and features and hence the induced messages through spurious factors. Meanwhile, message sufficiency view states that predictions are mediated by the received messages, ˜ ytv = gθ(Mtv ). This means that the message space can be used as a proxy for observing distributional shifts. And the causal component that is predictive of y but independent of environmental changes can be disentangled from each message.

Concretely, given the message indexed by i at target time t, a counterfactual family {m(k) i }Kk=0 is constructed according to Sec. 4.2, where m(0) i ≜ mi is factual and m(k) i (k ≥ 1) are counterfactual messages. Then a time-evolving network Dt is introduced to disentangle a causal component mc,(k) i that remains stable across k and a complementary spurious component that practically independent of y.

A time varying generation model in message space. We set that each message is generated from a time-dependent mixing between stable semantics and spurious biases:

where Ct is shared across k and S(k) t varies with the counterfactuals. The mixing mechanismΨt is indexed by t to capture temporal distribution shift.

Recoverable causal-spurious disentanglement. A time specific network Dt is constructed to disentangle each message into a causal and a spurious component:

with recover operator Rt:

Recoverability of Dt prevents trivial solutions that discard information, while counterfactual invariance (Eq. (15)) encourages mc,(k) to represent stable semantics.

Practical instantiation. Considering complexity and recoverability requirements, Dt is instantiated as an elementwise gate αti ∈ (0, 1)d which is calculated from the message by an evolving gate network:

where σ is the Sigmoid function and W(1) t ,W(2) t are timespecific parameters. This gate is then applied to disentangle every counterfactual variant:

where⊙ is the Hadamard product. Eq. (18) holds with Rt being simple addition m(k) i =mc,(k) i +ms,(k) i .

To track the shift in Ψt, we let the gate parameters evolve with t by an evolving parameter generator. {W(1) t ,W(2) t } are generated through a recurrent generator in parameter space:

where st is the recurrent state and ψ are generator parameters. This induces a smooth and associated temporal prior over the disentanglement mechanism without requiring an independent disentangler at each t.

Correlation regularization. To enforce the independence between mc and ms , we penalize their dependence using a centered correlation:

with ¯mc = 1 |M| ∑|M| i=1 mci and ¯ms = 1 |M| ∑|M| i=1 msi . Theorem 2 explains the contribution of independence to identifiable causal components.

The final objective. For each (v, t), we feed the raw causal message multiset Mc,t,(0) v to the DyGNN update via ˜htv = UPD(AGG(Mc,t,(0) v )). Then the supervised loss function is computed asLsup = 1

Hence the complete training objective is

where λinv and λcor are coefficients of counterfactual invariance and causal-spurious separation respectively. Theorem 2 provides theoretical guarantees for the approximate identifiability of disentanglement under this optimization objective.

5 Analysis

Two theories are proposed to support the effectiveness of messages extrapolation and evolving disentanglement. All the theories are rigorously proved in Appendix A.3.

Theorem 1 (Counterfactual causal-message invariance with uncertainty-guided probabilistic coverage). Assume: (i) the loss is bounded 0 ≤ ℓ ≤ Bℓ; (ii) for all y, the mapping M ↦→ ℓ(gθ(M), y) is L-Lipschitz w.r.t. the message input; (iii) the statistic-to-distribution regularity and target typicality assumptions hold. Then for any δ ∈ (0, 1), with probability at least 1-δ over the draw of N training instances and the uncertainty-guided sampling of K counterfactual summaries, the following holds:

where εN(·) is the standard uniform-convergence term and rK,δ(r) is a Gaussian K-sample covering radius.

Remark. The bound above decomposes into: the empirical supervised risk on factual causal messages, a stability term controlled by the counterfactual invariance loss, a standard generalization term, and a probabilistic coverage term quantifying how likely the K sampled counterfactual summaries contain one that is close to the target distribution. For fixed q and typicality radius r, the covering radius satisfies rK,δ(τ)= ˜O(K-1/q ). As a result, increasing K and decreasing Lsup and Linv may tighten the bound.

Theorem 1 gives perspective from the upper bound on test loss expectation, which explains the necessity for Lsup and Linv, and illustrates the effectiveness of uncertaintyconstrained message extrapolation for reducing the coverage error term and thus the upper bound on test loss expectation.

Theorem 2 (Approximate identifiability of invariant causal messages). Fix a time step t and assume the augmented messages follow an invertible mixing model Ψt. Let Dt be any recoverable disentangler as shown in Eq. (17). If the learned causal part is approximately counterfactually invariant,

and regularity holds non-degenerate spurious support and continuity, then there exists a decoder ht such that

In particular, as εinv→0, mct identifies Ct up to an invertible transform on its support.

Remark. The result is pointwise in t, hence compatible with time-evolving disentanglement; full constants and proof are provided in Appendix.

Theorem 2 formalizes that counterfactual invariance is the key condition for recovering stable semantics from the learned causal messages. The correlation penalty Lcor in Eq. (22) serves a complementary purpose: it discourages linear leakage between mc and ms , thereby improving the operational cleanliness of the split and stabilizing optimization under drift. A formal statement is given in Appendix.

6 Experiments

6.1 Experiment Setup

Datasets. We evaluate the performance of LEMD on dynamic graphs by six real world datasets under two tasks: future link prediction and node classification. For each dataset, we construct a dedicated distribution shift setting. COLLAB [Tang et al., 2012] is a citation network dataset for 16 years, comprising paper authors and their citation relations. Edges associated with the Data Mining category are regarded as distribution shift data. YELP [Sankar et al., 2020]

is a business-review dataset covering 24 months, consisting of interactions (reviews) between customers and businesses. Reviews for the Pizza category are regarded as distribution shift data. MOOC-ACT [Kumar et al., 2019] contains student interaction records on a MOOC platform over 30 days. WIKI [Kumar et al., 2019] records page-editing activities over 31 days on Wikipedia. REDDIT [Kumar et al., 2019] consists of posts made by users on Reddit subreddits over a month. According to general settings of strong baselines [Yuan et al., 2023], for MOOC-ACT, WIKI, and REDDIT, we perform K-means clustering on the interactions therein and randomly select a class as distribution shift. The above five datasets are used to evaluate future link prediction under distribution shifts. In addition, we adopt AMINER [Sinha et al., 2015; Tang et al., 2008] for node classification. AMINER is a longterm citation network spanning 18 years. Since the dynamics of the graph may evolve over time, we regard the final three years as distribution shift data.

Baselines. As in Sec. 2, our baseline includes three classes of methods. (i) Domain generalization methods contain IRM [Arjovsky et al., 2020], V-REx [Krueger et al., 2021] and GroupDRO [Sagawa et al., 2020]. (ii) Dynamic graph learning methods contain GCRN [Seo et al., 2018], DySAT [Sankar et al., 2020], EvolveGCN [Pareja et al., 2020] and SEIGN [Qin et al., 2023]. (iii) OOD generalization methods on dynamic graphs contain DIDA [Zhang et al., 2022], IB-D 2 [Wang et al., 2025], EAGLE [Yuan et al., 2023], EvoGD [Sun et al., 2025], DGIB [Yuan et al., 2024], SILD [Zhang et al., 2023] and DyAug [Zhang et al., 2024]. The specific method is introduced in the Appendix.

Implementation details. For the future link prediction task, the dataset is first divided into training, validation, and test sets based on the temporal order, and then categorized into an OOD section and a without OOD section based on the different domains of the edges. All pairs of nodes with existing edges and equally sampled non-existing edges in a future time slice are used as the ground truth of positive and negative samples. Test set results for the OOD section and w/o OOD section are reported, using AUC% as the evaluation metric. For the node classification task, model is used to predict the venues of the papers. Results of future three years are reported separately as OOD data, using percentage of accuracy (ACC%) as the evaluation metric.

6.2 Results

Link prediction under domain shifts. Experiments on future link prediction task under domain shifts are conducted on five widely used real world dynamic graph datasets. The results are shown in Table ??, where blue background with bold font indicates the optimal results, and light blue background with underline indicates the suboptimal results. Across all five datasets, SILD achieves the largest number of suboptimal results, which may be attributed to its strong capability of spectral disentanglement. Comparing all the most advanced baselines, LEMD achieves state of the art on all datasets under OOD setting. On the MOOC-ACT dataset, LEMD recorded the largest performance increase of 7.7% compared to the suboptimal baseline. LEMD achieves state-of-the-art performance on four datasets under the w/o OOD setting. Only under the w/o OOD setting on the COLLAB dataset does LEMD underperform the best baseline, SILD, by a margin of 0.1%. These results suggest that LEMD has strong representation and generalization ability across domain.

Node classification under domain shifts. Experiments on node classification task under domain shifts on AMINER dataset. Considering the impact of the outbreak of deep learning, the citation pattern of papers since 2015 has major changed. As the results shown in Table 2, LEMD achieves the state of the art in accuracy of node classification prediction in all three years. Averaged over three years, LEMD's node classification accuracy is 3% higher than the suboptimal SILD. These results demonstrate LEMD's universal representational capacity and its strong generalization across domain on a broader range of tasks.

Ablation Study. Ablation experiments are performed on our proposed module on six datasets within two tasks. Each subfigure in Figure 3 shows, in turn, complete LEMD, without Linv, without Lcor, without to Evolve disentanglement, and without message Extrapolation. The ablation results show that the removal of the two loss constraints and the two modules will respectively drop the model performance. For link prediction tasks under distribution shifts, it is usually the removal of message extrapolation impacts most severely, which reveals the effective modeling of message distribution extrapolation for data across domains. For node classification tasks under distribution shifts, the removal of evolve disentanglement impact most severely. This result is consistent with the context that the AMINER dataset is primarily affected by deep learning breakout, which creates shifts alone time.

Hyperparameter Sensitivity. We performed sensitivity analysis on all three most important hyperparameters λinv, λcor and K on the most representative dataset YELP, which results are shown in Figure 4. It can be found that moderate invariance regularity and correlation regularity (e.g. λinv = 1e - 3, λcor = 1e - 4) give the best model performance. Moreover, the model performance increases with increasing K, but there is a marginal decreasing effect.

7 Conclusion

This paper investigates dynamic graph domain generalization under distribution shifts, and motivates a message-sufficiency causal perspective in which predictions are fully mediated by received messages. We propose LEMD, a framework that improves robustness by intervening in message space. LEMD combines message extrapolation under uncertainty constraints to sample counterfactual message distributions with evolving causal-spurious message disentanglement via a recoverable gated network. Supported by theoretical analysis, LEMD achieves consistent OOD gains across six datasets, including +7.7% on future link prediction task and +3.0% average accuracy on node classification task.

Acknowledgments

This work is supported by the National Natural Science Foundation of China (No. 62272338). Chen Zhao, Xintao Wu, and Zhong Chen did not receive any financial support for this work and contributed only by developing the research ideas, participating in discussions, and providing feedback on the manuscript.

Contribution Statement

Minglai Shao is the corresponding author.