Multinomial Treatment Probabilities

Why Model Treatment as Multinomial?

A treatment may have several nominal options: for example, three drug regimens, different surgical procedures, or a choice among several care pathways. Nominal means that the categories have no inherent ranking. When treatment is recorded this way, a treatment model must estimate the conditional probability of every option given the covariates measured before assignment. Those probabilities are also the inputs needed to identify the probability of the treatment actually received by each person.

This tutorial focuses on fitting and using a baseline-category multinomial logistic model. It does not treat the categories as ordered or model a joint sequence of distinct treatment decisions. Let $A_i$ denote the treatment assigned to individual $i$, with possible categories $0,1,\ldots,K$, and let $X_i$ be a vector of covariates. The outcome of interest here is the conditional probability $P(A_i=k\mid X_i)$, for each category $k$.

Definition: A baseline-category multinomial logistic model uses category $0$ as a reference and specifies, for each $k=1,\ldots,K$, $$ \log\left\{\frac{P(A_i=k\mid X_i)}{P(A_i=0\mid X_i)}\right\} =\eta_{ik} =X_i^\mathsf{T}\beta_k. $$ Here $\beta_k$ is a category-specific coefficient vector. The corresponding fitted conditional probabilities are $$ p_{ik} =\frac{\exp(\eta_{ik})}{1+\sum_{j=1}^{K}\exp(\eta_{ij})} \quad (k=1,\ldots,K), \qquad p_{i0} =\frac{1}{1+\sum_{j=1}^{K}\exp(\eta_{ij})}. $$

The covariate vector normally includes an intercept, so category-specific intercepts are allowed. The model assumes that the treatment options are mutually exclusive and collectively exhaustive for the analysis population; each individual has exactly one observed category. It also assumes that the chosen linear predictors adequately represent the category-specific log odds. Covariates must be defined at the appropriate time: in a treatment-probability model they should be measured before the treatment assignment being modeled.

Turning Linear Predictors Into Probabilities

The probability formula follows directly from the baseline-category logits. Write $p_{ik}=P(A_i=k\mid X_i)$. The logit equations imply $p_{ik}=p_{i0}\exp(\eta_{ik})$ for $k\geq1$. Since all category probabilities sum to one, $p_{i0}\{1+\sum_{k=1}^{K}\exp(\eta_{ik})\}=1$. Solving for $p_{i0}$ gives the reference-category formula, and multiplying it by $\exp(\eta_{ik})$ gives every other category probability.

Result: For finite linear predictors, the baseline-category formula gives $p_{ik}>0$ for every category and $\sum_{k=0}^{K}p_{ik}=1$. It also recovers the modeled odds because $p_{ik}/p_{i0}=\exp(\eta_{ik})$.

To verify the result, let $D_i=1+\sum_{j=1}^{K}\exp(\eta_{ij})$, which is positive. Each numerator in the probability formula is positive, so every $p_{ik}>0$. Their sum is $$ p_{i0}+\sum_{k=1}^{K}p_{ik} =\frac{1+\sum_{k=1}^{K}\exp(\eta_{ik})}{D_i}=1. $$ Finally, $p_{ik}/p_{i0}=\exp(\eta_{ik})$, and taking logarithms recovers the specified log odds. Thus the formula defines valid probabilities with precisely the requested baseline-category logits. $\square$

Worked Example: Calculating Three Treatment Probabilities

Suppose treatment has categories $0,1,2$, category $0$ is the reference, and a single covariate $x$ gives the predictors $\eta_1=0.4x$ and $\eta_2=-0.3x$. Calculate all probabilities when $x=1$.

The three unnormalized quantities are $1$ for category $0$, $\exp(0.4)=1.4918$ for category $1$, and $\exp(-0.3)=0.7408$ for category $2$. Their denominator is $1+1.4918+0.7408=3.2326$. Therefore, to four decimal places, $$ p_{i0}=\frac{1}{3.2326}=0.3093,\qquad p_{i1}=\frac{1.4918}{3.2326}=0.4615,\qquad p_{i2}=\frac{0.7408}{3.2326}=0.2292. $$ The rounded probabilities sum to $0.3093+0.4615+0.2292=1.0000$. If this individual received category $2$, the model-based probability corresponding to the observed treatment is $0.2292$, not the probability for the reference category or the largest probability.

Likelihood and Estimation

