catnat-jax

JAX implementation of the CatNat parameterization for categorical random variables, introduced in “Beyond Softmax: A Natural Parameterization for Categorical Random Variables” (ICML 2026).

PyTorch implementation available at allemanenti/catnat-torch.

What is CatNat?

CatNat is a drop-in alternative to softmax for parameterizing categorical distributions over K outcomes. It is designed to improve optimization by using a natural parameterization whose Fisher Information Matrix (FIM) is diagonal.

While softmax maps K unconstrained logits to K probabilities by exponentiating and normalizing them, CatNat uses only K - 1 unconstrained scores. These scores are assigned to the internal nodes of a binary tree whose leaves correspond to the K categorical outcomes.

Each score $s_i \in \mathbb{R}$ is transformed by an activation function $a(\cdot)$ into the Bernoulli probability of taking one branch at the corresponding internal node. The probability of a leaf is then the product of the branch probabilities along the path from the root to that leaf. See the figure below for a visualisation of the binary tree structure.

Figure 1. Each score s_i is placed at an internal node. The activation a(s_i) gives the probability of taking the left branch, while 1 - a(s_i) gives the probability of taking the right branch. A leaf probability is the product of the branch probabilities along its path from the root.

This package provides two entry points:

• catnat(scores, activation="sigmoid", use_logprobs=True) — returns categorical log-probabilities by default (numerically stable, gradient-friendly), or plain probabilities when use_logprobs=False.
• CatNatCategorical — a distrax-compatible distribution wrapping the above. Since the advantage of CatNat is in the diagonal FIM, it is recommended to use CatNatCategorical as it does not re-normalise the probabilities.

Installation

From git:

pip install git+https://github.com/allemanenti/catnat-jax.git

Usage

1) catnat -- given K-1 scores, get categorical log-probabilities or probabilities

import jax.numpy as jnp
from catnat_jax import catnat

# K - 1 = 4 internal scores -> K = 5 categorical log-probabilities
scores = jnp.array([0.1, -0.2, 0.3, 0.4])
logprobs = catnat(scores)                       # default: sigmoid, log-space
logprobs_nat = catnat(scores, "natural")        # alternative built-in activation

# to get probabilities:
probs = jnp.exp(catnat(scores)) # Stable way, uses log-space in the binary tree
# or
probs = catnat(scores, use_logprobs=False)  # May be less stable for extreme scores, 
                                            # as it uses probabilities in the tree

You can also pass a custom callable instead of a string. In log mode the callable must return (log_p_left, log_p_right):

import jax

def log_sigmoid_split(x):
    return jax.nn.log_sigmoid(x), jax.nn.log_sigmoid(-x)

logprobs = catnat(scores, log_sigmoid_split)

Probability-space output

Set use_logprobs=False to get normalised probabilities. The callable signature becomes x -> p (Bernoulli probability):

probs = catnat(scores, use_logprobs=False)             # shape (5,), sums to 1.0
probs = catnat(scores, "natural", use_logprobs=False)
probs = catnat(scores, jax.nn.sigmoid, use_logprobs=False)

2) CatNatCategorical -- distrax-compatible distribution

CatNatCategorical does not re-normalise its inputs, so feeding it logits directly preserves the diagonal FIM that motivates CatNat:

import jax
from catnat_jax import catnat, CatNatCategorical

scores = jnp.array([0.1, -0.2, 0.3, 0.4])
dist = CatNatCategorical(logits=catnat(scores))        # default log-prob path

# Get samples
key = jax.random.PRNGKey(0)
samples = dist.sample(seed=key, sample_shape=(8,))     # sample 8 outcomes

# Get log-probabilities of samples
log_probs = dist.log_prob(samples)                     # log-prob of each sample

# Get entropy of the distribution
entropy = dist.entropy()

Notes

• CatNatCategorical imports private modules from distrax (distrax._src.*). The dependency is pinned to distrax>=0.1.5; future distrax releases may move these paths.
• This package depends on tfp-nightly (the JAX-compatible build of TFP) rather than the stable tensorflow-probability wheel, which lags JAX compatibility. If you have tensorflow-probability already installed in your environment, uninstall it before installing catnat-jax to avoid file-overlap conflicts.

Citation

If you use catnat-jax in your research, please cite the original paper:

@inproceedings{manenti2026beyond,
  title={Beyond Softmax: A Natural Parameterization for Categorical Random Variables},
  author={Manenti, Alessandro and Alippi, Cesare},
  booktitle={International Conference on Machine Learning},
  year={2026},
  organization={PMLR}
}

License

MIT License.