Enhancing Neural Training via a Correlated Dynamics Model

TL;DR

This paper introduces Correlation Mode Decomposition (CMD), a novel method that clusters neural network parameters based on their correlated dynamics during training, improving efficiency and generalization in large-scale models.

Contribution

The paper presents CMD, a new algorithm for capturing training dynamics through correlated parameter modes, enhancing model efficiency and generalization, and applicable to federated learning.

Findings

CMD outperforms state-of-the-art methods in modeling training dynamics.

CMD improves training efficiency and reduces communication overhead.

Preliminary experiments show benefits in federated learning scenarios.

Abstract

As neural networks grow in scale, their training becomes both computationally demanding and rich in dynamics. Amidst the flourishing interest in these training dynamics, we present a novel observation: Parameters during training exhibit intrinsic correlations over time. Capitalizing on this, we introduce Correlation Mode Decomposition (CMD). This algorithm clusters the parameter space into groups, termed modes, that display synchronized behavior across epochs. This enables CMD to efficiently represent the training dynamics of complex networks, like ResNets and Transformers, using only a few modes. Moreover, test set generalization is enhanced. We introduce an efficient CMD variant, designed to run concurrently with training. Our experiments indicate that CMD surpasses the state-of-the-art method for compactly modeled dynamics on image classification. Our modeling can improve training efficiency and lower communication overhead, as shown by our preliminary experiments in the context of federated learning.