Revisiting Node Affinity Prediction in Temporal Graphs

AI summary

Overview

Research area: Temporal graph learning, specifically future node affinity prediction with Temporal Graph Neural Networks (TGNNs), positioned at the intersection of sequence modeling, state space models, and ranking-based learning.

Technical level: Intermediate, with advanced theoretical sections. The architecture and empirical setup are accessible to readers familiar with graph neural networks, but the paper includes formal expressivity theorems (Theorems 1–3) whose proofs require comfort with recurrent cell equations and state space model formulations.

Scope: The paper diagnoses why simple heuristics outperform sophisticated TGNNs on node affinity prediction, then introduces a state space model-based architecture (NAViS) and a rank-based loss to close that gap on TGB and converted link prediction benchmarks.

What This Paper Is About

Future node affinity prediction asks a model to forecast a full ranked list of how strongly a node will interact with every other node at a future time, conditioned on the graph observed so far. This differs from future link prediction, which only asks whether a particular edge will appear. The paper's core problem is that recent state-of-the-art dynamic link property models, when adapted to this task, lose to trivial heuristics such as Persistent Forecast and Moving Average, and the authors set out to explain why and to build a TGNN that beats those heuristics.

Key Contributions

1. A theoretical bridge between heuristics and state space models. The authors prove that Persistent Forecast, Simple Moving Average, and Exponential Moving Average are each a special case of a linear state space model (Theorem 1: the set of heuristics is a strict subset of the set of maps realizable by linear SSMs), and use this equivalence to design an architecture that generalizes them.

2. A proof that standard recurrent memory cells cannot express Persistent Forecast. Theorem 2 shows that no parameterization of a standard RNN, LSTM, or GRU cell can produce the map where the hidden state equals the current input, because the elementwise nonlinearities bound the hidden state to ranges such as (-1, 1)^d or (0, 1)^d while the input is unbounded on R^d. This holds even when the target node is identified by its index in the source node's affinity vector.

3. A rank-based training objective for affinity prediction. The paper shows that cross-entropy is suboptimal for ranking: Theorem 3 gives a concrete example with ground truth y = [0.4, 0.6], correctly ordered logits s1 = [1, 2.4], and wrongly ordered logits s2 = [1, 0.6], where the cross-entropy loss of the correctly ranked prediction (0.78) is larger than that of the wrongly ranked one (0.75). NAViS is therefore trained with Lambda Loss plus a pairwise margin regularization term with margin hyperparameter Delta.

4. NAViS, a model combining per-node state with a virtual global state, evaluated on TGB and additional datasets. NAViS reports improvements over the best-performing TGNN, TGNv2, of +12.8% on tgbn-trade, and gains of +13.9% to +20% over the second-best TGNN on converted link prediction datasets.

Main Findings

• Heuristics beat TGNNs on this task. On TGB, Moving Avg reaches test NDCG@10 of 0.777 on tgbn-trade, 0.497 on tgbn-genre, 0.480 on tgbn-reddit, and 0.414 on tgbn-token, while TGNN baselines such as JODIE, TGAT, CAWN, TCL, GraphMixer, DyGFormer, DyGMamba, and DyRep sit far below (for example, DyGFormer at 0.388 on tgbn-trade versus Moving Avg at 0.777).

• NAViS leads on all four TGB datasets in the full-CTDG setting. NAViS test NDCG@10: 0.863 on tgbn-trade, 0.520 on tgbn-genre, 0.552 on tgbn-reddit, and 0.444 on tgbn-token, compared with TGNv2 at 0.735, 0.469, 0.507, and 0.294 respectively.

• NAViS also leads in the previous-labels-only setting. Using only prior ground-truth affinity vectors, NAViS test NDCG@10 is 0.863 on tgbn-trade, 0.528 on tgbn-genre, 0.569 on tgbn-reddit, and 0.513 on tgbn-token, compared with Moving Avg at 0.823, 0.509, 0.559, and 0.508 and Persistent Forecast at 0.855, 0.357, 0.369, and 0.430.

• Generalization to converted link prediction datasets. On Wikipedia, Flights, USLegis, and UNVote repurposed for affinity prediction, NAViS test NDCG@10 is 0.573, 0.499, 0.347, and 0.952 with full graph messages, and 0.566, 0.494, 0.326, and 0.953 using only previous labels. TGNNs again fall behind heuristics on several of these, even though they were shown to perform well on these datasets for future link prediction.

• Each design component contributes. The ablation study isolates the linear state update mechanism (versus a GRU), the global virtual vector g, and the proposed loss (versus cross-entropy); the paper states that each component contributes substantially, with full results in Appendix F.

• Dimensionality is controlled in practice. Parameters scale as O(N) with the number of nodes, so a sparsified affinity prediction pipeline retains only candidate target nodes. On tgbn-token, NAViS requires about 5,000 parameters even though the graph contains over 60,000 nodes.

• Batching is lossless by design. Because NAViS does not rely on neighbors' hidden states, it is compatible with the t-Batch mechanism, avoiding the within-batch state-transition omissions that affect memory-based TGNNs and the hard buffer truncation that affects non-memory methods.

Methodology in Plain English

The authors start by isolating four reasons TGNNs underperform simple heuristics: (i) expressivity, meaning the nonlinear updates and sampled neighborhoods of existing TGNNs cannot maintain the linear memory needed to reproduce a simple moving average; (ii) loss mismatch, since cross-entropy treats affinity outputs as a categorical distribution and ignores their ordinal ranking structure; (iii) global temporal dynamics, since affinities often depend on network-wide trends such as regime shifts that local sampling misses; and (iv) information loss, since memory-based models can drop intermediate updates inside a batch and non-memory models discard old events when a fixed-size buffer evicts them.

