[Paper Review] The Impact of Reinitialization on Generalization in Convolutional Neural Networks

Ibrahim Alabdulmohsin, Hartmut Maennel+1 more|arXiv (Cornell University)|Sep 1, 2021

Advanced Neural Network Applications46 references4 citations

TL;DR

This paper introduces a novel layerwise reinitialization method (lw) for convolutional neural networks that improves generalization, especially on small datasets, by successively reinitializing lower layers while maintaining weight norms and enhancing training margins. The method outperforms prior reinitialization techniques by promoting flatter loss landscapes and encouraging learning of general features in early layers.

ABSTRACT

Recent results suggest that reinitializing a subset of the parameters of a neural network during training can improve generalization, particularly for small training sets. We study the impact of different reinitialization methods in several convolutional architectures across 12 benchmark image classification datasets, analyzing their potential gains and highlighting limitations. We also introduce a new layerwise reinitialization algorithm that outperforms previous methods and suggest explanations of the observed improved generalization. First, we show that layerwise reinitialization increases the margin on the training examples without increasing the norm of the weights, hence leading to an improvement in margin-based generalization bounds for neural networks. Second, we demonstrate that it settles in flatter local minima of the loss surface. Third, it encourages learning general rules and discourages memorization by placing emphasis on the lower layers of the neural network. Our takeaway message is that the accuracy of convolutional neural networks can be improved for small datasets using bottom-up layerwise reinitialization, where the number of reinitialized layers may vary depending on the available compute budget.

Motivation & Objective

• To investigate whether reinitialization of neural network parameters during training can improve generalization, particularly in low-data regimes.
• To compare the effectiveness of various reinitialization strategies—random, magnitude-based, fixed, and fully-connected—on standard image classification benchmarks.
• To develop and evaluate a new reinitialization method, layerwise (lw), that systematically reinitializes layers from the bottom up with normalization.
• To explain the generalization gains through margin maximization, loss landscape flatness, and reduced memorization.
• To determine the conditions under which reinitialization provides benefits and to guide practical deployment based on compute constraints.

Proposed method

• The proposed lw method partitions a CNN into K convolutional blocks and reinitializes the lowest k blocks in sequence, starting from k=1 up to K.
• At each round k, the weights of the first k blocks are rescaled to their original initialization norms, while all subsequent layers are reinitialized.
• A normalization layer (batch normalization-like) is inserted or updated after block k to stabilize activations and maintain gradient flow.
• The process is repeated N times per block, with fine-tuning after each reinitialization round to allow convergence.
• The method explicitly maintains the Frobenius norm of early-layer weights while increasing margins on training examples.
• The approach is evaluated across 12 image classification datasets using multiple architectures, with ablation studies on normalization and rescaling components.

Experimental results

Research questions

• RQ1 Does layerwise reinitialization improve generalization in CNNs compared to standard training and other reinitialization methods?
• RQ2 How does reinitialization affect the margin on training examples and the norm of the weights?
• RQ3 To what extent does reinitialization lead to flatter local minima in the loss landscape?
• RQ4 Does reinitialization reduce memorization and encourage learning of general features in early layers?
• RQ5 How does the performance of lw vary with dataset size and available compute budget?

Key findings

• The layerwise reinitialization (lw) method consistently outperforms baseline training and other reinitialization strategies (welsr, dsd, wels, fc) across 12 benchmark datasets, especially on small training sets.
• lw increases the minimum margin on training examples without increasing the Frobenius norm of the weights, leading to improved margin-based generalization bounds.
• The final model trained with lw resides in a flatter region of the loss landscape, as evidenced by smaller changes in training loss under parameter perturbations compared to standard training.
• Ablation studies confirm that normalization and rescaling are critical components of lw, as their removal degrades performance.
• lw reduces memorization by emphasizing learning in early layers, where general features are typically learned, while discouraging overfitting to dataset-specific patterns.
• The benefits of reinitialization diminish on large datasets, suggesting its primary advantage lies in low-data generalization.