Paper: On the Generalization of Temporal Graph Learning with Theoretical Insights (ICLR-2024)

Reviewer comment (Theory):
Theorem 1: The dynamics with respect to m* are not discussed, hence do not allow to assess whether the chosen \( m \geq m^* \) in the experiments is relevant.

Author response:
**Q1. The dynamics w.r.t. $m^\star$ in Theorem.**

Thank you for your question. Increasing neural network hidden dimension $m$ will not affect the upper bound in Theorem 1, but it will only affect the closeness of the gradient descent solution to the random initialization, i.e., affect $\\| \widetilde{\boldsymbol{\theta}} - \boldsymbol{\theta}_0 \\|\_\mathrm{F}$. More specifically, the larger the $m$, the smaller the $\omega$ would be in Lemmas 1 to 6, and the closer $\widetilde{\boldsymbol{\theta}}$ to $\boldsymbol{\theta}_0$. This phenomenon is also know as "lazy learning" where weights only move very slightly during training in over-parameterized regime. Please refer to Section 2 of [1] for a more detailed discussion. We will clarify this in the revised version.

[1] Bietti, A., and Mairal, J., On the inductive bias of neural tangent kernels. NeurIPS 2019.

Reviewer comment (Experiments):
Could you further discuss the need for the LayerNormalization in SToNe? For example, instead of no normalization or BatchNormalization?

Author response:
**Q4. Discussion the need for Layer Normalization?**

Thank you for your question. We found that using layer normalization (LN) can help improve the stability and consistency of our model's performance when repeating experiments multiple times. This is because LN normalizes the activations within a neural network layer, making the model output less sensitive to changes in input data.
We also experimented with batch normalization, which gave similar results. However, since batch normalization computes scaling weights from training data, it may introduce potential temporal information leakage. To avoid this issue, we prefer to use LN instead of batch normalization. We will provide more details on this in our final version.

Reviewer comment (Novelty):
The motivation to define a new indicator for the generalization error is unclear: In the paper, the authors indicate that the generalization error bound decreases with respect to the number of training data, but increases with respect to the number of layers/steps in the neural-network-based methods and the feature-label alignment (FLA). Thus, FLA can be used as a proxy for the expressive power measurement. Although there is a connection between FLA and generalization error, using FLA can be unnecessary unless unique advantages of FLA exist, e.g., it is easier to measure.

Author response:
**Q1. Although there is a connection between FLA and generalization error, using FLA can be unnecessary unless unique advantages of FLA exist.**

Thank you for your comment. Please note that FLA is necessary if we want to understand the selection of input data and labels to the model performance. Existing generalization analysis cannot achieve this. For more information, please refer to our Section 2 paragraph “generalization and expressive power” and Appendix I.2.

Reviewer comment (Novelty):
The uniqueness of the temporal graph learning algorithms, compared to static graph learning algorithms, is not reflected: Based on the definition 1, FLA can be used for all gradient-based algorithms. It is not specific to temporal graph learning methods.

Author response:
**Q2. The uniqueness of the temporal graph learning algorithms, compared to static graph learning algorithms, is not reflected.**

We conducted a theoretical analysis of GNN-based, RNN-based, and Memory-based temporal graph learning methods. These methods are distinct from static graph learning algorithms because they take into account temporal information during training.

Reviewer comment (Novelty):
The proposed method is essentially a GNN method augmented with input data selection, hence has limited novelty as a new algorithm: Based on Section 4.1, the proposed method comprises three parts: input data preparation, encoding features via GNN and link prediction via MLP, which is similar to existing methods except for the feature engineering component.

Author response:
**Q3. The proposed method is essentially a GNN method augmented with input data selection, hence having limited novelty as a new algorithm.**

Our proposal is novel because (1) it is inspired by key insights of our theoretical analysis of temporal graph methods, (2) it has a small generalization error, (3) it performs well empirically with low model complexity. We have summarized the differences compared to existing methods in Section 4.2.

