A Locality Sensitive Hashing Based Algorithm to Accelerate Neighborhood Search in Graph Neural Operators

Neural Operators extend traditional neural networks to learn mappings between infinite-dimensional function spaces. Graph Neural Operators (GNOs), which operate on geometric structures, have shown strong performance in learning such mappings, particularly for solving partial differential equations (PDEs). These models employ message passing as part of their architecture to aggregate information from neighboring nodes, followed by a transformation step that updates node representations. However, the current approach for identifying neighboring nodes is computationally expensive, as it relies on pairwise distance calculations between all node pairs, which scales as O(N²) where N is the number of nodes. To address this, we leverage locality-sensitive hashing (LSH), a technique that hashes nearby points into the same bucket with high probability, allowing for approximate nearest neighbor search in O(N) time. We develop an LSH-based algorithm to efficiently identify node pairs within a given radius, reducing the preprocessing time for graph construction. Our method significantly accelerates neighborhood search, achieving a speedup of 29%.