🤖 machine learning

Invariant Graph Representations for Continuous-Time Dynamic Graphs Under Distribution Shifts

The paper proposes CIR, a framework leveraging a novel structural causal model (ICCM) and Normalized Weighted Geometric Mean approximation to efficiently learn invariant graph representations for continuous-time dynamic graphs, thereby achieving robust performance under out-of-distribution shifts where existing methods fail.

Original authors: Lanting Fang, Yulian Yang, Yawei Zhang, Shanshan Feng, Kaiyu Feng, Hanning Yuan

Published 2026-06-24

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Lanting Fang, Yulian Yang, Yawei Zhang, Shanshan Feng, Kaiyu Feng, Hanning Yuan

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Picture: Predicting the Future Without Getting Fooled

Imagine you are trying to predict whether two people, Alice and Bob, will become friends. You look at their history: who they talked to, when they talked, and who else was in the room.

Most computer programs (AI) that do this are very good at spotting patterns in the data they were trained on. But they are terrible at handling change. If the rules of the game shift slightly (a "distribution shift"), these programs often get confused and make bad guesses.

This paper introduces a new system called CIR (Continuous-time Invariant Representation). Its goal is to teach an AI to find the true reasons why things happen, so it can make good predictions even when the world changes around it.

The Problem: The "Fake Clue" Trap

To understand why this is hard, let's look at a specific example from the paper (Figure 1).

The Scenario:
You want to predict if Alice and Bob will link up.

The Real Reason (The Invariant Truth): Alice and Bob become friends because they both know a mutual friend, Charlie. This is a "triadic closure" (A knows B, B knows C, so A and C meet). This is a solid, unchanging rule of human behavior.
The Fake Clue (The Spurious Correlation): In your training data, every time Alice and Bob became friends, there was also a random stranger, Dave, who happened to wave at Alice right before.

The Mistake:
A standard AI looks at the data and says, "Aha! Every time Dave waves, Alice and Bob become friends. I will bet on Dave!"

In the training world: The AI is right 100% of the time.
In the real world (OOD - Out of Distribution): You test the AI on a new group of people where Dave never waves. The AI, having learned the "Dave" rule, now fails completely because it ignored the real reason (Charlie).

The paper calls this spurious correlation. The AI is memorizing shortcuts instead of learning the actual cause.

The Solution: CIR and the "Time-Traveling Detective"

The authors propose CIR, a system designed to be a "Time-Traveling Detective." It doesn't just look at what happened; it asks, "What caused this to happen, regardless of the weird coincidences?"

Here is how CIR works, broken down into three simple steps:

1. The Two Detective Teams (Structural & Temporal)

The paper argues that in a constantly changing world (Continuous-Time Dynamic Graphs), you need to look at two things separately:

The Structure Team: Who is connected to whom? (e.g., The triangle of Alice, Bob, and Charlie).
The Time Team: When did things happen? (e.g., Did the conversation happen right after a big event?).

CIR uses special tools to extract these two types of clues. It separates the "real causes" from the "background noise."

2. The "What If" Machine (Intervention)

This is the smartest part. The paper uses a concept from causal science called Intervention.

Imagine you have a photo of a crime scene. A normal AI looks at the photo and guesses.
CIR asks: "What if we magically erased the 'Dave' wave from the photo? Would Alice and Bob still become friends?"

If the answer is "Yes," then the wave wasn't the cause. If the answer is "No," then the wave was the cause.

The Problem with this approach:
In a real-world graph with millions of connections happening every second, you can't physically erase every single "Dave" and check the result. That would take forever (computationally impossible).

The Trick (The NWGM):
The authors use a mathematical shortcut called the Normalized Weighted Geometric Mean (NWGM).

Analogy: Instead of trying to interview every single person in a city to see how they would react to a new law, you interview a few representative groups and use a smart formula to guess the average reaction of the whole city.
This allows CIR to simulate "What if?" scenarios instantly without doing the impossible math of checking every single possibility.

3. The "Memory Bank" (Environment Memory)

The world changes. Sometimes the "fake clues" look like one thing, and sometimes they look like another.

CIR keeps a Memory Bank. Think of this as a notebook where the AI writes down different "types" of fake clues it has seen in the past (e.g., "Dave waving," "Rainy days," "Holidays").
When the AI sees a new situation, it checks its notebook to see if it's dealing with a familiar type of "fake clue." It then adjusts its prediction to ignore that specific noise.

Why is this better than what we have now?

The paper tested CIR against other top AI models on four real-world datasets (like Wikipedia edits, Reddit posts, and university course interactions).

The Old Models: When the "fake clues" (shortcuts) were removed or changed in the test data, the old models crashed. Their accuracy dropped by huge margins (sometimes losing 30-50% of their performance). They were too reliant on the specific patterns of the training data.
CIR: Even when the rules changed, CIR kept performing well. It successfully ignored the "Dave waves" and focused on the "Charlie connections."

Summary in One Sentence

CIR is a new AI framework that teaches computers to ignore coincidental patterns and focus only on the true, unchanging causes of events, allowing them to make accurate predictions even when the world changes in unexpected ways.

What the Paper Does Not Claim

