Neural Tangent Kernels Motivate Graph Neural Networks with Cross-Covariance Graphs

We thank the reviewer for their appreciation of the key theoretical contributions in our paper. Specific concerns have been addressed below.

Practical applications for specific use-cases in real-world scenarios. Our experiments originally considered the task of predicting brain functional activity in a future time step based on that in the current time step. This task forms the foundation of the time series forecasting problem in neuroimaging data. Thus, our findings suggest that the choice of cross-covariance graphs in GNNs will result in better prediction performance as compared to the GNNs that rely on the covariance graph. We expanded the scope of our experiments to test this intuition. Specifically, in revised Appendix H.5, we have provided a design of a time series prediction model that leverages information from an arbitrary length of data in the past to predict the future. For this specific application, the time series prediction improved by incorporating historical data and more importantly, GNNs with cross-covariance graphs indeed outperformed GNNs with input-based graphs.

Additional datasets. We have updated the Appendix H.7 with preliminary results on two other datasets (the PEMS07 traffic flow dataset and an ECG dataset.). For these settings, the cross-covariance based GNN models still retain a performance advantage over the alternatives.

The limits of our theoretical framework (primarily the fact that the input and output are vectors with the same dimension) certainly impose limits regarding the scope of practical tasks on which we could investigate our theoretical results. In future work, generalizing our theory beyond the setting considered in the paper will let us investigate the generalizability of our findings on a more holistic set of inference problems.

GNNs deeper than 2 layers. Our current experimental results suggest that our theoretical findings extend to GNNs deeper than two layers (see Fig. 4. that shows the effect of increasing the depth of the model). Even as the depth increases beyond two layers, the cross-covariance based GNNs consistently converge faster and achieve a smaller loss). Additionally, the fact that NTK-based results in the literature are usually generalizable to deeper models makes us optimistic that our results can also be generalized to deeper models. While this was not a focus of the current paper, it is an avenue we aim to explore theoretically in future work.

Future directions, limitations, and challenges. As discussed previously, extending the theoretical analyses to GNNs deeper than 2 layers is an immediate direction of interest. We expect this analysis to have novel theoretical challenges pertaining to intermediate GNN layers. Furthermore, we will aim to expand the scope of our results for inference tasks beyond the considered setting of the multi-variate regression problem (for instance, to node-level inference tasks and classification). Another limitation of our theoretical contributions is the lack of thorough analysis of the tightness of the lower bounds on alignment, i.e., the gap between ${\cal A}$ and ${\cal A}_L$.

We hope that we addressed the reviewer's concerns sufficiently, in which case, we would be grateful if your rating of our paper could be re-evaluated. We would be happy to clarify any additional concerns.

Missing literature. We thank the reviewer for suggesting the work in Huang et. al, 2021 as a relevant study to be discussed in our paper. We note that this study leverages the analyses of NTKs to investigate the impact of number of layers in GNNs on their inference. Broadly, this work leverages NTK of a GNN to motivate various architectural choices and in this respect, is similar in spirit to our work. We have discussed such works in our literature review and will add this paper to the discussion as well. However, we firmly believe that the analytical approaches and motivations relevant to our work and those in Huang et. al, 2021 are very distinct and do not overlap in scope.

Node classification task. While the extension of the definition of Alignment to tasks like node classification is straightforward, generalizing our results to such cases is not trivial and attempting to do so would vastly expand the scope of the paper. We have therefore left attempting such an extension for future work.

Ethical concern. We have provided the acknowledgement for HCP-YA dataset as per the data use agreement. The authors are not among the PI names listed there and have no collaboration with them for this work.

We hope that we addressed the reviewer's concerns sufficiently, in which case, we would be grateful if your rating of our paper could be re-evaluated. We would be happy to clarify any additional concerns.

We thank the reviewer for their insightful comments. Specific concerns have been addressed below.

Alignment, NTKs, and GNNs. We clarify that the notion of alignment (defined as $y^{\sf T} \Theta y$ for NTK $\Theta$ in equation (2) and Definition 1, where $y$ is the output) is a generic property of learning with gradient descent and not specific to only the GNN architecture or related inference tasks like node classification. Alignment dictates the convergence of gradient descent (as established in equation (9)). Our objective was to leverage (i) these facts about alignment and gradient descent, and (ii) specific analytical form of the NTKs for a GNN architecture; to demonstrate that optimizing alignment for a GNN is equivalent to optimizing the choice of graph shift operator for the supervised learning task. To the best of our knowledge, no existing work has motivated the practical choice of graph shift operator in GNNs using either these arguments or the theoretical analyses presented in our paper.

Same input and output dimensionality. The input and the output dimensionality were assumed to be the same for analytical tractability. Such settings are of practical relevance in spatio-temporal datasets, where the multiple correlated features evolve over time (for instance, in datasets that describe brain activity, traffic flow, stock prices, etc.). In these settings, the use of past data to predict the future data is a learning task of interest.

Experiments. Our theoretical results (Theorem 2 and 3) explicitly establish the relationship between the graph shift operator and cross-covariance graphs in the context of optimizing the alignment function. In this respect, our experiments were focused on validating this insight on a real dataset for a suitable practical learning task.

