Scalable GNN Training via Parameter Freeze and Layer Detachment
摘要
Graph neural networks (GNNs) have been widely used in various applications. However, training large-scale GNNs with minimal cost remains an ongoing challenge. In scenarios where GPU memory is constrained, sampling a minibatch for training becomes a common practice. Nonetheless, it introduces the neighbor explosion problem. While previous sampling-based methods have made efforts to mitigate this problem, they still suffer from drawbacks such as missing information or inaccurate node representation. To address this challenge, we present MENSA, a lightweight yet powerful framework aimed to alleviate significant time and memory costs. MENSA comprises two components: EMFD and SFAA. EMFD strategically freezes trainable parameters in the first few layers to avoid unnecessary gradient computation. Additionally, it detaches these layers from the entire computational graph, effectively reducing its depth and saving GPU memory. Since fixed parameters might degrade GNN’s capability, SFAA addresses this concern by sampling a considerably smaller number of neighbors, constructing a compact computational graph, and activating all parameters for message propagation. By alternately performing EMFD and SFAA in each epoch, MENSA achieves superior \(F_1\) scores compared to other sampling-based methods. Noticeably, it exhibits 1.5x-3.5x speedup in convergence speed on most datasets and yields savings of 50%–80% in GPU memory.