It does not claim this is for medical diagnosis or clinical use.
It does not claim this works for static graphs (graphs that don't change over time) better than existing methods; it is specifically built for graphs that change continuously.
It does not claim to solve all AI problems, only the specific problem of predicting links in dynamic graphs when the data distribution shifts.

Technical Summary: CIR for Invariant Representations in Continuous-Time Dynamic Graphs

Problem Formulation
Continuous-Time Dynamic Graphs (CTDGs) model evolving relational systems where nodes and edges are added or removed at arbitrary time points, capturing fine-grained structural and temporal changes. While existing representation learning methods for CTDGs have achieved success in in-distribution settings, they exhibit limited robustness under out-of-distribution (OOD) shifts. Specifically, standard models often exploit spurious correlations—such as specific temporal patterns or structural motifs that appear frequently in training data but are not causally linked to the target interaction. When these correlations shift in test environments (e.g., due to changes in the underlying data generation process), model performance degrades significantly.

The core challenge addressed in this work is the scarcity of methods capable of learning invariant representations for CTDGs under OOD shifts. Existing causal approaches, which rely on structural causal models (SCMs) and interventions, are primarily designed for static or discrete-time graphs. Applying these to CTDGs is computationally prohibitive due to the combinatorial explosion of possible structural and temporal subgraphs, making exhaustive intervention infeasible.

Methodology: The CIR Framework
To address these challenges, the authors propose CIR (Continuous-time Invariant Representation), a framework grounded in a novel Structural Causal Model (SCM) termed the Independent and Confounded Causal Model (ICCM).

The ICCM Framework:
The ICCM explicitly models the causal mechanisms in CTDGs by distinguishing between:
- Invariant Causal Subgraphs: $C_s$ (structural) and $C_t$ (temporal), which are encoded into latent representations $H_S$ and $H_T$ .
- Confounders: Structural ( $D_s$ ) and temporal ( $D_t$ ) confounders that introduce spurious correlations, particularly in OOD environments.
- Prediction Mechanisms: The model defines an in-distribution prediction $Y^I$ based solely on invariant features, and an OOD prediction $Y^C$ that accounts for confounders. The goal is to perform causal interventions (using Pearl's $do$-operation) to block backdoor paths induced by confounders, thereby isolating the invariant causal effect.
Efficient Intervention via NWGM:
Directly intervening on all possible structural and temporal subgraphs is computationally intractable. To overcome this, CIR leverages the Normalized Weighted Geometric Mean (NWGM) approximation. This technique allows the model to efficiently approximate the expected interventional outcome by moving the expectation operator inside the non-linearity (e.g., sigmoid), avoiding the need for exhaustive enumeration of subgraph configurations.
Neural Architecture:
The ICCM is instantiated within a deep learning architecture comprising:
- Dual Subgraph Extractors: Dedicated modules extract causal structural and temporal patterns.
  - Temporal Extraction: Uses continuous-time encoding (cosine-based) and cross-node attention to identify temporally relevant edges that causally support a target link, separating them from non-causal temporal noise.
  - Structural Extraction: Encodes node features based on $n$ -hop neighborhoods within a time window, using attention mechanisms to filter out background nodes that introduce spurious correlations.
- Environment Memory Bank: To approximate the distributions of latent confounders ( $D_s, D_t$ ) without observing them directly, the model maintains a memory bank of prototype vectors. These prototypes represent clusters of non-causal subgraphs and unobserved factors. The bank is updated via self-supervised alignment and diversity losses to ensure prototypes capture distinct confounding modes.
- Fusion and Prediction: A multi-head attention module aggregates the memory bank entries to estimate the expected contribution of confounders. These estimates are fused with the invariant representations ( $H_S, H_T$ ) to produce the final OOD prediction. The model is trained to optimize both in-distribution and interventional (OOD) objectives simultaneously.

Key Contributions

First Investigation of OOD in CTDGs: The paper claims to be the first to investigate invariant representation learning specifically for CTDGs under OOD shifts, addressing a gap where existing causal methods fail due to the continuous and non-stationary nature of the data.
ICCM and NWGM Approximation: The authors introduce the ICCM to characterize invariant causal subgraphs in CTDGs and develop an efficient intervention strategy using NWGM. This avoids the combinatorial explosion associated with traditional causal interventions in dynamic settings.
The CIR Framework: A practical deep learning architecture is designed to instantiate ICCM, enabling the joint learning of invariant structural/temporal representations and environment-specific features via a memory bank mechanism.

Experimental Results
Extensive experiments were conducted on four benchmark datasets (MOOC, REDDIT, WIKI, UCI) with simulated distribution shifts (introduced via structural perturbations).

Performance: CIR consistently outperforms state-of-the-art baselines, including CTDG-specific models (TGN, TGAT, GraphMixer) and robust DTDG models (DIDA, EAGLE, Ada-DyGNN).
Robustness: Under increasing levels of distribution shift (shortcut ratios from 0.2 to 0.6), existing models suffer significant performance degradation (e.g., GraphMixer's AUC drops by over 50% on MOOC). In contrast, CIR maintains high AUC and AP scores with low variance, demonstrating superior generalization.
Efficiency: CIR achieves high inference efficiency (approx. $5.4 \times 10^{-4}$ seconds per edge on REDDIT), comparable to or better than complex baselines like TGAT and TGN, and significantly faster than DIDA, which requires constructing intervention sets for every snapshot.
Ablation: Removing key components (structural/temporal encoders, the OOD prediction branch $Y^C$ , or the memory bank) leads to substantial performance drops, validating the necessity of the full architecture.

Significance
The paper positions CIR as a significant step toward developing robust, self-interpretable Graph Neural Networks for continuous-time dynamic systems. By shifting the focus from modeling temporal dynamics alone to identifying invariant causal structures under distribution shifts, the work provides a theoretical and practical solution for scenarios where data distributions are non-stationary. The authors claim that CIR effectively mitigates spurious correlations, ensuring that predictions rely on stable causal mechanisms rather than transient environmental artifacts, thereby enhancing the reliability of dynamic graph models in real-world applications.