Optimizing Convolution Neural Network Embedding Space: Utilizing Statistical Loss Function Pohang University of Science and Technology

Abstract

Convolutional Neural Networks (CNNs) have achieved remarkable success in various classification tasks. However, most existing CNN-based classification studies focus solely on minimizing the discrepancy between the predicted and ground truth distributions, such as cross-entropy loss, without explicitly considering the embedding space where features are represented. In this paper, we propose a novel approach to optimize the CNN embedding space by introducing a Statistical Loss function that leverages Welch's $t$-test to maximize inter-class separability. Our method introduces an additional embedding layer between the CNN encoder and the classifier, which transforms the features into a space more suitable for t-value computation. The Statistical Loss is calculated based on the t-values between different class pairs across each dimension of the embedding space, encouraging the model to learn more discriminative feature representations. We conducted experiments on various datasets, including the ADNI brain network dataset and CIFAR-10, to demonstrate the effectiveness of our approach. The results show that our method achieves improved performance compared to baseline models, especially on class-imbalanced datasets like CIFAR-10 Long-Tail. Furthermore, we investigate the impact of the Statistical Loss weight on model performance and provide visualizations of the optimized embedding space using t-SNE, revealing enhanced class separation and intra-class compactness. Our approach is easily integrable into existing CNN architectures with minimal modifications, offering a plug-and-play solution for embedding space optimization.