Towards Efficient Few-shot Graph Neural Architecture Search via Partitioning Gradient Contribution

2025-06-02

AI Summary

Graph Neural Architecture Search (GNAS) suffers from gradient conflicts and computational inefficiency due to weight coupling, while existing methods are restricted to a single GNN paradigm (e.g., MPNNs or Graph Transformers). Method: We propose UGAS—a unified, differentiable GNAS framework that for the first time models weight coupling via gradient direction similarity, introducing a differentiable Gradient Contribution (GC) metric to enable efficient sub-supernetwork partitioning. UGAS is the first framework supporting joint differentiable search across both MPNNs and Graph Transformers. Results: GC-based partitioning achieves state-of-the-art performance in both quality and efficiency. Architectures discovered by UGAS significantly outperform hand-crafted models and prior NAS approaches on multiple graph benchmarks, establishing new state-of-the-art results.

Technology Category

• Machine Learning: Graph-based Machine Learning
• Search and Optimization: Learning to Search
• Multiagent Systems: Adversarial Agents

Application Category

• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs
• Search and Retrieval-Augmented AI: Efficiency and scalability of Web search engines
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

Abstract

To address the weight coupling problem, certain studies introduced few-shot Neural Architecture Search (NAS) methods, which partition the supernet into multiple sub-supernets. However, these methods often suffer from computational inefficiency and tend to provide suboptimal partitioning schemes. To address this problem more effectively, we analyze the weight coupling problem from a novel perspective, which primarily stems from distinct modules in succeeding layers imposing conflicting gradient directions on the preceding layer modules. Based on this perspective, we propose the Gradient Contribution (GC) method that efficiently computes the cosine similarity of gradient directions among modules by decomposing the Vector-Jacobian Product during supernet backpropagation. Subsequently, the modules with conflicting gradient directions are allocated to distinct sub-supernets while similar ones are grouped together. To assess the advantages of GC and address the limitations of existing Graph Neural Architecture Search methods, which are limited to searching a single type of Graph Neural Networks (Message Passing Neural Networks (MPNNs) or Graph Transformers (GTs)), we propose the Unified Graph Neural Architecture Search (UGAS) framework, which explores optimal combinations of MPNNs and GTs. The experimental results demonstrate that GC achieves state-of-the-art (SOTA) performance in supernet partitioning quality and time efficiency. In addition, the architectures searched by UGAS+GC outperform both the manually designed GNNs and those obtained by existing NAS methods. Finally, ablation studies further demonstrate the effectiveness of all proposed methods.

Problem

• Addresses weight coupling in few-shot NAS via gradient analysis
• Improves supernet partitioning efficiency with Gradient Contribution method
• Unifies MPNN and GT search in Graph Neural Architecture

Innovation

• Partitioning supernet via Gradient Contribution (GC)
• Unified Graph Neural Architecture Search (UGAS)
• Cosine similarity of gradient directions