Numerical instabilities in Laplace EM for Poisson emissions with large spike counts

[slinderman]

@KeAWang and @kochalka and I are trying to debug a bunch of NaNs in the Newton update for the continuous states in Laplace EM. I'll follow up with more details, but it looks like the issues arise in cases where we have large spike counts and Poisson or OrthogonalPoisson emissions. I found a few places where it looked like we could get divide by zero errors in the likelihood and Hessian calculations, but haven't tracked it down yet.

[KeAWang]

I tried a more numerically stable version of inv_softplus (from cornellius-gp/gpytorch#776):

def inv_softplus(y):
    return y + np.log(-np.expm1(-y))

to replace what's in ssm.utils, but still getting errors.

[slinderman]

@KeAWang and @kochalka, as a workaround, it looks like it should work if you use lbfgs instead of newton for the continuous state update.

# Fit the model using Laplace-EM with a structured variational posterior
q_lem_elbos, q_lem = slds.fit(train_data, method="laplace_em",
                              variational_posterior="structured_meanfield",
                              num_iters=num_iters, alpha=0.0,
                              continuous_optimizer="lbfgs")

You might still have to add a little bias to the output of the softplus in order to ensure the rates aren't numerically zero. Otherwise you'll get nan's in the likelihoods.