When do Convolutional Neural Networks Stop Learning?

TL;DR

This work tackles the challenge of predicting when CNN training should stop without relying on a validation set. It introduces a layer-wise stability measure, formalized via $\alpha_n^t = \sigma(X_n^t)$, stability vectors $S_n^e$, and epoch-wise indicators $\mu_n^e$ and $\delta_n^e$, to detect near-optimal learning capacity when $\sum_{i=1}^n \delta_i^e =0$ across consecutive epochs. Implemented as a plug-and-play module with no trainable parameters, the approach was tested on six CNN variants across CIFAR10, CIFAR100, SVHN, and ten MedMNIST-V2 medical datasets, achieving average computational time savings of about $58.49\%$ on general images and $44.1\%$ on medical data without accuracy loss. The results are supported by ablation studies and generalization analyses showing train loss does not reliably predict generalization, and near-optimal epochs occur well before conventional 200/100-epoch baselines.

Abstract

Convolutional Neural Networks (CNNs) have demonstrated outstanding performance in computer vision tasks such as image classification, detection, segmentation, and medical image analysis. In general, an arbitrary number of epochs is used to train such neural networks. In a single epoch, the entire training data -- divided by batch size -- are fed to the network. In practice, validation error with training loss is used to estimate the neural network's generalization, which indicates the optimal learning capacity of the network. Current practice is to stop training when the training loss decreases and the gap between training and validation error increases (i.e., the generalization gap) to avoid overfitting. However, this is a trial-and-error-based approach which raises a critical question: Is it possible to estimate when neural networks stop learning based on training data? This research work introduces a hypothesis that analyzes the data variation across all the layers of a CNN variant to anticipate its near-optimal learning capacity. In the training phase, we use our hypothesis to anticipate the near-optimal learning capacity of a CNN variant without using any validation data. Our hypothesis can be deployed as a plug-and-play to any existing CNN variant without introducing additional trainable parameters to the network. We test our hypothesis on six different CNN variants and three different general image datasets (CIFAR10, CIFAR100, and SVHN). The result based on these CNN variants and datasets shows that our hypothesis saves 58.49\% of computational time (on average) in training. We further conduct our hypothesis on ten medical image datasets and compared with the MedMNIST-V2 benchmark. Based on our experimental result, we save $\approx$ 44.1\% of computational time without losing accuracy against the MedMNIST-V2 benchmark.

Figures (8)

• Figure 1: Top dotted box represents traditional steps of training a CNN variant. At each epoch, our plugin (bottom dotted box) measures data variation after convolution operations. Based on all the layers data variation, the plugin decides the continuity of training.
• Figure 2: At $t$-th iteration, the process of computing stability values $\alpha_1^t, \alpha_2^t, \ldots,\alpha_n^t$ for 1 to $n$ layers.
• Figure 3: At $e$-th epoch, the process of constructing stability vectors $S_1^e$, $S_2^e$, …, $S_n^e$ for 1 to $n$ layers.
• Figure 4: The cross entropy loss (top) and the validation error (bottom) are shown up to 200 epochs for ResNet18 on the CIFAR10 dataset.
• Figure 5: The horizontal axis shows the epoch number (ranging from 10--200) used to train the ResNet18, CNN, and VGG16 on the CIFAR10 dataset. The vertical axis shows the testing accuracy of those models. The X mark shows the testing accuracy and the epoch number to train a CNN variant based on the near-optimal learning capacity anticipated by our hypothesis (best viewed in color).