Choice of HCP-YA dataset. HCP-YA dataset is one of the larger datasets available for brain functional activity in terms of sample size and has been used extensively by the research community. Thus, this dataset was chosen for both its suitability to the considered setting in our theoretical analyses and in the spirit of reproducibility of our findings.

Validation of theoretical results. To validate the insights drawn from theoretical results, we considered the inference task of predicting the brain activity at the future using the current time step. This inference task forms the foundation for an algorithm that addresses a typical time series forecasting problem. Here, we compared the performance of GNNs modeled by cross-covariance graphs with GNNs that leveraged only the covariance graph estimated from the input data. As predicted by our theoretical results, the GNNs with cross-covariance graphs indeed outperformed those with covariance graphs.

Rigor of the Experiments. We note that we validated these findings across more than 1000 individual subjects in the dataset and a noticeable performance gain can be seen on average and has been reported in our results. The effects of changing different architectural parameters such as the number of layers, number of filter taps and the nonlinearity used in the GNNs were considered in the experiments reported in the appendices, where we observed that the benefit of using cross-covariance graphs remained consistent throughout.

Additional Experiments. We have further expanded our empirical evaluations (see Appendix H.6) by adding two more datasets and comparing GNNs with cross-covariance graphs with additional baselines.

• Additional baselines. These additional baselines include a graph based on the Pearson Correlation between the features at different nodes (which is commonly used in tasks pertaining to brain data), and another graph with edges based on the Euclidean distance between node features passed through a Gaussian kernel. (See ‘Data-driven graph construction and graph learning: A review’ by Qiao et al. for a comprehensive review of many such input-based graph construction methods). We observed that in general, the models with graphs constructed using only the input features with both linear and non-linear affinities exhibit similar performance and the models using the cross-covariance, consistently outperform the input-based graphs.

• More datasets. To ensure that these results don’t just pertain to one particular dataset, we have updated the Appendix (H7) with preliminary results on two other datasets (the PEMS07 traffic flow dataset and an ECG dataset.). While the performance gap between the different models isn’t as pronounced in these other datasets compared to HCP, the cross-covariance based models still retain an advantage.

We thank the reviewer for their insightful comments. Specific concerns are addressed below.

Theoretical analysis. Here, we emphasize that our theory has a bottom-up characteristic, where we have begun with the analysis of a simple graph convolutional filter$^1$ as the predictor. Due to inherent linearities in the graph filter, the NTK for a graph filter has, by default, a constant behavior with respect to learnable parameters. The analysis of alignment using NTK in this setting leads naturally to the conclusion of cross-covariance graphs being the optimal choice. Our subsequent analyses adds various different complexities, such as, non-linearities and extension to two layers for a GNN in practice. For the analysis of the GNN model, we operate under the standard assumptions of infinite width associated with NTKs and obtain conclusions similar to that for the simple graph filter setting. Therefore, due to the bottom up characteristic of our analysis and consistency of the conclusions for graph filters and GNNs, we do not believe that our theory overlooked any missing phenomenon that could have led to a drastically different choice of kernel for the considered regression task.

$^1$ This graph convolutional filter is the fundamental information processing block in the general GNN (or specifically, GCN) architecture considered in our paper.

Cross-covariance graphs versus graphs with non-linear affinities. Regarding the concern that graphs with non-linear affinities may outperform covariance-based graphs, we have performed additional experiments. We refer the reviewer to Appendix H.6 in the revised paper, where we have added the performance comparison between models with different graphs. These include a graph based on the Pearson Correlation between the features at different nodes (which is commonly used in tasks pertaining to brain data), and another graph with edges based on the Euclidean distance between node features passed through a Gaussian kernel. (See ‘Data-driven graph construction and graph learning: A review’ by Qiao et al. for a comprehensive review of many such input-based graph construction methods). We observed that in general, the models with graphs constructed using only the input features with both linear and non-linear affinities exhibit similar performance and the models using the cross-covariance consistently outperform the input-based graphs. As another baseline, we have also added comparison of the performance of our models to a Fully-connected Neural Network (FCNN; which has roughly 100 times more number of learnable parameters as compared with GNN models) and we observe that the two-layer GNN with cross-covariance graph exhibits comparable performance to the FCNN.

Choice of GNN architecture. We clarify that our analysis focuses on graph convolutional networks (GCNs) among the different variants of GNNs that exist in the literature. Using this setting, we have demonstrated how to leverage the theory-inspired insights into the construction for a graph from the data. Hence, our experiments have focused on the setting consistent with that studied in the theoretical analysis to validate our results. We conjecture that a similar performance advantage will appear when using the cross-covariance matrix as a graph in more complex graph-based architectures. Expanding our theory and experiments to show the advantage of using cross-covariance in more complex models is certainly a direction that we will pursue in future work.

Dear Reviewers,

Thanks again for your invaluable feedback.

Since the author-reviewer discussion period is near its conclusion, we hope to ensure that our responses have adequately addressed your concerns. If you have any further questions or require any additional information from us, please let us know.

Best regards,

The Authors