They then show formally that heuristics are instances of linear state space models, and that standard RNN, LSTM, and GRU cells cannot even reproduce Persistent Forecast. From this they build NAViS as a simple learnable linear state space model. Each node keeps a state h in R^d, where d is the affinity-space dimension, and there is additionally a virtual global state g in R^d. Two sigmoid gates, z_h and z_s, both constrained to [0, 1] to mirror the decay parameter alpha of an exponential moving average but computed at runtime rather than fixed, interpolate between the previous state and the new input when updating the node state and when producing the predicted affinity vector. The global vector g is computed by maintaining a buffer of the most recent previous affinity vectors across the graph and aggregating over it, with the goal of detecting a global trend before a specific node is queried.

When no previous affinity vector is available and only the interaction stream can be used, the model estimates it: initialize x-hat to zero, add each interaction weight into the entry corresponding to the destination node, normalize by the sum of entries at prediction time, and reset to zero afterward. For large graphs, entries are restricted to relevant candidate target nodes to keep parameter counts manageable.

Training uses Lambda Loss, which operates on pairs of nodes ordered by their ground-truth affinity and weights each pair by the rank displacement it would cause, approximating the gradient of non-differentiable ranking metrics. Because Lambda Loss alone proved insufficient and could push the model to shrink scores, the authors add a pairwise margin regularization term requiring a minimum margin Delta between the scores of each ordered pair.

Evaluation uses the standard protocol of a 70%-15%-15% chronological split, 50 epochs, batch size 200, reporting average NDCG@10 over three runs, across both the full fine-grained CTDG setting and the previous-ground-truth-labels-only setting. A synthetic experiment with a regime-switching latent process and nonlinear component g(t)^2 supports the argument: baselines relying only on per-node history (Persistent Forecast, SMA, EMA) or a local AR(1) SSM cannot recover the shared latent space, while NAViS, with its virtual global state, achieves the lowest error on both metrics. Full details of that experiment appear in Appendix C.

Why This Matters

For research, the paper reframes expressivity in temporal graphs around what it calls functional expressivity, the ability to represent specific mathematical operations, rather than the ability to distinguish non-isomorphic graphs under Weisfeiler-Lehman-style tests. It supplies a formal result that a widely used class of memory cells provably cannot represent the simplest competitive baseline, which is a strong argument for rethinking memory design in temporal graph learning. It also shows that the loss function, not only the architecture, is a first-order issue for ranking-style graph tasks.

Real-world applications the paper points to:

• Social platforms, where affinity between users shapes recommendations and feeds.
• Financial transaction networks, where interaction strengths between parties are forecast forward in time.
• Recommender systems and streaming services, where users and items such as movies or songs are nodes and the goal is each user's affinity to items rather than to other users.
• Token and content exchange datasets, represented by tgbn-token, which records user-token interactions.

Industry relevance comes from two directions. First, the sparsified pipeline keeps NAViS to roughly 5,000 parameters on tgbn-token despite a graph with over 60,000 nodes, which matters for deployment at scale. Second, the architecture's independence from neighbors' hidden states makes it compatible with t-Batch, meaning production-style batched processing does not silently drop state updates, a known failure mode in deployed memory-based models.

Future Directions

• Complex multi-hop dependencies. The authors explicitly state they do not claim to have solved the problem entirely and that the approach still inherits limitations, including modeling complex multi-hop dependencies. Extending the linear state design to capture such structure is a natural next step.

• Scaling to graphs with millions of nodes. Parameter count scales as O(N) in the number of nodes, and the paper's sparsification is a partial remedy based on domain knowledge about which targets matter. More general or automatic sparsification for massive graphs remains open.

• Broader exploration of loss design. The paper shows cross-entropy is misaligned with ranking in Theorem 3 and combines Lambda Loss with margin regularization, without demonstrating that this combination is optimal. Other rank-aware objectives and margin schedules could be examined.

• Transfer to adjacent tasks. The distinction between affinity prediction and weighted temporal link prediction, where a model predicts a future weighted adjacency matrix, is raised in the related work section, leaving open how much of the NAViS design transfers between the two settings.

Target Audience

This paper is most useful to researchers and practitioners working on temporal graph neural networks, dynamic link property prediction, and ranking-oriented recommendation. It suits readers who want both a theoretical explanation of why established architectures lose to moving averages on affinity tasks and a concrete, implementable alternative. Engineers deploying sequential recommendation or interaction-forecasting systems at scale will find the parameter-count and batching arguments relevant, while readers interested in expressivity theory will find Theorems 1 and 2 the most distinctive material. Source code is available at https://github.com/orfeld415/NAVIS.

Authors’ abstract

Node affinity prediction is a common task that is widely used in temporal graph learning with applications in social and financial networks, recommender systems, and more. Recent works have addressed this task by adapting state-of-the-art dynamic link property prediction models to node affinity prediction. However, simple heuristics, such as Persistent Forecast or Moving Average, outperform these models. In this work, we analyze the challenges in training current Temporal Graph Neural Networks for node affinity prediction and suggest appropriate solutions. Combining the solutions, we develop NAViS - Node Affinity prediction model using Virtual State, by exploiting the equivalence between heuristics and state space models. While promising, training NAViS is non-trivial. Therefore, we further introduce a novel loss function for node affinity prediction. We evaluate NAViS on TGB and show that it outperforms the state-of-the-art, including heuristics. Our source code is available at https://github.com/orfeld415/NAVIS