Leveraging the Graph Structure of Neural Network Training Dynamics

TL;DR

This paper introduces a temporal graph framework that captures neural network training dynamics, enabling early performance prediction and generalization across architectures and network sizes.

Contribution

It proposes a novel, scalable graph-based method to effectively model and analyze the evolving training dynamics of deep neural networks.

Findings

Accurately predicts task performance from early training epochs.

Captures generalizable training dynamics across different network sizes.

Effective across multiple architectures and datasets.

Abstract

Understanding the training dynamics of deep neural networks (DNNs) is important as it can lead to improved training efficiency and task performance. Recent works have demonstrated that representing the wirings of static graph cannot capture how DNNs change over the course of training. Thus, in this work, we propose a compact, expressive temporal graph framework that effectively captures the dynamics of many workhorse architectures in computer vision. Specifically, it extracts an informative summary of graph properties (e.g., eigenvector centrality) over a sequence of DNN graphs obtained during training. We demonstrate that our framework captures useful dynamics by accurately predicting trained, task performance when using a summary over early training epochs (<5) across four different architectures and two image datasets. Moreover, by using a novel, highly-scalable DNN graph representation, we also show that the proposed framework captures generalizable dynamics as summaries extracted from smaller-width networks are effective when evaluated on larger widths.