Learning Well-Separated Neural Representations with the Bhattacharyya Regularization

Ashish Jha, Senbao Zhang, Nikolay Kalmykov, Anastasia Sozykina, Razan Dibo, ANH-HUY PHAN

Keywords: Regularization, Deep Neural Networks, Bhattacharyya distance, Loss function, Optimization

Abstract: Deep neural networks demonstrate remarkable classification performance, yet traditional cross-entropy training provides limited control over the geometric structure of learned embeddings. We present a moment-based geometric regularization framework that shapes class-conditional feature distributions through distributional separation measures. Our approach estimates empirical class statistics and optimizes a differentiable objective derived from the Bhattacharyya distance to encourage inter-class separation while controlling intra-class scatter. We develop three computational variants tied covariance, diagonal approximation, and projected subspace methods that provide explicit complexity-accuracy trade-offs from $\mathcal{O}(D^3)$ to $\mathcal{O}(D)$ per training step. Comprehensive empirical validation across image and text classification benchmarks demonstrates consistent improvements in embedding quality metrics while maintaining competitive classification accuracy. Statistical analysis reveals that while learned embeddings deviate from perfect Gaussianity, the moment-based approximation provides an effective geometric regularization mechanism. The framework's modular design enables straightforward integration into existing architectures with minimal computational overhead.

Primary Area: unsupervised, self-supervised, semi-supervised, and supervised representation learning

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.

Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2026/AuthorGuide.

Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.

No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.

Submission Number: 17595