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.

On the Properties of the Softmax Function with Application in Game Theory and Reinforcement Learning.

Exploring the Frontiers of Softmax: Provable Optimization, Applications in Diffusion Model, and Beyond

Convergence Rate of Primal-Dual Approach to Constrained Reinforcement Learning with Softmax Policy

On Passivity and Reinforcement Learning in Finite Games.

On Passivity, Reinforcement Learning and Higher-Order Learning in Multi-Agent Finite Games

Convergence and Price of Anarchy Guarantees of the Softmax Policy Gradient in Markov Potential Games

Convex Bounds on the Softmax Function with Applications to Robustness Verification

Maxent in Experimental 2×2 Population Games

Elementary Analysis of Policy Gradient Methods

On convergence rates of game theoretic reinforcement learning algorithms

Simultaneous Transport Evolution for Minimax Equilibria on Measures

Maxent in Experimental 2$\times$2 Population Games

Entropic mean-field min-max problems via Best Response flow

Attention Scheme Inspired Softmax Regression

The Minimax Property in Infinite Two-Person Win-Lose Games

On The Statistical Representation Properties Of The Perturb-Softmax And The Perturb-Argmax Probability Distributions

Connected Superlevel Set in (Deep) Reinforcement Learning and its Application to Minimax Theorems

Master equations for finite state mean field games with nonlinear activations

LEAD: Min-Max Optimization from a Physical Perspective

Softmax Deep Double Deterministic Policy Gradients

Towards Sharper Risk Bounds for Minimax Problems