Diagonalizing the Softmax: Hadamard Initialization for Tractable Cross-Entropy Dynamics
Abstract
Cross-entropy (CE) loss is central to deep learning, but existing theory often relies on simplifications—such as squared loss or convex models—that miss key aspects of CE optimization. In this work, we study multi-class CE dynamics using a two-layer linear network with orthogonal inputs, the simplest non-convex setting where the CE implicit bias remains unresolved. This coincides with the unconstrained features model used to study neural collapse (NC). Our analysis is based on a key observation: Hadamard initialization diagonalizes the softmax operator. This allows us to extend the spectral initialization framework that Saxe et al. (2013, 2019) developed for squared loss. We prove convergence to NC under spectral CE training and give the first finite-time analysis in this setting via an explicit Lyapunov function that decreases monotonically to NC despite spurious critical points. We further identify CE-specific phenomena absent under squared loss—coupling, non-monotonic convergence, qualitative dependence on the number of classes, and exponential slowdowns. We also characterize the role of network width and show empirically that spectral dynamics qualitatively model small random initialization.