Beyond Softmax: A Natural Parameterization for Categorical Random Variables

Latent categorical variables are frequently found in deep learning architectures. They can model actions in discrete reinforcement-learning environments, represent categories in latent-variable models, or express relations in graph neural networks. Despite their widespread use, their discrete nature poses significant challenges to gradient-descent learning algorithms. While a substantial body of work has offered improved gradient estimation techniques, we take a complementary approach. Specifically, we: 1) revisit the ubiquitous softmax function and demonstrate its limitations from an information-geometric perspective; 2) replace the softmax with the catnat function, a function composed by a sequence of hierarchical binary splits; we prove that this choice offers significant advantages to gradient descent due to the resulting diagonal Fisher Information Matrix. A rich set of experiments - including graph structure learning, variational autoencoders, and reinforcement learning - empirically show that the proposed function improves the learning efficiency and yields models characterized by consistently higher test performance. Catnat is simple to implement and seamlessly integrates into existing codebases. Moreover, it remains compatible with standard training stabilization techniques and, as such, offers a better alternative to the softmax function.

Many AI systems need to make choices between categories, such as choosing an action, a group, or a relation between objects. These choices are usually parametrized with the softmax function, but we show that softmax can make learning harder than necessary. We explain this using Information Geometry, which studies how learning behaves on probability distributions. We introduce catnat, a simple alternative that can directly replace softmax in existing models. This small change gives the learning algorithm a better-shaped optimization problem. Across reinforcement learning, variational autoencoders, and graph learning tasks, replacing softmax with catnat leads to faster learning and better test performance.