Training Dynamics of Deep Network Linear Regions

Abstract: The study of Deep Network (DN) training dynamics has largely focused on the evolution of the loss function, evaluated on or around train and test set data points. In fact, many DN phenomenon were first introduced in literature with that respect, e.g., double descent, grokking. In this study, we look at the training dynamics of the input space partition or linear regions formed by continuous piecewise affine DNs, e.g., networks with (leaky)ReLU nonlinearities. First, we present a novel statistic that encompasses the local complexity (LC) of the DN based on the concentration of linear regions inside arbitrary dimensional neighborhoods around data points. We observe that during training, the LC around data points undergoes a number of phases, starting with a decreasing trend after initialization, followed by an ascent and ending with a final descending trend. Using exact visualization methods, we come across the perplexing observation that during the final LC descent phase of training, linear regions migrate away from training and test samples towards the decision boundary, making the DN input-output nearly linear everywhere else. We also observe that the different LC phases are closely related to the memorization and generalization performance of the DN, especially during grokking.

• The paper introduces a novel local complexity metric to quantify how linear regions in deep networks evolve during training.
• It identifies distinct phases where linear regions initially decline, then rise, and finally simplify as decision boundaries form.
• The study demonstrates that architecture choices and batch normalization significantly impact LC dynamics and overall network generalization.

An Analysis of Training Dynamics in Deep Network Linear Regions

The paper "Training Dynamics of Deep Network Linear Regions" by Ahmed Imtiaz Humayun and colleagues provides a comprehensive exploration of the input space partitioning dynamics during the training of deep networks (DNs) with continuous piecewise affine structures, such as those employing (leaky-)ReLU nonlinearities. The authors propose a novel metric, local complexity (LC), to quantify the complexity of DNs by examining the concentration and dynamics of linear regions within the input space throughout the model's training process.

Overview

The core focus of the paper is the investigation of the linear regions' dynamics during training. This is distinguished from most conventional studies that primarily analyze the loss function over training and test sets. The proposed local complexity measure captures the expressivity of the network as regions in the input space undergo transformations during training.

The paper identifies three key phases in the progression of local complexity around data points: an initial decrease post-initialization, a subsequent rise, and a final decline. Notably, the decline in the final phase is associated with the migration of linear regions towards the decision boundary, thus simplifying the network's input-output mapping elsewhere in the input space.

Strong Numerical Results and Claims

The empirical findings indicate that during the ascent phase, LC demonstrates substantial divergence between training and test datasets. This behavior points towards different memorization and generalization characteristics, particularly during grokking. The sharp ascent in LC for training data suggests substantial network memorization. Conversely, the final descent in LC correlates with emerging generalization capabilities as decision boundaries form and linear regions move away from data points.

Moreover, the study observes that architectures and regularizers impact LC dynamics markedly. For instance, deeper architectures diminish the sharpness of LC ascent and reduce train-test disparities in the ascent peak, reflecting lesser memorization propensity. Additionally, the use of batch normalization (BN) effectively eliminates the final LC descent by keeping region boundaries close to training samples, consistent with its role of stabilizing training dynamics and improving generalization.

Theoretical and Practical Implications

The theoretical implications of these findings suggest that understanding the dynamics of linear regions provides insights into the expressivity and functional capacity of DNs. Such understanding could guide the architectural design and training practices to harness the full potential of DNs, ensuring a balance between memorization and generalization.

Practically, the results hint at potential strategies to enhance model robustness and reduce overfitting by manipulating LC directly through hyperparameters and network architecture choices. The migration of linear regions towards decision boundaries could further be exploited for adversarial robustness, where a network's input space mapping is more predictably linear except near critical decision regions.

Future Directions

The study opens new avenues for further explorations, such as the development of advanced metrics to precisely quantify expressivity and decision boundary dynamics in deeper or more complex architectures beyond piecewise affine networks. Furthermore, expanding on this understanding could lead to improved regularization techniques that control partition dynamics intentionally, fostering more flexible and interpretable models in diverse AI applications.

In summary, this work provides substantial contributions to the understanding of DN training dynamics, emphasizing the significance of partition complexity in shaping model capabilities. It underscores the interconnectedness between training dynamics, network architecture, and regularization strategies, providing an enriched framework for the exploration of expressivity within DNs.