The Approximation - Big Picture

The saddlepoint approximation uses the cumulant-generating function (CGF) of a distribution to compute an approximate density at a given point. The neat part about CGFs is that the CGF of the sum of several variables is the sum of the individual CGFs! And CGFs are easy to come by, because the CGF is just the log of the moment-generating function and Wikipedia helpfully lists moment-generating functions for almost all distributions. Figuring out the CGF of almost any sum variable (including variables from different families) is thus relatively straightforward. The CGF can also easily be derived for general linear combinations of random variables.

The actual method for approximating density \(f\) at point \(x\), given the cumulant-generating function \(K\), and its first and second derivatives (\(K^\prime,K^{\prime\prime}\)) is as follows:

1. find the saddlepoint \(s_x\) by solving:

\[ K^\prime(s_x) = x \]

Generally, there is no closed-form solution for \(s_x\), but since \(K(x)\) is always convex, \(K^\prime\) is always increasing, making it a nice target for numerical solutions. Still, since a different solution is needed for each \(x\), finding \(s_x\) tends to be a computational bottleneck.

2. Once we have \(s_x\), we can approximate

\[ f(x) \simeq \frac1{\sqrt{2\pi K''(s_x)}} \exp(K(s_x) - x s_x) \]

The nice thing about the saddlepoint approximation is that it can easily produce approximations for both discrete and continous densities, and doesn’t constrain the approximation to be normal (unlike Laplace approximation). One thing to note is that the saddlepoint approximation in the form above does not necessarily integrate to 1, so a renormalization might be needed if you are interested in the actual density. But to use in Stan, unnormalized density is all that’s needed.

Saddlepoint for Sum of NBs

The moment-generating function of NB distribution parametrized by number of failures \(r\) and probability of success \(p\) is:

\[ M(t) = \left( \frac{1 - p}{1 - p e^t} \right)^r \]

So, taking the log and summing over \(n\) independent NB variables, the cumulant of sum of NB is:

\[ K(t) = \sum_{i=1}^{n} r_i \left[ \log(1-p_i) - \log(1 - p_i e^t) \right] \]

We now transform to the more useful parametrization of NB via mean \(\mu\) and precision \(\phi\) (i.e. \(Var(X) = \mu + \frac{\mu^2}{\phi}\)), where we have:

\[ r_i = \phi_i \\ p_i = \frac{\mu_i}{\phi_i + \mu_i} \\ K(t) = \sum_{i=1}^{n} \phi_i \left[ \log \frac{\phi_i}{\phi_i + \mu_i} - \log \left(1 - \frac{\mu_i e^t}{\phi_i + \mu_i} \right) \right] = \\ =\sum_{i=1}^{n} \phi_i \left[ \log(\phi_i) - \log(\phi_i + \mu_i ( 1 - e^t)) \right] \]

Note that \(K(t)\) does exist only when \(\forall i:\phi_i + \mu_i ( 1 - e^t) > 0\) this constrains \(t\) such that:

\[ \begin{align} \tag{*} \forall i : t &< log \left(\frac{\phi_i}{\mu_i} + 1 \right) \end{align} \]

The first and second derivatives of \(K\) are:

\[ K^\prime (t) = \sum_{i=1}^{n} \frac{\phi_i \mu_i e^t}{\phi_i + \mu_i (1 - e^t)} \\ K^{\prime\prime} (t) = \sum_{i=1}^{n} \frac{\phi_i \mu_i (\phi_i + \mu_i) e^t}{(\phi_i + \mu_i (1 - e ^t))^2} \\ \]

It turns out that the saddlepoint \(s_x\) is not defined when \(x = 0\), since the numerator of \(K^\prime(t)\) is positive for all \(t\) and the denominator has to be positive for \(K\) to exist. But for this special case, the density can be easily computed, as \(f(0) = \prod_i P(X_i =0) = \prod_i NB(0 | \mu_i,\phi_i)\). The non-existence of the saddlepoint solution for boundaries of the domain is actually a recurring theme, as the existence of the solution is guaranteed only for the inner points, so it is useful to check for this when developing your approximations.

Implementing the Approximation in Stan

This has all been a nice math excercise, but how can we translate that into a piece of code we could use? The only problematic part is solving for \(s_x\), once we have it, the rest is a simple math that Stan will digest easily. Luckily, Stan has the built-in algebra_solver that can solve equations AND provide derivatives of the solution wrt. parameters. There is only a minor problem - we have an upper bound on \(s_x\) from the equation \((*)\) and algebra_solver turns out not to work when there are boundaries (even when initialized within the boundaries). Instead we use the same method Stan uses for bounds on parameters and solve for unbounded \(y_x\) where:

\[ s_x = \min_i{log \left(\frac{\phi_i}{\mu_i} + 1 \right)} -e^{y_x} \]

