Fourier Diagonalization of Softmax Dynamics: Convergence to Neural Collapse

In this work, we consider cross-entropy (CE) dynamics in a two-layer linear network with orthogonal inputs. This coincides with the unconstrained feature model used to study neural collapse, a phenomenon occurring in the representations of deep classification networks. Recent work by Garrod et al. (2025) has extended the spectral initialization framework of Saxe et al. (2013) used to explore feature learning to this setting leveraging the fact that Sylvester-Hadamard matrices diagonalize the softmax operator. However, this construction only works when the number of classes is a power of 2. Here we show that Fourier matrices can in fact diagonalize any nonlinearity: any matrix diagonalizable in the Fourier basis remains so after the application of an elementwise nonlinear activation function, or columnwise softmax. We use this to characterize the spectral dynamics of linear CE loss, proving an implicit bias towards neural collapse in the case of real spectra.