Reviewer comment (Evaluation):
The performance resulting from the proposed method in the experiment is not impressive empirically: some existing methods, e.g., TGN and GraphMixer, have a similar performance. Their average precision on the first five datasets are very close. Although the proposed method has better performance on the last dataset UCI, it is unclear whether this last dataset has some special data distribution or setting.

Author response:
**Q4. The performance resulting from the proposed method in the experiment is not impressive empirically.**

Please note that our method not only closely matches the performance of state-of-the-art temporal graph learning methods, but also offers a better theoretical understanding. Additionally, it has significantly lower model complexity, as shown in Table 2.

Reviewer comment (Evaluation):
What’s the advantage of using FLA? In experiments, the authors still use the generalization gap, i.e., the absolute difference between the training and validation average precision scores, to compare different algorithms.

Author response:
**Q5. What’s the advantage of using FLA? In experiments, the authors still use the generalization gap.**

Please note that FLA is a component of the generalization bound (i.e., the variable $R$ in Theorem 1), which tells us how data selection affect generalization. 
Other terms in Theorem 1 provide insights into the effect of model architecture on generalization ability.
Therefore, to gain a more comprehensive understanding, we must also take into account the generalization gap.

Reviewer comment (Theory):
Can FLA be used for the generalization error of static graph learning algorithms? If so, what’s the uniqueness of the temporal graph learning algorithms in the aspect of generalization error?

Author response:
**Q2. The uniqueness of the temporal graph learning algorithms, compared to static graph learning algorithms, is not reflected.**

We conducted a theoretical analysis of GNN-based, RNN-based, and Memory-based temporal graph learning methods. These methods are distinct from static graph learning algorithms because they take into account temporal information during training.

Reviewer comment (Novelty):
I think the discussion of the existing generalized bounds for GNN and so on should be elaborated on more in the literature review part and the comparison with the novel bound derived in this paper should be provided.

Author response:
Thank you for the suggestion. Most of the existing generalization bounds (e.g., Rademacher complexity, PAC-Bayesian, and Uniform Stability bounds) are data-independent and rely solely on factors like the number of layers or hidden dimensions. In contrast, our feature-label alignment generalization bound, grounded in deep learning theory, is data-dependent and operates under different assumptions. For instance, the aforementioned bounds require assumptions on weight matrix norms during training, whereas our deep learning theory bounds do not necessitate such assumptions. Consequently, direct comparisons between these results may not be practical.

Reviewer comment (Experiments):
Based on your findings, SToNe outperforms other methods due to its smaller FLA and thus smaller generalization bound. I am curious about the comparison of FLA of the different methods in those real datasets. The current experiments to illustrate the effect of FLA are all on SToNe itself.

Author response:
Please refer to Figure 3. As shown in Figure 3, we know  

(1) JODIE has a relatively larger FLA score than most of the other TGL methods. This is potentially due to “stop gradient” operation that prevents gradients from flowing through the memory blocks and could potentially hurt the expressive power.  

(2) APAN and TGN alleviate the expressive power degradation issue by applying a single layer GNN on top of the memory blocks, resulting in a smaller FLA than the pure memory-based method.  

(3)TGAT has a relatively smaller FLA score than other methods,which is expected since GNN has achieved outstanding performance on static graphs.  

(4) DySAT is originally designed for snapshot graphs, so its FLA score might be highly dependent on the selection of time-span size when converting a temporal graph to snapshot graphs. A non-optimal choice of time-span might cause information loss, which partially explains why DySAT’s FLA is large.

Reviewer comment (Theory):
The proof of lemma 5 relies on the upper bound of \( \|\partial loss_i(\theta)/\partial W^{(l)}\| \) for \( \theta \) from the iteration of SGD. This upper bound is obtained in lemma 4, which is built on lemma 8 from Zhu et al. (2022). However, in that paper, they assumed that \( W \) is a gaussian random matrix, which is different in this paper (note that in this case, \( W \) is obtained from SGD). Hence, I don't think that lemma can be applied here and there are some flaws in upper bounding \( \|\partial loss_i(\theta)/\partial W^{(l)}\| \).