So let us get our hands dirty and show some code, starting with how to write the saddlepoint equation in a way that the algebra_solver can handle. Since \(K^\prime\) is always positive, we transform the equation to log scale - partly because we might have some big \(\sum\mu_i\) out there and partly because it seems nice - I didn’t test the non-log version. So the equation we are actually solving for \(s_x\) is:

\[ \log \sum_{i=1}^{n} \exp \left( \log\phi_i + \log \mu_i + s_x - \log(\phi_i + \mu_i - \mu_i \exp(s_x) \right) - x = 0 \] Translated into Stan we get:

  //Compute upper bound on s - there is a singularity at this point
  real max_s(vector mus, vector phis) {
    return min(log(phis ./ mus + 1));
  }

  //Transform from unbounded y to s upper-bounded by max_s
  vector s_transform(vector y, vector mus, vector phis) {
    return -exp(y) + max_s(mus, phis);
  }

  vector nb_sum_log_Kd_eq(vector y, vector theta, real[] x_r, int[] x_i) {
    int G = rows(theta) / 2;
    vector[G] mus = theta[1:G];
    vector[G] phis = theta[(G + 1) : (2 * G)];

    real s = s_transform(y, mus, phis)[1];
    real sum_y = x_i[1];
    vector[G] log_phis_mus = log(phis) + log(mus);
    vector[G] phis_mus = phis + mus;
    
    real value = log_sum_exp(log_phis_mus + s - log(phis_mus - mus * exp(s))) - log(sum_y);
    return to_vector({value});
  }

Above, y are the unconstrained unknowns, which we transform via s_transform to the constrained space. Further we extract \(\mu_i\) and \(\phi_i\) from theta which can be parameters while x_i contains the observed sums (data). Since we have no real number data, x_r is ignored. The algebra_solver will try to find y such that value is 0 which is exactly when s is the solution to the saddlepoint equation.

We use the nb_sum_log_Kd_eq function to compute the actual saddlepoint density:

  real neg_binomial_sum_saddlepoint_lpmf(int[] sum_y, vector mus, vector phis, real[] dummy_x_r) {
    int N = size(sum_y);

    int G = rows(mus);

    // Solve the saddlepoint equation
    vector[2 * G] solver_params = append_row(mus, phis);

    vector[N] s_vec_raw;
    vector[1] solver_guess = to_vector({0});
    for(n in 1:N) {
      if(sum_y[n] != 0) {
        //Saddlepoint is defined only for non-zero values
        s_vec_raw[n] = algebra_solver(nb_sum_log_Kd_eq, solver_guess, solver_params, dummy_x_r,  {sum_y[n]})[1];
      } else {
        //This will be ignored, but needed to pass to s_transform without problems
        s_vec_raw[n] = 0;
      }
    }

    {
      vector[N] s = s_transform(s_vec_raw, mus, phis);
      //Calculate the saddlepoint mass
      vector[N] lpmf;
      vector[G] log_mus = log(mus);

      for(n in 1:N) {
        if(sum_y[n] != 0) {
          vector[G] log_denominator_s = log(phis + mus - mus * exp(s[n]));
          real K_s = sum(phis .* (log(phis) - log_denominator_s));
          real log_Kdd_s = log_sum_exp(log(phis) + log_mus + log(phis + mus) + s[n] - 2 * log_denominator_s);
          lpmf[n] = -0.5 * (log(2*pi()) + log_Kdd_s) + K_s - s[n] * sum_y[n] ;
        } else {
          //For zero values, the probability is simply that of all NBs giving 0 
          lpmf[n] = neg_binomial_2_lpmf(rep_array(0, G) | mus, phis);
        }
      }
      
      return sum(lpmf);
    }
  }

The above shows how the algebra_solver is called - we combine \(\mu_i\) and \(\phi_i\) as params, pass a guess (0 works great, so we don’t need to worry about details). The only weird part is dummy_x_r - I want it to be just an empty array, but it has to be of type real and has to be data. And I didn’t find a way to make the compiler understand that unless I pass dummy_x_r from outside as in

transformed data {
  real dummy_x_r[0];
}

...

model {
  sums ~ neg_binomial_sum_lpmf(mus, phis, dummy_x_r);
}

A Simple Baseline

To assess, how useful the saddlepoint approximation is in practice, we’ll compare it to a straightforward application of Method of moments. This is just a fancy name for choosing a distribution family and choosing it’s parameters so that mean, variance (and possibly higher moments) match those of the desired distribution. In case of NBs, when \(Y_i \sim NB(\mu_i, \phi_i)\) then

\[ E \left(\sum Y_i \right) = \sum \mu_i \\ Var \left(\sum Y_i \right) = \sum \left( \mu_i + \frac{\mu_i^2}{\phi_i} \right) \]

Simply because both mean and variance are linear operators. Maybe sum of NBs isn’t that different from a NB distribution, so let’s approximate

\[ \sum Y_i \approx NB(\bar\mu, \bar\phi) \]

Solving for \(\bar\mu\) and \(\bar\phi\) by matching the mean and variance of the approximate distribution gives:

\[ \bar \mu = \sum \mu_i \\ \bar \phi = \frac{ \left(\sum \mu_i \right)^2 }{\sum \frac{\mu_i^2}{\phi_i}} \]

This can be implemented very directly in Stan as:

  real neg_binomial_sum_moments_lpmf(int[] sum_y, vector mus, vector phis) {
    real mu_approx = sum(mus);
    real phi_approx = square(mu_approx) / sum(square(mus) ./ phis);

    return neg_binomial_2_lpmf(sum_y | mu_approx, phi_approx);
  }

Eyeballing Masses

As a first look, we will see how well do both approximations match the empirical mass function - we simulate a lot of sums of NBs, bin the range of observed values and compute empirical mass as the proportion of the samples that fits in each bin. For the approximations, we sum the mass for all values belonging to the bins.

The saddlepoint approximation improves notably over moments when the Fano factors of the summed variables are vastly different and we do not sum a large number of values, below we show mass and log mass for the case when \(\mu = \{800, 1600 \}\) and \(\phi = \{10, 1 \}\):

It is visible that the saddlepoint mass tracks the empirical mass very tightly both in the bulk and in the tail (visible better on the log mass) - note that the tail of the empirical log mass is jittery due to low number of samples in the tail.

On the other hand, when we sum a lot of variables which are not very different and/or when \(\phi_i\) are large, the sum becomes normal-ish and both approximation work well - let us for example look at the case when \(\mu = \{50, 100, 1300, 2000 \}\) and \(\phi = \{10, 10, 10, 10 \}\):

This gives us some intuition where to expect differences.

Evaluating Performance

We evaluate the model using Simulation-based Calibration (SBC). The main idea is that when I generate data exactly the way the model assumes, then for any \(c\) the \(c\%\) posterior interval should contain the true value an unobserved parameter in exactly \(c\%\) of the cases. In other words the quantile in which the true value is found in the posterior distribution should be uniformly distributed. There are some caveats to this, read the paper for details.

I am using my own implementation of SBC which is in my not very well documented, likely-never-on-CRAN package rstanmodeldev. We run 500 simulations for each of the test cases. If you want to see under the hood, the code for this post is available at the GitHub repo of this blog.

The first test case I will use is that I observe the sum of \(G+1\) variables where I know \(\mu_i\) and \(\phi_i\) for \(i \in {1 .. G}\) while \(\mu_{G+1}\) and \(\phi_{G+1}\) is unknown and has to be infered from \(N\) observations of the sum.

In all cases, both observed and unobserved \(\phi_i\) are drawn following the advice of Dan simpson, i.e.:

\[ \phi_{raw} \sim HalfN(0, 1) \\ \phi = \frac{1}{\sqrt{\phi_{raw}}} \]

This is how the model looks-like in Stan ( test_sum_nb.stan ):

#include /sum_nb_functions.stan

data {
  int<lower=1> N;
  int<lower=0> sums[N]; 
  int<lower=1> G;
  vector[G] mus;
  vector[G] phis;

  //0 - saddlepoint, 1 - method of moments
  int<lower = 0, upper = 1> method;
  real mu_prior_mean;
  real<lower = 0> mu_prior_sd;
  
}

transformed data {
  real dummy_x_r[0];
}

parameters {
  real log_extra_mu_raw;
  real<lower=0> extra_phi_raw;
}

transformed parameters {
  real<lower=0> extra_mu = exp(log_extra_mu_raw * mu_prior_sd + mu_prior_mean);
  real<lower=0> extra_phi =  inv(sqrt(extra_phi_raw));
}

model {
  vector[G + 1] all_mus = append_row(mus, to_vector({extra_mu}));
  vector[G + 1] all_phis = append_row(phis, to_vector({extra_phi}));

  if(method == 0) {
    sums ~ neg_binomial_sum_saddlepoint_lpmf(all_mus, all_phis, dummy_x_r);
  } else {
    sums ~ neg_binomial_sum_moments_lpmf(all_mus, all_phis);
  }

  log_extra_mu_raw ~ normal(0, 1);
  extra_phi_raw ~ normal(0,1);
}

Most notably, the way the sum of NBs is implemented is given as data. The sum_nb_functions.stan include contains the functions shown above.

And this is an R method to generate simulated data - this is a function that given parameters of the observed data gives a function that on each call generates both true and observed data in a format that matches the Stan model:

generator <- function(G, N, method = "saddlepoint", observed_mean_mus, observed_sd_mus, mu_prior_mean, mu_prior_sd) {
  if(method == "saddlepoint") {
    method_id = 0
  } else if (method == "moments") {
    method_id = 1
  } else {
    stop("Invalid method")
  }
  
  function() {
    all_mus <- rlnorm(G + 1, observed_mean_mus, observed_sd_mus)
    all_mus[G + 1] <- rlnorm(1, mu_prior_mean, mu_prior_sd)
    all_phis <- 1 / sqrt(abs(rnorm(G + 1)))
    sums <- array(-1, N)
    for(n in 1:N) {
      sums[n] <- sum(rnbinom(G + 1, mu = all_mus, size = all_phis))
    }
    list(
      observed = list(
        N = N,
        sums = sums,
        G = G,
        method = method_id,
        mus = array(all_mus[1:G], G),
        phis = array(all_phis[1:G], G),
        mu_prior_mean = mu_prior_mean,
        mu_prior_sd = mu_prior_sd
      ),
      true = list(
        extra_mu = all_mus[G+1],
        extra_phi = all_phis[G+1]
      )
    )
  }
}

#include /sum_nb_functions.stan

data {
  int<lower=1> N;
  int<lower=0> sums[N]; 
  int<lower=1> G;

  //0 - saddlepoint, 1 - method of moments
  int<lower = 0, upper = 1> method;
  
  real mu_prior_mean;
  real<lower = 0> mu_prior_sd;
  
  real<lower=0> mu_series_coeff;
  real<lower=0> phi_series_coeff;
}

transformed data {
  real dummy_x_r[0];
  vector[G] mu_coeffs;
  vector[G] phi_coeffs;
  
  mu_coeffs[1] = 1;
  phi_coeffs[1] = 1;
  for(g in 2:G) {
    mu_coeffs[g] = mu_coeffs[g - 1] * mu_series_coeff;
    phi_coeffs[g] = phi_coeffs[g - 1] * phi_series_coeff;
  }
}

parameters {
  real log_mu_raw;
  real<lower=0> phi_raw;
}

transformed parameters {
  real<lower=0> mu = exp(log_mu_raw * mu_prior_sd + mu_prior_mean);
  real<lower=0> phi =  inv(sqrt(phi_raw));
}

model {
  vector[G] all_mus = mu * mu_coeffs;
  vector[G] all_phis = phi * phi_coeffs;

  if(method == 0) {
    sums ~ neg_binomial_sum_saddlepoint_lpmf(all_mus, all_phis, dummy_x_r);
  } else {
    sums ~ neg_binomial_sum_moments_lpmf(all_mus, all_phis);
  }

  log_mu_raw ~ normal(0, 1);
  phi_raw ~ normal(0,1);
}

generator_series <- function(G, N, method = "saddlepoint", mu_prior_mean, mu_prior_sd, mu_series_coeff, phi_series_coeff) {
  if(method == "saddlepoint") {
    method_id = 0
  } else if (method == "moments") {
    method_id = 1
  } else {
    stop("Invalid method")
  }
  
  function() {
    mu <- rlnorm(1, mu_prior_mean, mu_prior_sd)
    phi <- 1 / sqrt(abs(rnorm(1)))
    all_mus <- mu * mu_series_coeff ^ (0:(G - 1))
    all_phis <- phi * phi_series_coeff ^ (0:(G - 1))
    sums <- array(-1, N)
    for(n in 1:N) {
      sums[n] <- sum(rnbinom(G, mu = all_mus, size = all_phis))
    }
    list(
      observed = list(
        N = N,
        sums = sums,
        G = G,
        method = method_id,
        mu_prior_mean = mu_prior_mean,
        mu_prior_sd = mu_prior_sd,
        mu_series_coeff = mu_series_coeff,
        phi_series_coeff = phi_series_coeff
      ),
      true = list(
        mu = mu,
        phi = phi
      )
    )
  }
}

Summing up

We see that in the regimes we tested, the saddlepoint approximation for sum of negative binomials provides somewhat better inferences for small number of variables at the cost of much increased computation times. For sums of large number of variables, it may even be worse than the moments method. So it is probably not very practical unless you have few variables you need that extra bit of precision. But it is a neat mathematical trick and of interest on its own. It is also possible that for some low-mean regimes the difference is bigger.

Saddlepoint Approximations for Other Families

If you want to use saddlepoint approximation for other than NB variables, but don’t want to do the math on your own, there are some worked out on the Internet:

• Sum of Gamma variables: Answer on Cross Validated
• Sum of binomials: Liu & Quertermous: Approximating the Sum of Independent Non-Identical Binomial Random Variables

Thanks for reading!