17. Multinomial regression*

import pymc as pm
import numpy as np
import arviz as az
import pytensor.tensor as pt

rng = np.random.default_rng(6420)

17. Multinomial regression*#

Adapted from Unit 7: NHANESmulti.odc

Problem statement#

The National Health and Nutrition Examination Survey (NHANES) is a program of studies designed to assess the health and nutritional status of adults and children in the United States. The survey is unique in that it combines interviews and physical examinations.

Assume that \(N\) subjects select a choice from \(K\) categories. The \(i\text{-th}\) subject is characterized by 3 covariates x[i, 1], x[i, 2], and x[i, 3]. Given the covariates, model the probability of a subject selecting the category \(k\), where \(k \in \{1,...,K\}\).

# data
# fmt: off
y = np.array([[1, 0, 0, 0, 0],
              [0, 1, 0, 0, 0],
              [1, 0, 0, 0, 0],
              [0, 0, 1, 0, 0],
              [0, 1, 0, 0, 0],
              [0, 0, 1, 0, 0],
              [0, 0, 0, 1, 0],
              [0, 0, 0, 0, 1],
              [0, 0, 0, 0, 1],
              [0, 0, 0, 1, 0]])

X = np.array([[2, 4, 9],
              [1, 5, 10],
              [1, 6, 14],
              [2, 4, 21],
              [2, 4, 22],
              [2, 6, 30],
              [3, 3, 33],
              [3, 2, 36],
              [3, 1, 40],
              [4, 1, 44]])
# fmt: on

# N = 10, P = 4  (intercept + 3 predictors)
X_aug = np.concatenate((np.ones((X.shape[0], 1)), X), axis=1)
N, P = X_aug.shape
K = y.shape[1]

The Multinomial distribution has two parameters, \(n\) and \(p\). See the PyMC Documentation for the Multinomial distribution class

Argument

What it means

In this model

n

Total number of independent trials per replicate. A Multinomial with \(n = 1\) is a one‑hot vector. Must equal the sum (per-row) of the corresponding observed values.

Each survey respondent makes exactly one choice, so n = 1.

p

Vector (or matrix) of category probabilities. Must sum to 1 along the last axis.

We generate p with a softmax of the linear predictor \(\eta = X\beta\). Its shape is \((N, K)\), one probability vector per observation.

observed

Either a one‑hot matrix (\(n = 1\)) or a count matrix (if \(n > 1\)). Shape must match p.

For this example every row of y is one‑hot encoded.

 
with pm.Model() as m:
    y_data = pm.Data("y", y)
    X_data = pm.Data("X", X_aug)

    _beta = pm.Normal("_beta", mu=0, tau=0.1, shape=(P, K - 1))
    beta = pt.concatenate([pt.zeros((P, 1)), _beta], axis=1)

    eta = pm.math.dot(X_data, beta)
    p = pm.Deterministic(
        "p", pm.math.softmax(eta, axis=1)
    )  # saving for probability prediction as well

    pm.Multinomial("likelihood", n=1, p=p, observed=y_data, shape=X_data.shape)

    trace = pm.sample(10000, random_seed=rng)

NUTS[nutpie]: [_beta]











X_new = np.array([1, 3, 3, 30]).reshape((1, 4))

with m:
    pm.set_data({"X": X_new})
    ppc = pm.sample_posterior_predictive(
        trace, var_names=["p", "likelihood"], predictions=True
    )









Sampling: [likelihood]











Notice the more useful credible intervals on the probabilities (p) than with the predictions for y/likelihood.






az.summary(ppc, group="predictions")










  
    

      
      mean
      sd
      eti89_lb
      eti89_ub
      ess_bulk
      ess_tail
      r_hat
      mcse_mean
      mcse_sd
    

  
  
    

      p[0, 0]
      0.0079
      0.042
      1.4e-10
      0.026
      11486
      13223
      1.00
      0.00023
      0.0013
    

    

      p[0, 1]
      0.039
      0.108
      5e-07
      0.22
      35026
      30613
      1.00
      0.00055
      0.0013
    

    

      p[0, 2]
      0.118
      0.19
      5.5e-05
      0.55
      36996
      32827
      1.00
      0.00097
      0.0012
    

    

      p[0, 3]
      0.682
      0.288
      0.12
      0.99
      35569
      33227
      1.00
      0.0015
      0.0009
    

    

      p[0, 4]
      0.154
      0.21
      0.00093
      0.63
      32743
      30727
      1.00
      0.0012
      0.0012
    

    

      likelihood[0, 0]
      0.0081
      0.09
      0
      0
      38805
      38805
      1.00
      0.00046
      0.0025
    

    

      likelihood[0, 1]
      0.038
      0.192
      0
      0
      39811
      39811
      1.00
      0.00096
      0.0023
    

    

      likelihood[0, 2]
      0.118
      0.323
      0
      1
      39191
      39191
      1.00
      0.0016
      0.0019
    

    

      likelihood[0, 3]
      0.682
      0.466
      0
      1
      38573
      38573
      1.00
      0.0024
      0.00093
    

    

      likelihood[0, 4]
      0.153
      0.36
      0
      1
      38030
      38030
      1.00
      0.0018
      0.0018
    

  









%load_ext watermark
%watermark -n -u -v -iv









Last updated: Sun, 18 Oct 2026

Python implementation: CPython
Python version       : 3.14.7
IPython version      : 9.16.1

arviz   : 1.3.0
numpy   : 2.4.6
pymc    : 6.3.1
pytensor: 3.3.0