For independent individuals, conditional on their covariates, the likelihood is the product of the probability assigned by the model to each observed category. Define $I(A_i=k)$ to be one when individual $i$ received category $k$, and zero otherwise. The likelihood and log-likelihood are $$ L(\beta_1,\ldots,\beta_K) =\prod_{i=1}^{n}\prod_{k=0}^{K}p_{ik}^{I(A_i=k)}, \qquad \ell(\beta_1,\ldots,\beta_K) =\sum_{i=1}^{n}\sum_{k=0}^{K}I(A_i=k)\log(p_{ik}). $$ In practice, coefficients are usually estimated by maximizing this log-likelihood numerically. A finite maximum need not exist: for example, covariates can perfectly or nearly perfectly predict a category. Identifiability also requires that the covariate design contain enough variation to distinguish the coefficients.

Result: For a coefficient vector $\beta_k$ in the category-$k$ predictor $X_i^\mathsf{T}\beta_k$, the log-likelihood score is $$ \frac{\partial \ell}{\partial \beta_k} =\sum_{i=1}^{n}X_i\{I(A_i=k)-p_{ik}\}. $$ This follows by differentiating the log probability contributions: the contribution for an observed category $r$ has derivative $X_i\{I(r=k)-p_{ik}\}$, because the derivative of the log denominator with respect to $\beta_k$ is $X_i p_{ik}$. Summing over individuals gives the stated score. $\square$

The score compares observed category indicators with their fitted probabilities. At an interior maximum, the score is zero. In particular, with an intercept in each category-specific predictor, the fitted expected number in category $k$ equals the observed number in that category: the intercept score equation is $\sum_i\{I(A_i=k)-p_{ik}\}=0$. With covariates, the score equations also balance the corresponding covariate-weighted totals.

Worked Example: A Saturated Model for Two Covariate Strata

Suppose a binary covariate $x$ defines two groups. In the $x=0$ group, the observed counts in categories $0,1,2$ are $6,3,1$. In the $x=1$ group, the counts are $2,5,3$. Fit a model that allows a separate probability for each category in each stratum.

For one stratum with counts $n_0,n_1,n_2$, the likelihood for probabilities $q_0,q_1,q_2$ is proportional to $\prod_k q_k^{n_k}$, subject to $\sum_kq_k=1$. If all counts are positive, the log-likelihood is $\sum_k n_k\log q_k$. A Lagrange multiplier $\lambda$ gives $n_k/q_k=\lambda$, hence $q_k=n_k/\lambda$. Summing over categories yields $\lambda=\sum_kn_k=n$, so the maximum occurs at $q_k=n_k/n$. This establishes the saturated-stratum estimates rather than assuming them.

The first stratum has $n=10$, giving $(\hat q_0,\hat q_1,\hat q_2)=(0.6,0.3,0.1)$. Its two fitted baseline logits are $\log(0.3/0.6)=\log(0.5)=-0.6931$ and $\log(0.1/0.6)=\log(1/6)=-1.7918$. The second stratum also has $n=10$, giving $(0.2,0.5,0.3)$. Its baseline logits are $\log(0.5/0.2)=\log(2.5)=0.9163$ and $\log(0.3/0.2)=\log(1.5)=0.4055$.

A linear predictor in $x$ exactly represents these two logits. For category $1$, the intercept is $-0.6931$, and the coefficient of $x$ is $0.9163-(-0.6931)=1.6094$. For category $2$, the intercept is $-1.7918$, and the coefficient is $0.4055-(-1.7918)=2.1973$. At $x=1$, the exponentiated predictors are approximately $2.5$ and $1.5$; with the reference numerator $1$, the denominator is $5$. This recovers probabilities $1/5=0.2$, $2.5/5=0.5$, and $1.5/5=0.3$, as required.

Extracting Probabilities for Observed Treatments

After fitting a model, calculate all $K+1$ probabilities for each individual using that individual's covariates and the estimated coefficients. Then select the probability whose category matches the observed treatment. Equivalently, for observed category $A_i$, the selected value is $\hat p_{i,A_i}=\prod_{k=0}^{K}\hat p_{ik}^{I(A_i=k)}$, since exactly one indicator equals one. Keeping this selection explicit prevents a common indexing error: a person's probability of receiving category $1$ is not necessarily the probability needed when that person actually received category $2$.

Worked Example: Selecting the Observed-Category Probabilities

Use the model $\eta_1=0.4x,\ \eta_2=-0.3x$ from the first example for four independent individuals. Their covariate values are $-1,0,1,1$, and their observed categories are $2,1,2,0$, respectively.

