Normalized Conditional Mutual Information Surrogate Loss for Deep Neural Classifiers

TL;DR

This work introduces the normalized conditional mutual information (NCMI) as a principled, information-theoretic surrogate loss for deep classifier training, defined as $\hat{I}(X; \mathcal{P}| Y) = \frac{I(X; \mathcal{P}|Y)}{\Gamma}$ and minimized via an alternating optimization that employs dummy centroids $\boldsymbol{q}^y$. By modeling classification as a three-state Markov chain and focusing on reducing $I(X; \mathcal{P}| Y}$ while maintaining class separation through $\Gamma$, the authors demonstrate that NCMI can surpass cross-entropy and other surrogates on CIFAR-100, ImageNet, and whole-slide imaging benchmarks with computational costs comparable to CE. The methodology relies on a network mapping to simplex-valued outputs through a normalization step (NSF), feature centering, and a temperature scaling, enabling stable optimization and effective centroid-based evaluation. Empirically, NCMI yields notable gains (e.g., a $2.77\%$ top-1 increase on ImageNet with ResNet-50 and an $\approx 8.6\%$ macro-$F_1$ improvement on CAMELYON-17) and demonstrates robustness across architectures and batch sizes, highlighting its practical potential as a CE alternative for diverse domains.

Abstract

In this paper, we propose a novel information theoretic surrogate loss; normalized conditional mutual information (NCMI); as a drop in alternative to the de facto cross-entropy (CE) for training deep neural network (DNN) based classifiers. We first observe that the model's NCMI is inversely proportional to its accuracy. Building on this insight, we introduce an alternating algorithm to efficiently minimize the NCMI. Across image recognition and whole-slide imaging (WSI) subtyping benchmarks, NCMI-trained models surpass state of the art losses by substantial margins at a computational cost comparable to that of CE. Notably, on ImageNet, NCMI yields a 2.77% top-1 accuracy improvement with ResNet-50 comparing to the CE; on CAMELYON-17, replacing CE with NCMI improves the macro-F1 by 8.6% over the strongest baseline. Gains are consistent across various architectures and batch sizes, suggesting that NCMI is a practical and competitive alternative to CE.

Figures (5)

• Figure 1: Mappings from the label space $Y$ to the input space $X$, and from the input space to a output space $\hat{Y}$. Input $\boldsymbol{x}$ are sampled from the class $Y=y$ according to the $P_{X|Y}(\cdot|y)$. This is further mapped by a DNN and a simplex-valued function to an output probability distribution $\boldsymbol{p}\in\mathcal{P}$.
• Figure 2: The accuracy vs NCMI value over the validation set of pre-trained ResNet models on the ImageNet dataset.
• Figure 3: ResNet-50 test accuracy on CIFAR-100 as a function of batch size. We evaluate batch sizes {16, 32, 64, 128, 256, 1024}; NCMI consistently outperforms CE and SupCon across all settings.
• Figure 4: The evolution curves of ResNet-18 on CIFAR-100 under all combinations of NSF and feature centering (enabled/disabled). Shown are the epoch-wise trajectories of (a) CMI, (b) $\Gamma$, (c) NCMI, and (d) accuracy.
• Figure 5: Ablation of feature centering and the normalized sigmoid (NSF). We ablate each component by enabling or disabling it: w.o.C/w.C denote without/with centering, and SM/NSF denote applying softmax/normalized sigmoid function. (a) Training curves of CMI, $\Gamma$, NCMI, and top-1 accuracy for ResNet-18 on CIFAR-100. (b) t-SNE of features from three randomly selected classes at epochs 60, 120, and 200; black crosses mark constant-valued vectors (all entries equal), which map via $\sigma$ to the uniform distribution. (c) t-SNE trajectories of feature centers and their EMA updates across training under all settings.