Exploring the Frontiers of Softmax: Provable Optimization, Applications in Diffusion Model, and Beyond
Yang Cao, Yingyu Liang, Zhenmei Shi, Zhao Song
TL;DR
This work provides the first NTK-based theoretical analysis of two-layer softmax networks, showing that the softmax normalization induces favorable perturbation properties in the NTK, yielding a large convex region in the loss landscape and enabling learning in the over-parameterized regime. The authors prove that, with poly$(nd)$ neurons, such networks can fit any $n imes d$ regression data and offer explicit over-parameterization and time-horizon bounds that align with ReLU and exponential activations. They extend the framework to diffusion-model score estimation, establishing provable accuracy for learning score functions with noisy labels via a kernel-regression lens and RKHS concepts, thereby demonstrating broad applicability beyond NLP. The results illuminate why softmax-based architectures, including self-attention, can exhibit strong optimization and generalization properties and suggest principled guidance for diffusion-based generative modeling implementations. Overall, the work deepens understanding of softmax effectiveness and provides a pathway to rigorous guarantees in NLP-relevant regimes and diffusion-model learning.
Abstract
The softmax activation function plays a crucial role in the success of large language models (LLMs), particularly in the self-attention mechanism of the widely adopted Transformer architecture. However, the underlying learning dynamics that contribute to the effectiveness of softmax remain largely unexplored. As a step towards better understanding, this paper provides a theoretical study of the optimization and generalization properties of two-layer softmax neural networks, providing theoretical insights into their superior performance as other activation functions, such as ReLU and exponential. Leveraging the Neural Tangent Kernel (NTK) framework, our analysis reveals that the normalization effect of the softmax function leads to a good perturbation property of the induced NTK matrix, resulting in a good convex region of the loss landscape. Consequently, softmax neural networks can learn the target function in the over-parametrization regime. To demonstrate the broad applicability of our theoretical findings, we apply them to the task of learning score estimation functions in diffusion models, a promising approach for generative modeling. Our analysis shows that gradient-based algorithms can learn the score function with a provable accuracy. Our work provides a deeper understanding of the effectiveness of softmax neural networks and their potential in various domains, paving the way for further advancements in natural language processing and beyond.