# Beyond Softmax: A Natural Parameterization for Categorical Random Variables

arXiv: https://arxiv.org/abs/2509.24728 · HTML: https://mlsanity.com/paper/2509.24728/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 35 of its 74 references that are in the map.
- First posted on arXiv 29 Sep 2025.

## Builds on (most cited first, 30 of 35)

- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Semi-Supervised Classification with Graph Convolutional Networks](https://mlsanity.com/paper/1609.02907/) (2016)
- [Neural Discrete Representation Learning](https://mlsanity.com/paper/1711.00937/) (2017)
- [Playing Atari with Deep Reinforcement Learning](https://mlsanity.com/paper/1312.5602/) (2013)
- [Categorical Reparameterization with Gumbel-Softmax](https://mlsanity.com/paper/1611.01144/) (2016)
- [Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation](https://mlsanity.com/paper/1308.3432/) (2013)
- [High-Dimensional Continuous Control Using Generalized Advantage Estimation](https://mlsanity.com/paper/1506.02438/) (2015)
- [The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables](https://mlsanity.com/paper/1611.00712/) (2016)
- [Importance Weighted Autoencoders](https://mlsanity.com/paper/1509.00519/) (2015)
- [New insights and perspectives on the natural gradient method](https://mlsanity.com/paper/1412.1193/) (2014)
- [From Softmax to Sparsemax: A Sparse Model of Attention and Multi-Label Classification](https://mlsanity.com/paper/1602.02068/) (2016)
- [Neural Relational Inference for Interacting Systems](https://mlsanity.com/paper/1802.04687/) (2018)
- [Monte Carlo Gradient Estimation in Machine Learning](https://mlsanity.com/paper/1906.10652/) (2019)
- [Learning to Explain: An Information-Theoretic Perspective on Model Interpretation](https://mlsanity.com/paper/1802.07814/) (2018)
- [Learning Discrete Structures for Graph Neural Networks](https://mlsanity.com/paper/1903.11960/) (2019)
- [Revisiting Natural Gradient for Deep Networks](https://mlsanity.com/paper/1301.3584/) (2013)
- [A Kronecker-factored approximate Fisher matrix for convolution layers](https://mlsanity.com/paper/1602.01407/) (2016)
- [Neural Variational Inference and Learning in Belief Networks](https://mlsanity.com/paper/1402.0030/) (2014)
- [SLAPS: Self-Supervision Improves Structure Learning for Graph Neural Networks](https://mlsanity.com/paper/2102.05034/) (2021)
- [Backpropagation through the Void: Optimizing control variates for black-box gradient estimation](https://mlsanity.com/paper/1711.00123/) (2017)
- [Implicit MLE: Backpropagating Through Discrete Exponential Family Distributions](https://mlsanity.com/paper/2106.01798/) (2021)
- [A Review of the Gumbel-max Trick and its Extensions for Discrete Stochasticity in Machine Learning](https://mlsanity.com/paper/2110.01515/) (2021)
- [REBAR: Low-variance, unbiased gradient estimates for discrete latent variable models](https://mlsanity.com/paper/1703.07370/) (2017)
- [Discovering Discrete Latent Topics with Neural Variational Inference](https://mlsanity.com/paper/1706.00359/) (2017)
- [Gradient Estimation with Stochastic Softmax Tricks](https://mlsanity.com/paper/2006.08063/) (2020)
- [MuProp: Unbiased Backpropagation for Stochastic Neural Networks](https://mlsanity.com/paper/1511.05176/) (2015)
- [SIMPLE: A Gradient Estimator for k-Subset Sampling](https://mlsanity.com/paper/2210.01941/) (2022)
- [Estimating Gradients for Discrete Random Variables by Sampling without Replacement](https://mlsanity.com/paper/2002.06043/) (2020)
- [Fisher Information and Natural Gradient Learning of Random Deep Networks](https://mlsanity.com/paper/1808.07172/) (2018)
- [Rao-Blackwellized Stochastic Gradients for Discrete Distributions](https://mlsanity.com/paper/1810.04777/) (2018)

## Abstract (arXiv)

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 textitsoftmax function and demonstrate its limitations from an information-geometric perspective; 2) replace the textitsoftmax with the textitcatnat function, a function composed of 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. textitCatnat 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 textitsoftmax function.