Haste Makes Waste: A Simple Approach for Scaling Graph Neural Networks

Anonymous link: https://anonymous.4open.science/r/REST_ICML2025-0972/REST_ICML_2025_Reviewer_y9S9.pdf

A1 Counterexample of Theorem 3.1

We sincerely appreciate the reviewer's careful and thoughtful comments, and we are happy to clarify this counterexample.

Based on our understanding of the review, a special graph was constructed with only two nodes ($v_1, v_2$), where each induced subgraph contains exactly one node. Under this extreme scenario, each subgraph includes only the node itself. The reviewer further assumes that the historical embeddings exactly match the true embeddings, which leads to the gradient becoming zero under these special conditions:

$\nabla_{\hat{h}_2^{(1)}}\hat{h}_1^{(2)}= 0$

In the full batch training scenario, due to the existence of actual message passing, we have:

$\nabla_{h_2^{(1)}} h_1^{(2)} \neq 0$

Thus, the statement argues that the assumption of Lipschitz continuity may not hold.

However, we would like to clearly point out that this scenario does not align with the historical embedding or neighbor aggregation logic used in our method (and in the broader family of historical embedding methods):

Firstly, we want to clarify the definition of "out-of-batch" nodes commonly used in all historical embedding methods. Existing methods typically first sample a mini-batch (or subgraph) of nodes termed "in-batch" nodes, whose embeddings are updated every iteration. Then, they select all direct one-hop neighbors of these "in-batch" nodes as "direct one-hop out-of-batch" nodes (simply referred to as "out-of-batch" nodes in existing literature) and use their historical embeddings to approximate the true embeddings during aggregations to reduce computation bias. All other nodes that are not directly connected to the batch are simply discarded and referred to as "other out-of-batch nodes." Therefore, since the message passing is given by $h_1^{(2)} = W_2 h_2^{(1)} = W_2 \sigma(W_1 h_1^{(0)})$, this indicates that there is an edge between $v_1$ and $v_2$. if we consider $v_1$ as an in-batch node, then $v_2$ should exactly be treated as a "direct one-hop out-of-batch node" whose historical embedding is used in aggregation—even though they are not in the same subgraph—rather than simply discarding it as an "other out-of-batch node." We also provide a figure in the anonymous link that better explains the definition.

Specifically, in most historical embedding methods, the absence of node $v_2$ explicitly from the subgraph does not imply the complete loss of its influence on node $v_1$. Instead, historical embedding methods utilize historical embeddings $\hat{h}_2^{(1)}$ to preserve neighbor information, and gradients in the backward pass will still propagate through these stored embeddings. For "out-of-batch" nodes, their previously stored activation values are retrieved from memory to ensure their influence can still be incorporated. A typical form can be represented as:

$h_i^{(l)} = \text{Agg}\left(\\\{h_j^{(l-1)}\mid j \in \text{in-batch}\\\} \cup \\\{\hat{h}_k^{(l-1)}\mid k \in \text{out-of-batch}\\\}\right)$

Thus, the update of $v_1$'s embeddings still references $\hat{h}_2^{(1)}$, even if $v_2$ is not explicitly included in the sampled mini-batch/subgraph. Consequently, node $v_1$ will indeed incorporate $\hat{h}_2^{(1)}$, preserving

$\nabla_{\hat{h}_2^{(1)}}(\hat{h}_1^{(2)})\neq 0$,

and $\nabla_{W_1}L$ does not simply vanish but instead includes the gradient contributions from $\hat{h}_2^{(1)}$. Thus, the proposed counterexample does not hold under historical embedding methods.

Moreover, the Lipschitz continuity assumption regarding the gradients is a general and widely accepted assumption in the literature, such as in LMC[1]. It has consistently proven valid empirically across a wide range of realistic datasets and practical training scenarios.

A2 The standard deviation in Table 1.

We conducted multiple runs and included the standard deviation in Table 1. Please see the updated table in the anonymous link.

We appreciate your feedback and have worked diligently to address all your concerns. If you have any further questions, please let us know, and we kindly request that you consider adjusting the scores in light of our revisions.

[1] LMC: Fast Training of GNNs via Subgraph Sampling with Provable Convergence

Anonymous link: https://anonymous.4open.science/r/REST_ICML2025-0972/REST_ICML_2025_Reviewer_1EQa.pdf

