Topological deep learning scales efficiently on large graphs with local batching
Batch Before You Lift: Scalable Topological Deep Learning on Large Graphs
Machine LearningArtificial Intelligence
Summary
Working with large and complex networks can be very slow because a required step called graph lifting needs a lot of memory and time. The authors developed a method called Cluster-TNN that splits these big networks into smaller parts. It then processes these parts one at a time, which cuts down the computing power needed without losing accuracy. This lets them run advanced learning methods on huge datasets that were previously too big to handle.
What this means in practice
- •For graph machine learning engineers: Train topological neural networks on very large graphs using smaller batches to reduce memory demand and enable scalable model development.
- •For social network analytics teams: Analyze complex social networks with advanced topological methods by processing partitions of massive user interaction graphs efficiently.
Authors
David Leko, Luka Benić, Guillermo Bernárdez, Nina Miolane, Olga Fink, Lev Telyatnikov
Abstract
Topological Deep Learning extends graph-based learning to higher-order domains, such as hypergraphs, cellular, and simplicial complexes. These domains are typically constructed from patterns in an input graph through a process of graph lifting. Full-domain training constructs and stores the complete lifted representation before model execution. On large and dense datasets like Reddit (233k nodes and 57.3M edges), this global materialization becomes a severe computational bottleneck, often rendering training infeasible. To address this limitation, we introduce Cluster-TNN, a domain-agnostic framework that avoids this bottleneck by lifting locally instead. After partitioning the input graph during preprocessing, at runtime Cluster-TNN dynamically samples groups of node clusters, reconstructs their induced subgraphs to form mini-batches, and applies the chosen lifting within each mini-batch. Retaining all edges among sampled nodes preserves the connectivity needed to construct higher-order structures across clusters, producing topological mini-batches that existing Topological Neural Networks can process directly. Across 21 matched comparisons with full-graph execution, Cluster-TNN reduces peak GPU memory in every configuration, by 83.2% on average while maintaining competitive predictive performance. Notably, such a reduction enables, to our knowledge, the first training of multiple different higher-order Topological Neural Networks on large datasets such as Reddit and OGBN Products. These results establish Cluster-TNN as a general strategy for scaling Topological Deep Learning beyond the limitations of global domain construction.