📊 statistics

Beyond Softmax: A Natural Parameterization for Categorical Random Variables

This paper proposes replacing the standard softmax function with a novel "catnat" parameterization based on hierarchical binary splits to overcome gradient descent limitations in categorical latent variables, demonstrating through information geometry and extensive experiments that this approach yields a diagonal Fisher Information Matrix, improved learning efficiency, and superior test performance across diverse deep learning tasks.

Original authors: Alessandro Manenti, Cesare Alippi

Published 2026-05-14

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Alessandro Manenti, Cesare Alippi

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are teaching a robot to make choices. Sometimes, the robot has to pick one option out of many, like choosing a word in a sentence, picking a move in a video game, or deciding which connections exist in a social network. In the world of machine learning, these are called categorical variables.

For a long time, the standard way to teach the robot how to make these choices has been a mathematical tool called Softmax. Think of Softmax as a very popular, well-worn map that guides the robot's learning process. It works, but the authors of this paper argue that this map is a bit "bumpy" and confusing. It creates a landscape where the robot's learning path (gradient descent) tends to zigzag and oscillate, making it harder to find the best solution quickly and accurately.

The Problem: A Bumpy Road

The authors use a concept from mathematics called Information Geometry to explain why Softmax is tricky. Imagine the robot's learning space as a hilly terrain.

With Softmax: The hills and valleys are curved in a complex way. When the robot tries to roll down the hill to find the lowest point (the best solution), the curvature forces it to bounce back and forth. It's like trying to roll a ball down a twisted, curved slide; it takes a long, winding path to get to the bottom.
The Result: This "bounciness" makes training unstable and can lead to the robot settling for a "good enough" answer rather than the best answer.

The Solution: The "Catnat" Map

The authors propose a new way to map these choices, which they call Catnat.

Instead of asking the robot to pick one option out of ten all at once (which is what Softmax does), Catnat breaks the decision down into a series of simple Yes/No questions, arranged like a family tree or a flowchart.

The Analogy: Imagine you are trying to find a specific book in a library.
- Softmax approach: You look at every single book on the shelf at once and try to guess which one is the right one. It's overwhelming and messy.
- Catnat approach: You ask a series of simple questions: "Is it in the top half?" (Yes/No). "Is it in the left quarter?" (Yes/No). "Is it in the first row?" (Yes/No). You keep splitting the options in half until you find the book.

This "hierarchical binary split" approach changes the shape of the learning landscape. The authors prove mathematically that this new map creates a straight, flat road (a diagonal structure) instead of a curved, bumpy one.

Why This Matters

Because the road is straighter and flatter:

The robot learns faster: It doesn't waste time zigzagging.
The robot learns better: It finds more accurate solutions.
It's easy to swap: You can replace the old Softmax tool with Catnat in existing code without rewriting the whole program.

What They Tested

The authors didn't just do math; they tested this new tool in three very different real-world scenarios to see if it actually worked:

Learning Graph Structures: Trying to figure out how things are connected (like a social network or a brain map). Catnat helped the model figure out the connections more accurately.
Variational Autoencoders (VAEs): These are AI models that learn to compress and recreate images (like turning a photo of a cat into a code and back). Using Catnat, the models recreated images with higher quality and fewer errors.
Reinforcement Learning: Training AI to play video games (specifically Breakout and Seaquest). The AI using Catnat scored higher points and played better than the AI using the standard Softmax method.

The Bottom Line

The paper claims that by changing how we ask the AI to make choices—switching from a complex "pick one of many" method to a simple "yes/no" tree structure—we can make the learning process smoother, more stable, and more effective. It's a simple switch that yields consistently better results across different types of AI tasks.

Technical Summary: Beyond Softmax: A Natural Parameterization for Categorical Random Variables

Problem Formulation

Latent categorical random variables are fundamental components in deep learning architectures, utilized for modeling discrete actions in reinforcement learning (RL), representing topics in latent-variable models, and defining graph structures in graph neural networks (GNNs). However, the discrete nature of these variables poses significant challenges for gradient-descent learning algorithms. While the training signal is often derived from downstream tasks (e.g., reconstruction loss in Variational Autoencoders or rewards in RL), the lack of direct supervisory signals for the latent variables necessitates the use of gradient estimators.

Standard approaches often rely on the Score Function (REINFORCE) estimator, which is unbiased but suffers from high variance, leading to unstable training. Alternative methods, such as the Gumbel-Softmax trick, introduce continuous relaxations to enable backpropagation but may introduce bias. A complementary perspective, explored in this work, focuses on the parameterization of the categorical distribution itself. The authors argue that the ubiquitous softmax function, while standard, induces a dense Fisher Information Matrix (FIM). This density creates geometric distortions in the parameter space, causing gradient-based optimizers to follow oscillatory and less direct trajectories, thereby hindering convergence and stability.

Methodology: The Catnat Parameterization

The paper proposes catnat, a new parameterization for categorical random variables designed to align with principles of information geometry.