A1 Proof of Theorem 3.2

Please check the detailed proof in the anonymous link . We would like to clarify that the main result in Theorem 3.2 shares foundational assumptions with previous works, which we explicitly cited in line 253. To avoid misinterpretation, we initially chose to cite the references rather than showing our proofs in submission.

A2 Difference between GAS and REST

We believe that the essential technical aspects of GAS are adequately covered in Section 2.1. Given the space constraints of the rebuttal, please refer to our main text, and we are happy to address any further questions you may have.

Difference: Overall, GAS and other related works are specific historical embedding methods, whereas REST is a novel, general training framework that effectively solves the most notable staleness issue in these methods, boosting performance and accelerating convergence. In GAS, most node embeddings remain stale because these methods inherently suffer from a frequency mismatch between cache updates and model parameter updates. Our work tackles this issue by introducing a general training framework that can be integrated with any related approach, enabling more frequent refreshing of historical embeddings and thereby eliminating staleness without altering the model architecture—rather than merely designing a model that differs from GAS.

Thus, the target of our work is not the introduction of a new model, but the development of a widely applicable training framework that addresses a key limitation—staleness—in all related works. Moreover, it provides promising solutions for other tasks involving asynchronous update scenarios in machine learning.

A3 Specific Phrases

In this work, we follow the definitions used in existing related works. Due to space constraints, we did not reiterate these definitions in our submission. However, please refer to the figure in the anonymous link for a more visual explanation.

A4 Contribution Highlight

We would like to respectfully disagree with the statement. We hope to clarify any misunderstandings regarding our objectives and highlight the significant novelty and contributions:

First, the main focus of our work is not to introduce a new historical embedding model architecture, but rather to propose a general, effective, efficient, and simple training framework that fundamentally resolves the staleness issue present in all existing methods. In other word, our model aims to address the critical shortcomings in GAS and other related works, rather than designing a model that differs from GAS. The requirement of minor implementation modification precisely demonstrates our model's strength rather than its weakness. REST exhibits strong generalizability and can significantly improve performance and accelerate convergence by eliminating staleness in historical embeddings, without requiring any changes to the underlying model.

Second, our work is the first to provide a comprehensive analysis of staleness and to reveal that a mismatch in update frequencies is its root cause—a research aspect that has not been explored before. We not only employ Theorem 1 to directly link the frequency mismatch to staleness, but also present extensive empirical results demonstrating that existing methods inevitably suffer from severe staleness issues.

Third, REST introduces a previously unexplored method: decoupling forward and backward prop so that the memory table can refresh at a frequency different than the parameter updates. To the best of our knowledge, this core innovation does not appear in GAS or any other existing method. Rather than merely being a modification, our approach can be viewed as a more general form of GAS. Our experiments demonstrate substantial performance improvements as well as significantly faster convergence. These empirical results not only validate our theoretical claims but also highlight practical advantages that GAS does not achieve. Moreover, this idea can be applied to any staleness issue arising from asynchronous updates—a common challenge in various machine learning tasks—providing promising solutions for future work that hold significant value for the entire community. Furthermore, we introduce a novel variant, REST-IS, which employs a unique importance sampling strategy to further mitigate staleness, showcasing additional methodological innovation.

Fourth, Theorem1 in our paper diverges from GAS. We also introduce Theorem 3.2, which substantiates the empirical observation that REST achieves faster convergence. This rigorous theoretical contribution clearly differentiates our work from GAS. Please refer to the anonymous link for proof.

We believe our work offers a more substantial advancement and provides promising solutions for future research, rather than merely proposing a specific model design.

We sincerely appreciate your insightful comments and thoughtful questions.

A1 Other Metrics of performance:

We choose to use accuracy as the metric in our submission since all other baselines use it, ensuring a fair comparison. Since the major task is a multi-class classification problem, we follow your suggestion and further include micro and macro F1 scores as additional metrics to better highlight the advantages of REST/REST-IS. Due to the limited time available for the rebuttal, we currently report the performance of state-of-the-art baselines, GAS and GraphFM on ogbn-arxiv and ogbn-products. We use GCN as our GNN backbone for all methods including REST. We will include similar results for other baselines on other large-scale datasets in our revision of the main text.

