More Is More -- Narrowing the Generalization Gap by Adding Classification Heads

TL;DR

This paper introduces TransNet, an architecture enhancement that learns useful invariance levels from data, improving generalization and reducing overfitting in image classification by adding classification heads and enforcing soft invariance.

Contribution

The paper proposes TransNet, a novel architecture with multiple classification heads and a training algorithm that learns invariance levels, enhancing generalization without increasing inference complexity.

Findings

Improved accuracy on various datasets.

Enhanced generalization performance.

Effective pruning for deployment.

Abstract

Overfit is a fundamental problem in machine learning in general, and in deep learning in particular. In order to reduce overfit and improve generalization in the classification of images, some employ invariance to a group of transformations, such as rotations and reflections. However, since not all objects exhibit necessarily the same invariance, it seems desirable to allow the network to learn the useful level of invariance from the data. To this end, motivated by self-supervision, we introduce an architecture enhancement for existing neural network models based on input transformations, termed 'TransNet', together with a training algorithm suitable for it. Our model can be employed during training time only and then pruned for prediction, resulting in an equivalent architecture to the base model. Thus pruned, we show that our model improves performance on various data-sets while exhibiting improved generalization, which is achieved in turn by enforcing soft invariance on the convolutional kernels of the last layer in the base model. Theoretical analysis is provided to support the proposed method.