Author response:
Please note that Lemma 8 of Zhu et al. (2022) only requires that the weight parameters be sufficiently close to the initialized weight parameters. In fact, their Lemma 8 is also applied to the output of SGD parameters in their Lemma 11.  

We apologize for any confusion and aim to provide additional clarity as outlined below (primarily due to incorrect notation and missing details in the proofs of [1], but the final statements in [1] remain unchanged).  

As outlined in our Appendix C.1, the gist of the proof is:  

- (i) We show that all the iterates of SGD $\{\boldsymbol{\theta}_i|i=1,\ldots, N\}$  and optimal solution $\boldsymbol{\theta}^*$ lie in a ball of radius $\omega$ around the initial weights $\boldsymbol{\theta}_0$ (which is a random matrix and its characteristics can be quantified with high probability, e.g., $\|\mathbf{D}\|_2 \leq \rho, \|\boldsymbol{f}_{i, \ell-1}\| \leq \Theta(1), \|\mathbf{W}_0^{(\ell)}\|_2 \leq 3, \forall \ell \in [L-1], \|\boldsymbol{W}_0^{(L)}\|_2 \leq \sqrt{2}$,   

- (ii) We utilize the the bounds in (i) to upper bound the desired quantities about SGD iterates (e.g., for any $\boldsymbol{\theta}_i$ we can show $\| \mathbf{W}_i^{(\ell)} \|_2 = \| \mathbf{W}_i^{(\ell)} - \mathbf{W}_0^{(\ell)} + \mathbf{W}_0^{(\ell)} \|_2 \leq \|\mathbf{W}_i^{(\ell)} - \mathbf{W}_0^{(\ell)} \|_{\textrm{F}}+ \|\mathbf{W}_0^{(\ell)} \|_2 \leq (\omega + 3)$ or similarly $\|\mathbf{W}_i^{(L)}\| \leq (\omega + \sqrt{2})$, similarly for other quantities,   

- (iii) By proper choice of number of neurons etc we can pick $\omega$  to achieve the desired accuracy. Therefore, by guaranteeing the $\omega$-closeness of SGD iterates to  initialization, we only require the randomness of initial weights to bound the necessary quantities related to SGD iterates.  

The Lemma 8 of [1] states that for any $\mathbf{W}, \mathbf{W}^\prime$ close to randomly initialized weight parameters $\mathbf{W}^{(0)}$ with a distance of at most $\omega=1/(3\text{Lip}+1)^{L-1}$,  with high probability, neural network function is almost linear in terms of its weights, i.e., $|f(x_i, \mathbf{W}^\prime) - f(x_i, \mathbf{W}) - \langle \nabla f(x_i, \mathbf{W}), \mathbf{W}^\prime - \mathbf{W} \rangle | = \mathcal{O}(1) $.   
However, in the proof  where the authors utilize the chi-square distribution to bound $\|\mathbf{W}_L\|_2$,  it should be about $\|\mathbf{W}^{(0)}_L\|_2$ instead of $\|\mathbf{W}_L\|_2$, which caused the confusion you pointed out.   
Please note that in the final inequality of Eq. 41 in [1], we need to bound $\|\mathbf{W}_L\|_2$, $\|\mathbf{W}'_L\|_2$, $\|\mathbf{D}_{i,r}\|_2$,  $\|\mathbf{W}'_{\ell} - \mathbf{W}_{\ell}\|_2$, and $\|\boldsymbol{f}_{i, \ell-1}\|_2$ etc,   
where $\mathbf{W}$ and $\mathbf{W}'$ are arbitrary matrices (SGD parameters) but close enough to the random initialization weights $\mathbf{W}^{(0)}$. By having a bound on $\|\mathbf{W}^{(0)}_L\|_2$, $\|\mathbf{W}^{(0)}_r\|_2, r \in [L-1]$, $\|\mathbf{D}\|_2$ (which can bounded at initial random parameters with high probability),  under the assumption that $\mathbf{W}, \mathbf{W}' \in \mathcal{B}(\mathbf{W}^{(0)}, \omega)$, we can bound these quantities about $\mathbf{W}, \mathbf{W}' $ as pointed in step (ii) above, which leads to the final inequality at the end of the proof of Lemma 8 in [1]. A similar argument holds for bounding $\| \nabla_{\mathbf{W}_\ell} L_i(\mathbf{W}^{(i)}) \|_2$ as in the proof of Lemma 10 in [1] (here authors of [1] used the correct notation $\|\mathbf{W}^{(0)}_r\|_2 \leq 3$ and $\|\mathbf{W}^{(0)}_L\|_2 \leq \sqrt{2}$ etc to bound $\|\mathbf{W}^{}_L\|_2$ and other quantities about arbitrary $\mathbf{W} \in \mathcal{B}(\mathbf{W}^{(0)}, \omega)$ to get the final upper bound). We note that an alternative approach is to directly bound these quantities using the machinery developed in [2] (as in lazy training or NTK), but we found the bounding through $\mathbf{W}^{(0)}$ more convenient.  We'd gladly incorporate comprehensive proof details for these lemmas, ensuring correct notation for thoroughness.  

[1] Generalization Properties of NAS under Activation and Skip Connection Search  

[2] A Convergence Theory for Deep Learning via Over-Parameterization

Reviewer comment (Theory):
In the proof of Theorem 2, the part of constructing \( f'\in \mathcal{F}(\theta_0,R) \) such that \( \psi \) can be upper bounded by \( 1/\sqrt{N} \) is confusing (e.g., why \( \langle \nabla_\theta f_i(\theta_0),\theta \rangle =y_i(B+B') \)). I can't get the point how \( R \) can make a difference in the result.