1. Information Geometric Perspective

The authors analyze the geometry of the parameter space using the Fisher Information Matrix (FIM). For a categorical distribution parameterized by softmax, the FIM is dense, with off-diagonal entries $-p_i p_j$ that couple all scores. This curvature complicates optimization. In contrast, the authors propose a parameterization that yields a diagonal FIM, effectively flattening the statistical manifold and allowing optimizers to follow more direct paths to the solution.

2. Hierarchical Binary Splits

The catnat function models the categorical probability distribution as the outcome of a sequence of hierarchical binary decisions, structured as a binary tree.

Structure: Given $K$ classes, the parameterization uses a hierarchy of depth $H = \log_2(K)$ . Each unnormalized score $s_i$ corresponds to a unique node in this tree.
Mechanism: An activation function $a(\cdot)$ is applied to each score to produce a binary probability. The final probability $p_k$ for a category $k$ (identified by a binary string $\vec{b}$ ) is the product of probabilities along the path from the root to the leaf node representing $k$ .
Notation: The authors introduce a notation splitting the binary path into a hierarchy level and an identifier, allowing for a compact mathematical representation of the probabilities.

3. Theoretical Properties

The paper proves that the catnat parameterization induces a diagonal Fisher Information Matrix with respect to the intermediate scores $\vec{s}$ .

Theorem 4.2: The FIM $G_a(\vec{s})$ for catnat is diagonal, where the diagonal entries depend on the probability of reaching node $i$ ( $P(a_i)$ ) and the derivative of the activation function.
Natural Activation: To further simplify the optimization landscape, the authors propose a specific "natural activation function" $\nu(x)$ , a smoothed step function. When used with catnat, this activation renders the diagonal entries of the FIM constant (dependent only on the ancestor probabilities), effectively eliminating the dependence on the local score values. This aligns the parameterization with the objectives of natural gradient methods without requiring the explicit computation and inversion of the FIM.

Key Contributions

Revisiting Softmax: The paper provides an information-geometric critique of the softmax function, demonstrating that its dense FIM induces unnecessary geometric distortions in the optimization landscape.
Catnat Proposal: It introduces catnat, a hierarchical binary-split parameterization that guarantees a diagonal FIM with respect to the input scores.
Natural Activation: It defines a specific activation function $\nu$ that, when paired with catnat, simplifies the FIM further, removing local score dependencies from the diagonal terms.
Implementation: The method is presented as a drop-in replacement for softmax in existing codebases, compatible with standard training stabilization techniques (e.g., control variates, Gumbel-Softmax) and various optimizers.

Experimental Results

The authors evaluate catnat against softmax, sparsemax, and catnat with sigmoid activation across three distinct domains:

1. Graph Structure Learning (GSL)

Setup: Learning latent graph structures where edges are modeled as Bernoulli variables. The task involves recovering the true graph structure and GNN parameters.
Results: The natural activation catnat consistently outperformed softmax, sparsemax, and sigmoid-based catnat across all metrics (Energy Score, Point Prediction MAE/MSE, and parameter estimation MAE). It demonstrated superior accuracy in recovering the true latent distribution parameters.

2. Variational Autoencoders (VAEs)

Setup: Discrete VAEs on MNIST and Binary MNIST datasets with varying numbers of latent variables ( $N$ ) and categories ( $K$ ).
Results: Both catnat variants (sigmoid and natural activation) achieved lower test negative log-likelihoods compared to softmax and sparsemax across all configurations ( $N \in \{10, 20, 30\}$ , $K \in \{8, 16, 32\}$ ). The natural activation showed slight improvements in most settings. These results held true for both Adam and SGD optimizers, indicating the benefits are not optimizer-specific.

3. Reinforcement Learning (RL)

Setup: Policy learning in Atari environments (Breakout and Seaquest) using Proximal Policy Optimization (PPO).
Results: The catnat parameterization with natural activation yielded higher episodic returns than the standard softmax. The improvement was modest in Breakout but substantial in the more complex Seaquest environment, suggesting robustness in high-dimensional, sequential decision-making tasks.

Significance and Claims

The paper claims that replacing the standard softmax with the catnat parameterization offers a principled, theoretically grounded improvement for training models with latent categorical variables. By inducing a diagonal Fisher Information Matrix, catnat reduces geometric distortions in the parameter space, facilitating more stable and efficient gradient-based optimization.

The authors emphasize that:

The approach is complementary to existing variance-reduction techniques (e.g., control variates) and can be combined with them.
The method is simple to implement and integrates seamlessly into existing deep learning pipelines.
The improvements are robust, observed across diverse architectures, optimizers, and task complexities.

The paper concludes that while the study focuses on categorical distributions, the findings suggest that parameterizations inducing diagonal FIMs generally improve performance, opening avenues for extending this approach to other families of distributions. The authors position catnat as a better alternative to softmax for improving training stability and final test performance in models relying on discrete latent variables.