Deriving the MAP estimate for Multinomial-Dirichlet

Question

Let $X = (X_1, X_2,...,X_k) \sim \text{Multinomial}(n, \theta)$ on $k$ elements (i.e. each $X_i$ is the count of the element $i$ on $n$ trials). Let $\theta \sim \text{Dirichlet}(\alpha)$, where $\alpha = (\alpha_1, \alpha_2, ..., \alpha_k)$. I'm trying to derive the MAP estimate $\hat \theta_{MAP}$ for $\theta$.

My attempt:

By Bayes' theorem, $P(\theta | x) = \frac{\mathcal{L}(\theta) \pi(\theta)}{P(x)}$, where $\mathcal{L}(\theta)$ is the likelihood, $\pi(\theta)$ is the prior and $P(x)$ is the evidence. The MAP is the argmax of this expression (we can drop this since it doesn't depend on $\theta$). Replacing the definitions of the model's likelihood and prior, we get

$$ \hat \theta_{MAP} = \arg\max\big( \frac{\Gamma(\sum_i x_i + 1)}{\prod_i \Gamma(x_i+1)} \prod_{i=1}^k \theta_i^{x_i} \frac{\prod_{i=1}^K \Gamma(\alpha_i)}{\Gamma\bigl(\sum_{i=1}^K \alpha_i\bigr)} \prod_{i=1}^k \theta_i^{\alpha_i - 1} \big) $$

Dropping the terms that don't depend on $\theta$, we get

$$ \hat\theta_{MAP} = \arg\max\big( \prod_{i=1}^k \theta_i^{x_i} \prod_{i=1}^k \theta_i^{\alpha_i - 1} \big) \\ = \arg\max\big( \prod_{i=1}^k \theta_i^{x_i + \alpha_i - 1}\big) $$

Applying the logarithm:

$$ \hat\theta_{MAP} = \arg\max\big(\log\big( \prod_{i=1}^k \theta_i^{x_i + \alpha_i - 1}\big)\big)\\ =\arg\max\big(\sum_{i=1}^k \log \big(\theta_i^{x_i + \alpha_i - 1}\big)\big) \\ =\arg\max\big(\sum_{i=1}^k (x_i + \alpha_i - 1)\log \big(\theta_i\big)\big) $$

Deriving with respect to $\theta_i$ and equating to $0$: $$ 0 = \frac{\partial}{\partial \theta_i} \big(\sum_{i=1}^k (x_i + \alpha_i - 1)\log \big(\theta_i\big)\big) \\ = \sum_{i=1}^k \frac{(x_i + \alpha_i - 1)}{\theta_i} $$

And here I'm stuck since $\theta_i$ cancels. Where is my error?

Answer 1

I see two mistakes in your steps.

First of all, taking the derivative with respect to a single $\theta_m$ cancels out all the other terms where $i \ne m$ in the sum. That is, $$ \frac{\partial}{\partial \theta_m} \sum_i (x_i + \alpha_i - 1) \log(\theta_i) = \frac{x_m + \alpha_m - 1}{\theta_m}, $$ as opposed to what you said. You now have $k$ equations to solve, one for each $m.$

Secondly, it's not enough to simply maximize the expression $\sum_i (x_i + \alpha_i - 1) \log(\theta_i).$ The solution is obviously $\theta_i \rightarrow \infty,$ because $\log$ monotonically increases. As you well know, we have the constraint that $\sum_i \theta_i = 1.$ You have to encode that constraint into your optimization via a Lagrange multiplier so that the correct solution isn't obviously infinity. That is, we must find the stationary point of

$$ L(\theta, \lambda) = \sum_{i=1}^k (x_i + \alpha_i - 1) \log(\theta_i) - \lambda \left[ \sum_{i=1}^k \theta_i - 1 \right]. $$

Taking the derivative of the above expression with respect to each $\theta_i$ and WRT the Lagrange multiplier $\lambda,$ and setting them to zero, you have to simultaneously solve,

$$ \frac{x_m + \alpha_m - 1}{\theta_m} = \lambda $$ for each $m,$ and $$ \sum_i \theta_i = 1. $$

That's not very hard to do. Rearrange the first equation above to get an expression of $\theta_m$ in terms of $\lambda.$ Then, you can sensibly find the value of $\lambda$ that correctly imposes the constraint.

Answer 2

You can impose the constraint $\sum \theta_i = 1$ by specifying $\theta_k = 1 - \sum_{i<k}\theta_i$ in your likelihood function. This results in:

$$l(\theta) = \sum_{i=1}^{k-1}(x_i+a_i-1)\log \theta_i + (x_k+a_k-1)\log (1-\sum_{j=1}^{k-1}\theta_j)$$

Taking the derivative gives:

$$\frac{\partial l(\theta)}{\partial\theta_i} = (x_i+a_i-1)/\theta_i -(x_k+a_k-1)/(1-\sum_{j=1}^{k-1}\theta_j)$$

Substituting $\theta_k$ for $1-\sum_{j=1}^{k-1}\theta_j$ gives:

$$\frac{\partial l(\theta)}{\partial\theta_i} = (x_i+a_i-1)/\theta_i -(x_k+a_k-1)/\theta_k$$

which evidently implies that $\hat{\theta}_i \propto (x_i+a_i-1)$ for all $i$.

Combining this with the constraint that the $\hat{\theta}_i$ sum to 1 leads quickly to:

$$\hat{\theta}_i = \frac{x_i+a_i-1}{\sum_j(x_j+a_j-1)}$$

Answer 3

Notice that an optimizer to the optimization problem $$\arg\max_{y, \sum_{i=1}^Ky_i=1}\prod_{j=1}^K y_j^{\alpha_j-1}$$

satisfy $$(y_1, \ldots, y_n)$$ where $$y_i=\frac{\alpha_i-1}{\sum_{j=1}^ky_j -K}$$ as we can see from the wikipedia page of Dirichlet Distribution.

As you have shown the MAP have to maximize $$\arg\max_{\theta, \sum_{i=1}^K\theta_i=1}\prod_{j=1}^K \theta_j^{\alpha_j+x_j-1}$$

We are just adding data to the prior pseudo-count, hence the corresponding mode would be

$$\theta_i = \frac{\alpha_i+x_i-1}{\sum_{j=1}^k(\alpha_j+x_j) -K}$$