Initialize Dirichlet fitting without sample variable

[andreasnoack]

This avoids dividing by zero. The new initialization is based on the paper by Thomas Minka linked in the comment. This is an alternative to #610 based on some of the ideas in the paper. It also avoids the extra allocation. Of course, at the cost of some special function evaluations.

Fixes #602

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 86.74%. Comparing base (2697768) to head (6d7c210).

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #2077      +/-   ##
==========================================
- Coverage   86.74%   86.74%   -0.01%     
==========================================
  Files         149      149              
  Lines        8883     8881       -2     
==========================================
- Hits         7706     7704       -2     
  Misses       1177     1177              

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[devmotion]

AFAICT the PR mixes two approaches presented by Minka, the EM-like fixed-point iteration and the Newton algorithm. It's not immediately clear to me why one would want to combine those two approaches and switch after a seemingly arbitrary amount of 5 fixed-point iterations to the Newton algorithm. IMO the crucial and currently in a few cases problematic part is just the initialization of the algorithms. As Minka mentions as well, there's many different possible approaches for initializing the algorithms, e.g., based on method of moments (obviously, different choices of moments will lead to different initializations, Minka already discusses two different approaches) or approximations of the likelihood function, e.g., using Stirling's approximation.

Exploiting the moment equations $E\left[p_k\right] = \alpha_k / \alpha_0$ for $k = 1, \ldots, K$ (as currently done) seems natural. However, since one of these $K$ equations is redundant, it leaves the question what additional equation/constraint to include to obtain an equation system that can be solved for the $K$ parameters $\alpha_1, \ldots, \alpha_K$. Minka presents different approaches in eq. 21, 23 and 42. The approaches in eq. 21 and 23 break down if the sample variance of $p_1$ (eq 21) or any $p_k$ (eq 23) is zero. Eq. 42 doesn't suffer from this problem, so as suggested in this PR it would be a more stable approach. Alternatively, one could consider other moments such as

$$E\left[\sum_k p_k^2\right] = \frac{1 + \alpha_0 \sum_k {(\alpha_k / \alpha_0)}^2}{1 + \alpha_0}$$

which AFAICT would give moment estimates

$$\alpha_k = \hat{E}\left[p_k\right] \left(\frac{1 - \sum_{\tilde{k}} \hat{E}\left[p_{\tilde{k}}^2\right]}{\sum_{\tilde{k}} \hat{E}\left[p_{\tilde{k}}^2\right] - \sum_{\tilde{k}} {\hat{E}\left[p_{\tilde{k}}\right]}^2} \right) = \hat{E}[p_k] \frac{\sum_{\tilde{k}} \hat{E}\left[p_{\tilde{k}}(1-p_{\tilde{k}})\right]}{\sum_{\tilde{k}} \hat{var}\left[p_{\tilde{k}}\right]} $$

(seemingly a symmetrized version of eq. 21?). This would only break down if the sample variance of all $p_k$ is zero, ie., all samples are identical, in which case the estimation breaks down anyway.

So, taken together, my impression is: Using eq. 42 would give a less brittle initialization but one could also consider other approaches, e.g., based on different moments; but it's not clear to me why one would want to combine the EM-like fixed-point iteration with the Newton algorithm, using a different more stable initialization for the Newton algorithm alone (or the EM-like fixed-point iteration, if one prefers that approach) seems sufficient.

[andreasnoack]

The problem is that the parameters must be positive, so the initialization has to be sufficiently good to ensure that the Newton step is feasible. The fixed-point iterations will always produce positive values, but it has linear convergence. I also tried the initialization suggested in https://www.sciencedirect.com/science/article/abs/pii/S0167947307002848, but then I can trigger a maxiter error by making the repeated value small enough.

[devmotion]

To me it seems there's two separate problems you're facing and trying to address: the problematic initialization (the original issue reported in #602) and infeasible iterates of the Newton algorithm. Seemingly the first problem could be fixed by an initialization that doesn't break down if the sample variance of some $p_k$ is zero, e.g., by making use of eq. 42 in Minka's report, the estimate for $\alpha_0$ in the linked paper or possibly by considering different moments as outlined above. The second problem seems merely a consequence of applying (unconstrained) Newton method on a constrained space, and could maybe be approached by e.g. applying Newton method on the log space of $\log \alpha_k$ (AFAICT the efficient structure of the Hessian would be preserved) or using a penalty method.

[andreasnoack]

The two issues are closely related. In theory, we could just initialize all the values at one, except for the risk of stepping to an infeasible argument. Hence, we need to find a sufficiently good initialization before switching to Newton's method and thereby ensure fast convergence through the concave objective function.

Seemingly the first problem could be fixed by an initialization that doesn't break down if the sample variance of some $p_k$ is zero, e.g., by making use of eq. 42 in Minka's report, the estimate for $α_0$ in the linked paper or possibly by considering different moments as outlined above.

Yes, but as mentioned above, I could trigger maxiter issues with these initializers.

...e.g. applying Newton method on the log space of $log(α_k)$ (AFAICT the efficient structure of the Hessian would be preserved)

The global concavity of the objective function is not preserved under this reparametrization. We could try a constrained optimizer, but that would be a larger change than the current one.

[devmotion]

To avoid the concavity issue, we could use natural gradient descent in log space instead of of Newton's method. In the standard parameter space, Newton's method is exactly the natural gradient method with step size 1 (the Hessian is constant with respect to the data and hence identical to the FIM). AFAICT this would correspond to iterations $\alpha \to \alpha \odot \exp.{\left(\gamma s(\alpha) ./ \alpha\right)}$ with step size $\gamma$ and $s(\alpha)$ being the step direction of the corresponding Newton iteration in non-log space $\alpha \to \alpha + s(\alpha)$.