Author response:
The feature-label alignment  $R$ controls the range we select  $f^\prime \in \mathcal{F}(\boldsymbol{\theta}\_0, R)$ such that $\frac{1}{N} \sum\_{i=1}^N \psi(y\_i f\_i^\prime) = O(1/\sqrt{N})$.  
We need $\frac{1}{N} \sum\_{i=1}^N \psi(y\_i f\_i^\prime) = O(1/\sqrt{N})$ because the second and third term in Lemma 6 is $O(1/\sqrt{N})$.  

To achieve this goal, according to the definition of $\mathcal{F}(\boldsymbol{\theta}\_0, R)=\{ f(\boldsymbol{\theta}\_0) + \langle \nabla f(\boldsymbol{\theta}\_0) , \boldsymbol{\theta}\rangle | \boldsymbol{\theta} \in \mathcal{B}(0, R/\sqrt{m}) \}$, we have to select $R$ and $\boldsymbol{\theta} \in \mathcal{B}(0, R/\sqrt{m})$ such that   
$$
    \langle \nabla f\_i(\boldsymbol{\theta}\_0) , \boldsymbol{\theta}\rangle = y\_i \Big( \underbrace{- \log(\exp(N^{-1/2})-1)}\_{B} + \underbrace{\max_i |f_i(\boldsymbol{\theta}_0)|}\_{B^\prime} \Big)
$$  
As detailed in the proof, one solution of $R$ and $\boldsymbol{\theta}$ is   
\begin{equation*}
    \text{vec}(\boldsymbol{\theta}) = \mathbf{Q} \Lambda^{-1} \mathbf{P}^\top \hat{\mathbf{y}} \text{ and }
    R = \mathcal{O}\left(\sqrt{\mathbf{y}^\top (\mathbf{J\mathbf{J}^\top})^{-1} \mathbf{y}}\right),
\end{equation*}
where this $R$ is our feature-label alignment.

Reviewer comment (Theory):
I do not see the formal and strong theoretic claim that the submission frames its theoretic contribution as strong. Specifically, since the given result only established an upper bound, I believe strictly speaking it can merely indicate desirable aspects of methods for temporal graph learning, and not really explain the difference in performance. This is because we have no lower bound relating the generalization error to e.g. the number of layers. This means a small number of layers gives us a good bound on the expected error with high probability, but a larger number of layers does not necessarily lead to a large expected error.