At $x=-1$, the unnormalized values are $1,\exp(-0.4)=0.670320,\exp(0.3)=1.349859$, with denominator $3.020179$. Thus the probabilities for categories $0,1,2$ are approximately $0.33110,0.22195,0.44695$, respectively, and the observed category $2$ contributes $0.44695$.

At $x=0$, all three unnormalized values equal one, so all probabilities are $1/3$. The observed category $1$ contributes $1/3$. At $x=1$, the probabilities are $0.3093,0.4615,0.2292$ as calculated above. The third individual received category $2$, contributing $0.2292$; the fourth received category $0$, contributing $0.3093$.

The likelihood contribution for these four observed assignments is therefore approximately $0.44695\times(1/3)\times0.2292\times0.3093=0.01056$. Each factor is a probability for the treatment actually observed at that individual's covariate value. This product is a likelihood contribution, not a probability that the fitted model is correct.

Diagnostics, Alternatives, and Limitations

A fitted multinomial model should be checked before its probabilities are used. Verify that each row of predicted probabilities is finite, lies between zero and one, and sums to one (allowing for rounding). Inspect category counts across important covariate patterns, predicted probabilities by category, convergence messages, and coefficient standard errors. Very large coefficients, unstable estimates across minor model changes, or fitted probabilities extremely close to zero can indicate sparse data or separation. Check whether the model captures important nonlinearities and interactions; a convenient linear predictor is not automatically an adequate conditional probability model.

The estimated probability for an observed treatment can be very small even if the model converges. Such observations signal a practical positivity concern: the observed treatment is rarely represented among people with similar covariates. Inverse use of that probability can make later weighted analyses unstable. Probability estimates do not establish conditional exchangeability, correct covariate timing, or absence of unmeasured confounding. Those are substantive design and causal-identification requirements, not consequences of fitting a multinomial regression.

A baseline-category model is a natural single model for unordered options because it estimates a coherent probability vector directly. Separate one-versus-rest binary regressions are an alternative, but their independently fitted probabilities need not sum to one; normalization changes the fitted values and does not generally reproduce multinomial maximum likelihood. Coding categories as a numeric score in an ordinary binary logistic regression is also inappropriate when the labels have no ordering. The model's reference category is a parameterization choice: changing the reference changes coefficient descriptions, while the underlying probability vector remains the quantity of interest.

The model is limited by its functional form, available sample size, and support in the data. A category with few observations may have poorly estimated probabilities, and a category absent from a covariate region cannot be rescued by the algebraic fact that a finite softmax probability is positive. If treatment choices have a genuine order, an ordinal model may be more suitable; if decisions are actually separate components or occur sequentially, they require a different treatment representation. Neither situation is solved simply by relabeling categories in a nominal model.

Common Mistake

Wrong reasoning: “The model gives the largest probability to category $1$, so I will use that probability for everyone.” This substitutes the most likely category for the category actually received.
Correction: Calculate the full probability vector for each person, then select $\hat p_{i,A_i}$. In the first example, a person with $x=1$ who received category $2$ contributes $0.2292$, even though category $1$ has the largest probability, $0.4615$.
Another warning: A set of normalized probabilities is not evidence by itself that the covariate model is well specified or that causal assumptions hold.

Key takeaway: A multinomial treatment model must produce one conditional probability per nominal category. The probability relevant to an observed assignment is the fitted value for that person's actual category, and its reliability depends on model adequacy, data support, and the causal assumptions of the analysis.

Check Your Understanding

Answer each question, then compare with the short solutions below.

For three treatment categories with predictors $\eta_1=\log(2)$ and $\eta_2=0$, what are the probabilities for categories $0,1,2$?
An individual has fitted probabilities $(0.2,0.5,0.3)$ for categories $0,1,2$, and received category $2$. Which probability is selected?
Why can independently fitted one-versus-rest probabilities fail to be a valid probability vector?
What does an extremely small fitted probability for an observed treatment suggest, and what does it not establish?

Solutions. (1) The unnormalized values are $1,2,1$, so the denominator is $4$, and the probabilities are $(0.25,0.50,0.25)$. (2) Select $0.3$, because category $2$ was observed. (3) Each regression models its own event separately, so the predicted probabilities have no constraint forcing their sum to equal one. (4) It suggests weak empirical support for that observed assignment at the person's covariates and possible instability if inverse probabilities are used. It does not, by itself, prove a causal assumption is violated or that the model is misspecified.