Training speedups via batching for geometric learning: an analysis of static and dynamic algorithms
This addresses training efficiency for GNN users in domains like materials science and chemistry, but it is incremental as it adapts known batching techniques from neural networks to GNNs.
The paper tackles the problem of training speed for graph neural networks (GNNs) by analyzing static and dynamic batching algorithms, finding that batching can provide up to a 2.7x speedup but the optimal algorithm depends on data and model specifics.
Graph neural networks (GNN) have shown promising results for several domains such as materials science, chemistry, and the social sciences. GNN models often contain millions of parameters, and like other neural network (NN) models, are often fed only a fraction of the graphs that make up the training dataset in batches to update model parameters. The effect of batching algorithms on training time and model performance has been thoroughly explored for NNs but not yet for GNNs. We analyze two different batching algorithms for graph based models, namely static and dynamic batching for two datasets, the QM9 dataset of small molecules and the AFLOW materials database. Our experiments show that changing the batching algorithm can provide up to a 2.7x speedup, but the fastest algorithm depends on the data, model, batch size, hardware, and number of training steps run. Experiments show that for a select number of combinations of batch size, dataset, and model, significant differences in model learning metrics are observed between static and dynamic batching algorithms.