Diagonalizing the Softmax: Hadamard Initialization for Tractable Cross-Entropy Dynamics

TL;DR

<3-5 sentence high-level summary> The paper tackles the challenge of understanding implicit bias and optimization dynamics in multi-class cross-entropy training within non-convex, overparameterized models. It introduces a Hadamard Initialization that diagonalizes the softmax, enabling a tractable reduction of CE dynamics to a low-dimensional evolution of logit singular values and proving convergence to neural collapse via a KL-based Lyapunov function. The authors show that CE dynamics are fundamentally coupled and can exhibit non-monotonic convergence and exponential slowdowns, in contrast to mean-squared-error dynamics, but standard initializations bias trajectories toward NC by driving the system toward the simplex ETF subspace. Empirically, the results are corroborated, and the framework promises extensions to deeper, regularized, and imbalanced settings, offering a path to analyze CE dynamics beyond simple linear or MSE regimes.

Abstract

Cross-entropy (CE) training loss dominates deep learning practice, yet existing theory often relies on simplifications, either replacing it with squared loss or restricting to convex models, that miss essential behavior. CE and squared loss generate fundamentally different dynamics, and convex linear models cannot capture the complexities of non-convex optimization. We provide an in-depth characterization of multi-class CE optimization dynamics beyond the convex regime by analyzing a canonical two-layer linear neural network with standard-basis vectors as inputs: the simplest non-convex extension for which the implicit bias remained unknown. This model coincides with the unconstrained features model used to study neural collapse, making our work the first to prove that gradient flow on CE converges to the neural collapse geometry. We construct an explicit Lyapunov function that establishes global convergence, despite the presence of spurious critical points in the non-convex landscape. A key insight underlying our analysis is an inconspicuous finding: Hadamard Initialization diagonalizes the softmax operator, freezing the singular vectors of the weight matrices and reducing the dynamics entirely to their singular values. This technique opens a pathway for analyzing CE training dynamics well beyond our specific setting considered here.

Figures (10)

• Figure 1: Training of the logit singular values under the evolution given in Eq. \ref{['eq:logit_sing_evol']} with $K=8$, from the initialization $a(0)=[1,1,1,1,0.001,0.001,0.001]$. Left: The evolution of the singular values $\hat{a}_i(t)$. Middle: The evolution of the KL divergence $D_{\mathrm{KL}}(\frac{1}{K-1} 1_{K-1} \,\|\, \hat{a})$. Right: The evolution of the metric $M$, given in Eq. \ref{['eq:metric']}.
• Figure 2: Comparison of Hadamard Init., Random Init., and the theoretical ODE dynamics. Plots show metrics vs. training iteration on a log scale. Shaded regions denote standard error over 10 realizations. Sec. \ref{['sec:experiments']} for detailed discussion.
• Figure 3: Plots of the normalized logit matrices generated by the vector $a^* = [1,1,1,1,0,0,0]$, which is uniform on its support, and by the non-uniform vector $a_{\textrm{conv}}=[1,1,1,2,0,0,0]$. The smallest normalized margin achieved by $a_{\textrm{conv}}$ is visibly larger, indicating that this configuration outperforms the uniform one as the logit norm diverges.
• Figure 4: Comparison of Hadamard Init., Random Init., and theoretical ODE for $K=8$ classes.
• Figure 5: Comparison of Hadamard Init., Random Init., and theoretical ODE for $K=16$ classes.

Theorems & Definitions (12)

• Definition 1: Neural Collapse
• Definition 2: Sylvester Hadamard Matrix
• Theorem 1
• Theorem 2
• Theorem 3
• Theorem 4
• Theorem 5
• Theorem 6
• Theorem 7
• Theorem 8