Complex-Valued Unitary Representations as Classification Heads for Improved Uncertainty Quantification in Deep Neural Networks

TL;DR

This work tackles the persistent miscalibration of deep neural networks by introducing complex-valued unitary classification heads that map backbone features into a complex Hilbert space and evolve them under a learned Cayley unitary. The most effective variant, a magnitude-based readout, achieves a 2.4x improvement in Expected Calibration Error on CIFAR-10 compared to a standard softmax head, while the Born-rule readout, though providing better alignment with human uncertainty (CIFAR-10H), degrades calibration due to an information bottleneck. The authors provide a theoretical bound showing that norm-preserving unitary dynamics cap logit magnitudes, thereby reducing overconfidence, and they report targeted negative results in OOD detection and sentiment analysis to delineate the approach’s scope. A hybrid backbone–head experimental design isolates head-related effects, and the NoBorn magnitude head offers a practical drop-in calibration improvement for pretrained models. The work also highlights the trade-off between calibration and human-alignment and discusses practical implications for safety-critical AI, with code released for reproducibility.

Abstract

Modern deep neural networks achieve high predictive accuracy but remain poorly calibrated: their confidence scores do not reliably reflect the true probability of correctness. We propose a quantum-inspired classification head architecture that projects backbone features into a complex-valued Hilbert space and evolves them under a learned unitary transformation parameterised via the Cayley map. Through a controlled hybrid experimental design - training a single shared backbone and comparing lightweight interchangeable heads - we isolate the effect of complex-valued unitary representations on calibration. Our ablation study on CIFAR-10 reveals that the unitary magnitude head (complex features evolved under a Cayley unitary, read out via magnitude and softmax) achieves an Expected Calibration Error (ECE) of 0.0146, representing a 2.4x improvement over a standard softmax head (0.0355) and a 3.5x improvement over temperature scaling (0.0510). Surprisingly, replacing the softmax readout with a Born rule measurement layer - the quantum-mechanically motivated approach - degrades calibration to an ECE of 0.0819. On the CIFAR-10H human-uncertainty benchmark, the wave function head achieves the lowest KL-divergence (0.336) to human soft labels among all compared methods, indicating that complex-valued representations better capture the structure of human perceptual ambiguity. We provide theoretical analysis connecting norm-preserving unitary dynamics to calibration through feature-space geometry, report negative results on out-of-distribution detection and sentiment analysis to delineate the method's scope, and discuss practical implications for safety-critical applications. Code is publicly available.

Figures (12)

• Figure 1: Architecture of the proposed complex-valued unitary classification head. Real-valued backbone features are projected into a complex Hilbert space, normalised to unit norm, evolved under a learned Cayley unitary, and read out via either magnitude--softmax (proposed best variant) or Born rule measurement.
• Figure 2: Experiment 1: Predicted samples overlaid on ground-truth data. Wave MLP captures both branches through superposition (Qual=0.214), outperforming Real MLP (0.043) and matching Deep Ensemble diversity, though MDN (0.431) achieves the sharpest concentration near true modes.
• Figure 3: Experiment 1: Mode coverage (left) and quality (right). All models achieve full coverage; Wave MLP's quality is $5\times$ that of Real MLP, demonstrating the superposition advantage for multi-modality, though MDN's parametric mixture dominates.
• Figure 4: Experiment 1: Predicted density profiles at four $x$ values. Wave MLP (left, red) produces broad multi-peaked distributions via Born rule. Real MLP and Deep Ensemble produce sharp unimodal peaks at each mode. MDN captures both modes with tight Gaussian components.
• Figure 5: Experiment 1: (a) Training loss convergence. Wave MLP plateaus at higher loss, indicating optimisation difficulty in complex-valued density estimation. (b) Phase structure of the evolved wave function at four $x$ values, showing non-trivial phase diversity across Hilbert space dimensions.

Theorems & Definitions (8)

• Definition 1: Complex Hilbert Space
• Proposition 1: Unitarity of Cayley Map
• proof
• Theorem 1: Logit Magnitude Bound under Unitary Evolution
• proof
• Remark 1
• Proposition 2: Information Bottleneck of Born Rule Readout
• proof