Author response:
**Q1. No lower bound relating the generalization error to e.g. the number of layers. The framing of the theoretical contributions: the results only have an upper bound, but lower bound is also expected.**


Thank you for your suggestion. We fully agree that studying the lower bound is also an interesting and important future direction. 
However, it's worth noting that lower bound analysis is specific to each case (e.g., a specific number of layers, architecture), which may not be feasible if we want to examine the effects of neural architecture selection and data selection on generalization ability.
To the best of our knowledge, there are no previous works on the generalization lower bound of deep neural networks (previous studies only applies to 2-layer MLP [1] or decision tree [2]), let alone temporal graph learning methods.
Our paper primarily focuses on the upper bound of generalization, which captures the worst-case generalization ability of various methods. 

[1] Lower Bounds on the Generalization Error of Nonlinear Learning Models

[2] Margin-Based Generalization Lower Bounds for Boosted Classifiers

Reviewer comment (Theory):
Can you discuss my concerns about the framing of the theoretical contributions?

Author response:
We are happy that you found our response to the second question convincing. 

With regards to implications of upper bound, we totally agree with you that the statement of our claim should be rewritten as we are not aware of any lower bound  to avoid conveying the wrong message. We will gladly incorporate your comment and make the implications of obtained upper bound  crystal clear in the subsequent version.

Reviewer comment (Theory):
Can you give a formal argument of why memory blocks increase the FLA? Is this always true? (In particular my question relates to the second bullet on the top of page 7.)

Author response:
**Q2: Give a formal argument of why memory blocks increase the FLA.**

Feature-label alignment (FLA) increases in memory-based methods potentially because they (1) utilize the stop-gradient operation and (2) aggregate an excessive amount of temporal neighbor information.

As a reminder, FLA is defined as $\mathbf{y}^\top (\mathbf{J} \mathbf{J}^\top)^{-1} \mathbf{y}$, where $\mathbf{J}$ represents a stack of gradients computed for each data point using the randomly initialized weight parameters. Since memory-based methods aggregate information across all time steps, they involve that random weight parameters significantly more times than GNN-based or RNN-based methods. Consequently, their FLA is expected to have a larger lower bound because $\mathbf{y}^\top (\mathbf{J} \mathbf{J}^\top)^{-1} \mathbf{y} \geq 1/\lambda_{\max}(\mathbf{J} \mathbf{J}^\top)$  and $\lambda_{\max}(\mathbf{J} \mathbf{J}^\top)$ is smaller.
To illustrate why involving more random initialized weight parameters could result in a smaller $\lambda_{\max}(\mathbf{J} \mathbf{J}^\top)$, we use a toy example with the stop-gradient operator:

$$
\hat{y}\_i = \text{stopgrad}(\mathbf{s}(\mathbf{x}\_i))^\top \mathbf{v},  \mathbf{s}(\mathbf{x}\_i) = \sigma(\ldots \sigma( \sigma(\mathbf{x}\_i \mathbf{W}) \mathbf{W})\ldots \mathbf{W} )
$$

By definition, the feature-label alignment is computed as:
$$
    \mathbf{y}^\top(\mathbf{J} \mathbf{J}^\top)^{-1} \mathbf{y} = \mathbf{y}^\top\left( \sum\_{i=1}^N \mathbf{s}(\mathbf{x}\_i)\mathbf{s}(\mathbf{x}\_i)^\top \right)^{-1} \mathbf{y} \geq \| \mathbf{y} \|\_2^2 / \lambda\_{\max}\left(\sum\_{i=1}^N \mathbf{s}(\mathbf{x}\_i)\mathbf{s}(\mathbf{x}\_i)^\top\right)
$$
In general, the more random initialized weight parameters $\mathbf{W}$ are involved, the more randomness there is in $\mathbf{s}(\mathbf{x}\_i)$, leading to a smaller largest eigenvalue of $\sum\_{i=1}^N \mathbf{s}(\mathbf{x}\_i) \mathbf{s}(\mathbf{x}\_i)^\top$.