ZClassifier: Temperature Tuning and Manifold Approximation via KL Divergence on Logit Space

TL;DR

ZClassifier introduces a probabilistic approach to classification by modeling logits as Gaussian distributions, improving robustness and calibration through KL divergence minimization, and unifying uncertainty and latent control.

Contribution

It presents a novel probabilistic classification framework that replaces deterministic logits with Gaussian distributions, addressing temperature scaling and manifold approximation simultaneously.

Findings

Improves robustness over softmax classifiers.

Enhances calibration and latent separation.

Shows consistent benefits across datasets.

Abstract

We introduce a novel classification framework, ZClassifier, that replaces conventional deterministic logits with diagonal Gaussian-distributed logits. Our method simultaneously addresses temperature scaling and manifold approximation by minimizing the KL divergence between the predicted Gaussian distributions and a unit isotropic Gaussian. This unifies uncertainty calibration and latent control in a principled probabilistic manner, enabling a natural interpretation of class confidence and geometric consistency. Experiments on CIFAR-10 and CIFAR-100 demonstrate that ZClassifier improves over softmax classifiers in robustness, calibration, and latent separation, with consistent benefits across small-scale and large-scale classification settings.