From the results in the following table, we observe that REST outperforms state-of-the-art baselines across various evaluation metrics on all datasets, demonstrating superior performance in large-scale training. This is consistent with the conclusions in our main submission. Moreover, while historical embedding methods still suffer from the staleness problem—resulting in significant performance drops on large-scale datasets such as ogbn-products—REST addresses the problem at its source, delivering significant performance enhancements that further validate our main text's findings.

We appreciate your feedback and have worked diligently to address all your concerns. If you have any further questions, please let us know, and we kindly request that you consider adjusting the scores in light of our revisions.

A1 Memory Efficiency:

First, we highlight that REST consistently demonstrate significant advantages over GraphSAGE in terms of memory efficiency across all datasets as shown in Table 4 and 5.

When comparing REST with GAS and VR-GCN, REST maintains a similar memory cost. This is because our work aims to introduce a novel, general training framework that can be applied to any historical embedding model to address the serious staleness issue by adjusting the execution frequency of forward propagation, rather than altering the model architecture, thereby naturally preserving the memory efficiency of the chosen baselines. This design choice makes our approach compatible with any historical embedding methods.

We want to emphasis that baseline methods come with more notable drawbacks— deteriorate performance and slow convergence—than memory efficiency. Therefore, our work focuses on resolving these critical issues to benefit all existing methods and provide a promising solution for future research, rather than further reducing memory cost.

A2 Computational Complexity:

We would like to emphasize that REST still demonstrates significant computational advantages over the baselines. In terms of time complexity, our model consistently achieves much faster convergence and requires considerably less running time compared with SOTA historical embedding baselines such as GAS (Table 4, 12; Figure 5, 6, 10, 11), VR-GCN (Table 5; Figure 12, 13), LMC (Table 7). It maintains a comparable running time to GraphSAGE while achieving superior performance. Regarding memory complexity, please refer to the previous answer. Overall, since related works are more severely impacted by staleness than by memory cost, our work focuses on addressing critical staleness problems in these methods. REST aims to enhance performance and accelerate convergence while preserving the high scalability that these methods have achieved.

A3 Applicability of REST:

We would like to clarify that the memory table being updated at a different frequency from the model is a key advantage of our method rather than a limitation. Based on our finding that frequency mismatch is the root cause of staleness, our model is designed to decouple the forward and backward operations, enabling the memory table to be refreshed at a flexible frequency that existing works cannot achieve. In other words, REST actively assigns different update frequencies, thereby alleviating the staleness. If we set the frequency $f$ to 0, it degenerates into the conventional training mode. Thus, our approach is actually an general form of the existing methods rather than an imposed restriction.

In summary, REST is not limited to any specific case; rather, it represents a novel strategy that enables flexible adjustment of the forward and backward operations to address the root cause of staleness. Moreover, it is adaptable to existing training frameworks and broader techniques, making it both highly valuable and widely applicable.

A4 Broad Impact and asynchronous scenario:

REST addresses the general issue of staleness arising from asynchronous updates—a common challenge in various scenarios, not just limit to the historical embeddings, even though our focus is on large-scale GNN training in this work. We provide several other common scenarios in which REST can be applied:

(1) Distributed Training: Asynchronous updates often arise in distributed training due to communication overhead. Nodes may operate on potentially stale global parameters from central servers, leading to parameter staleness. Applying REST in this scenario means performing extra forward passes without gradient calculation asynchronously, which refreshes local embeddings (or activation caches) more frequently. Consequently, when the node eventually computes gradients or synchronizes with the central server, the resulting update is less affected by staleness, improving overall convergence.

(2) Federated Learning: Federated learning suffers from parameter staleness due to infrequent client–server communications. REST’s decoupling strategy can be applied by letting the clients (or server) perform additional forward-only steps between global synchronizations. These extra forward passes keep local representations up to date with the client’s current model version, so that when the client finally computes gradients and communicates them back, the cached embeddings are no longer heavily stale. These additional asynchronous forward updates serve as opportunities for refreshing beyond global synchronization. This mitigates the mismatch arising from stale parameters and promotes faster convergence in federated learning.

This broader applicability indicates REST is not merely a method specific to historical embeddings, but a generally beneficial framework for addressing asynchronous update issues widely prevalent in machine learning training. We believe that this idea can offer new insights for various scenarios.