Collaborative Learning: Joint Training of Multiple Classifier Heads Improves Generalization of Deep Neural Networks

This article introduces the TensorFlow implementation of the NeurIPS 2018 paper Collaborative Learning for Deep Neural Networks. The core idea is to improve model generalization performance without increasing inference cost by jointly training multiple classifier heads to share intermediate layer representations. This implementation comes from a GitHub repository (by authors like qq456cvb) and has been validated on CIFAR-10 using the DenseNet architecture.

Deep neural networks often face the problem of performing well on training data but struggling to generalize to new data. Traditional ensemble learning improves performance but multiplies inference costs; knowledge distillation requires pre-training a teacher model, which is a complex process. The 2018 NeurIPS paper proposed collaborative learning to address these pain points.

The core of collaborative learning is the consensus mechanism of multiple classifier heads: the same network is split into multiple heads that share underlying representations. During training, each head learns from both real hard labels and soft labels from other heads (consensus predictions after temperature scaling). Technical details include: hierarchical head splitting based on DenseNet-40-12 (recursive bisection to form 4 heads); backpropagation gradient rescaling (dividing gradients at split points by the number of child nodes to stabilize training); and a loss function that is a weighted combination of hard loss (β=0.5) and soft loss (KL divergence).

Training is based on CIFAR-10: optimizer SGD (momentum 0.9, weight decay 1e-4), 300 training epochs (learning rate decayed in a stepwise manner), data augmentation (random cropping, flipping). The result is a test error of 6.09%, which is close to the value reported in the paper. The slight gap may be due to the position of the split point (in the implementation, splitting occurs after BatchNorm+ReLU, which was not explicitly stated in the paper).

The prominent advantages of collaborative learning are: only one head needs to be retained during testing, resulting in zero additional inference cost; robustness to label noise, which can mitigate the impact of abnormal samples; and friendly for engineering deployment (implemented with TensorFlow 1.x + Tensorpack, code concentrated in main.py, CIFAR-10 data downloaded automatically).

Collaborative learning is an elegant multi-head collaborative training paradigm, which proves that one can benefit from multi-perspective learning without sacrificing inference efficiency. It is recommended for practitioners to try, especially in scenarios with limited data or low-quality annotations, where its consensus mechanism may bring significant generalization improvements.