Sum-to-zero with a covariance matrix

Hi all,

Given the surge of posts about sum-to-zero recently I thought I’d add another one. I’ve posted about this on Slack before but it might be of broader interest.

My question is about implementing sum-to-zero vectors with a covariance matrix using non-centered parameterisation. Consider the following Stan program:

data {
  int<lower=0> J;  // number of observations
  array[J] real x;  // input
  int<lower=0, upper=1> flag;
}
parameters {
  sum_to_zero_vector[J] z;
  real<lower=0> tau;
  vector<lower=0>[flag ? 1 : 0] rho;
}
transformed parameters {
  vector[J] theta = tau * z;  // s2z Gaussian
  if (flag) {
    matrix[J, J] K = gp_exp_quad_cov(x, 1, rho[1]),
                 L = cholesky_decompose(K);
    theta = L * theta; 
  }
}
model {
  // also add priors for tau, rho
  z ~ std_normal();  // or normal(0, sqrt(J * inv(J - 1)).
}

When flag = 0, the sum-to-zero vector theta is scaled by tau so the vector still sums to 0. When flag = 0, the vector is additionally scaled by a covariance matrix so the vector doesn’t still sum to zero. I often write Stan programs this way to I can toggle complexity on and off.

My question is mostly about the “validity” of this sum-to-zero approach. I suspect there’s still sampling efficiency by enforcing the constraint at the level of the z-scores, but it’s not the same as putting a centered parameteristion like this:

model {
  matrix[J, J] K = gp_exp_quad_cov(x, tau, rho[1]),
               L = cholesky_decompose(K);
  z ~ multi_normal_cholesky(L);
}

because her, I believe, z is constrained to sum to zero while having the GP prior.

I’m mostly interested in the implications of having the z-scores summing to 1 and then multiplying by a Cholesky factor, vs. having the multi normal prior while also summing to 0.

Cheers,

Matt

I don’t see a problem putting a sum-to-zero constraint on the standardized form. We transform sum-to-zero vectors all the time.

P.S. It took me a few passes to parse this and even realize it was valid Stan code:

  vector[J] theta = tau * z;  // s2z Gaussian
  if (flag) {
    matrix[J, J] K = gp_exp_quad_cov(x, 1, rho[1]),
                 L = cholesky_decompose(K);
    theta = L * theta; 

I’d recommend the following form without all the intermediate declarations:

  vector[J] theta
    = flag
    ? cholesky_decompose(gp_exp_quad_cov(x, 1, rho[1])) * (tau * z);
    : tau * z;

The parens are necessary in the second line for efficiency to ensure that tau is multiplied by a vector rather than a matrix.

P.P.S.

can just be

vector<lower=0>[flag];

because if flag in { 0, 1 }, then flag ? 1 : 0 == flag.

I may not follow, but after multiplying the sum-to-zero standard-normals by a Cholesky factor of a covariance matrix, the resulting vector doesn’t sum to zero anymore. I’m curious about the implications of this vs. putting a multi-normal-Cholesky prior on the sum-to-zero vector itself. I believe these result in different things?

Ha! Here I thought I was writing it more legibly for people to follow; your proposed form is how I tend to write it my programs. I believe your suggestion can be further simplied to:

cholesky_decompose(gp_exp_quad_cov(x, tau, rho[1])) * z;

Yes, they’re different things. It might be useful to work out the distribution of the random variables X, Y \in \mathbb{R}^N to see for yourself.

  1. Y \sim \textrm{multi-normal}(0, \Sigma) s.t. \textrm{sum}(Y) = 0.

  2. X = \textrm{cholesky}(\Sigma) \cdot Z, where Z \sim \textrm{multi-normal}(0, \textrm{I}) s.t. \textrm{sum}(Z) = 0.

As you note, they have different support (different domains considered as random variables), so they’re clearly different. For Y, the model says the effects sum to zero. For X, the model says the standardized effects sum to zero. Both approaches identify the variable if it’s being used as a varying effect in a regression with an intercept. You can’t add or subtract from all the variables either way and preserve the sum-to-zero property. Also, both have the correct N - 1 degrees of freedom for identifiability.

It’s a good question as to which would would better empirically. Maybe from the perspective of what kind of posterior geometry do they generate. (1) samples Y and (2) samples Z. That makes me think that (2) would be better.

Let’s ask people better at this than me by pinging them: @spinkney, @andrewgelman, and @avehtari.

Let’s derive the centered vs non-centered versions of these. @Bob_Carpenter is right about the differences between models 1 and 2.

Let g_{\text{free}} be a typical, unconstrained GP

g_{\text{free}} \sim \mathcal{N}(0, \tau^2 K).

One key element to my Stancon 2026 talk is that we derive the distribution of the sum of g_{\text{free}} where we will call it S_{\text{free}} = \sum g_{\text{free}}. The distribution is

S_{\text{free}} \sim \mathcal{N}(0, \tau^2 \bf{1}^{\top} K \bf{1})

where \bf{1} is a row vector of ones. The covariance between the vector g_{\text{free}} and S_{\text{free}} is

\begin{aligned} \mathrm{Cov} &= \mathrm{Cov}(g, \mathbf{1}^{\top} g) \\ &=\mathrm{Cov}(g, g) \mathbf{1} \\ &= \tau^2 K \mathbf{1}. \end{aligned}

We can do two cool things here. The first is that we can regress g on the sum S (dropping the free subscripts for now) to get

a = \frac{\mathrm{Cov}(g, S)}{\mathrm{Var}(S)} = \frac{\tau^2 K \mathbf{1}}{\tau^2 \mathbf{1}^{\top} K \mathbf{1}} = \frac{K \mathbf{1}}{\mathbf{1}^{\top} K \mathbf{1}}.

The residual sum of squares is 0 thus let the residual squares be

r = g - a S

then \sum r = 0. We can derive the covariance matrix of r as

\begin{aligned} \mathrm{Cov}(r) &= \mathrm{Cov}(g - a S) \\ &= \mathrm{Var}(g) + a \mathrm{Var}(S) a^{\top} - \mathrm{Cov}(g, S) a^{\top} - a \mathrm{Cov}(g, S) \\ &= \tau^2 \underbrace{ \left(K - \frac{K \mathbf{1}\mathbf{1}^{\top} K}{\mathbf{1}^{\top} K \mathbf{1}} \right )}_{K_0}. \end{aligned}

Now, with those preliminaries out of the way, we get to our sum-to-zero vector by applying the s2z transform: z = H y_{\text{free}} where y is size J - 1 and z is size J. We act on the J - 1 subspace of J where everything sums to 0. The covariance happens to be the same as derived above with the projection to J that lies on the sum-to-zero plane:

\begin{aligned} C &= H^{\top} \mathrm{Cov}(r) H \\ &= \tau^2 H^{\top} K_0 H \end{aligned}

The centered parameterization has one issue where if we use the s2z vector we’ve placed an additional Gaussian element that we need to remove which we can below:

data {
  int<lower=2> J;
  array[J] real x;
}

transformed data {
  vector[J] q = rep_vector(inv_sqrt(1.0 * J), J);
 vector[J] mu = rep_vector(0.0, J);
}
parameters {
  sum_to_zero_vector[J] z;

  real<lower=0> tau;
  real<lower=0> rho;
}
transformed parameters {
  matrix[J, J] K = gp_exp_quad_cov(x, 1.0, rho);
  matrix[J, J] L_K = cholesky_decompose(K);
}

model {
  tau ~ normal(0, 1);
  rho ~ lognormal(0, 1);

  z ~ multi_normal_cholesky(mu, tau * L_K);

  /*
   * Normalize the Gaussian density after conditioning on q' z = 0.
   * This matters when tau or rho is estimated.
   */
  target -= normal_lpdf(
    0.0 | 0.0,
    tau * sqrt(dot_product(q, K * q))
  );
}

generated quantities {
  real sum_z = sum(z);
}

The non-centered

functions {
  /**
   * Orthonormal basis for the sum-to-zero subspace.
   *
   * H' * H = I_(J-1)
   * H' * 1 = 0
   */
  matrix s2z_basis(int J) {
    matrix[J, J - 1] H = rep_matrix(0.0, J, J - 1);

    for (k in 1:(J - 1)) {
      real a = inv_sqrt(k * (k + 1.0));

      for (j in 1:k) {
        H[j, k] = a;
      }

      H[k + 1, k] = -k * a;
    }

    return H;
  }
}

data {
  int<lower=2> J;
  array[J] real x;
}

transformed data {
  matrix[J, J - 1] H = s2z_basis(J);
  vector[J] q = rep_vector(inv_sqrt(1.0 * J), J);
}

parameters {
  vector[J - 1] eta;
  real<lower=0> tau;
  real<lower=0> rho;
}
transformed parameters {
  matrix[J, J] K = gp_exp_quad_cov(x, 1.0, rho);
  vector[J] Kq =  K * q;
  matrix[J, J] K_zero =   K - (Kq * Kq') / dot_product(q, Kq);
  matrix[J - 1, J - 1] C = cholesky_decompose(quad_form(H, K_zero));
  matrix[J - 1, J - 1] L_C;

  vector[J - 1] y =  tau * (L_C * eta);
  vector[J] z = sum_to_zero_constrain(y); //should be same as H * y
}
model {
  eta ~ std_normal();
  tau ~ std_normal();
  rho ~ lognormal(0, 1);
}
generated quantities {
  real sum_z = sum(z);
}

I bet there’s a faster way to do quad_form(H, K_zero) since we know H.

Thanks, @spinkney.

You’d win that bet, because

quad_form(H, K_zero) = A' * A,

where L_H = cholesky_decompose(H) and A = L_H' * K_zero. This doesn’t help with the Cholesky factorization, but if you QR decompose A, then R is the Cholesky factor of quad_form(H, K_zero). I don’t know how optimized Eigen is for these specializations.

I think it should be quad_form(K_zero, H) because it’s H^T * K_zero * H.

Some back and forth with Claude landed on

functions {
  // H' (K - K 1 1' K / 1'K1) H for symmetric K: the covariance of a GP
  // conditioned on 1'f = 0, compressed to the sum-to-zero subspace.
  // Helmert basis h_k = w_k (1_k, -k, 0), w_k = 1 / sqrt(k (k + 1)).
  // Two scalar loops, no J x J intermediates.
  matrix quad_form_helmert_schur(matrix K, data vector w) {
    int J = rows(K);
    matrix[J - 1, J - 1] B;
    vector[J] r = K * rep_vector(1.0, J);      // K 1
    real d = sum(r);                            // 1' K 1
    vector[J - 1] v;                            // H' K 1
    vector[J] p = rep_vector(0, J);             // p[i] = sum_{j <= l} K[i, j]

    {
      real s = 0;
      for (k in 1:(J - 1)) {
        s += r[k];
        v[k] = w[k] * (s - k * r[k + 1]);
      }
    }

    for (l in 1:(J - 1)) {
      real acc = 0;                             // sum_{i' < i} c[i']
      real vl_d = v[l] / d;
      for (i in 1:J) {
        p[i] += K[i, l];
        if (i <= l + 1) {
          real c = p[i] - l * K[i, l + 1];      // (K H)[i, l] / w[l]
          if (i > 1) {
            B[i - 1, l] = (w[i - 1] * w[l]) * (acc - (i - 1) * c)
                          - v[i - 1] * vl_d;
            B[l, i - 1] = B[i - 1, l];
          }
          acc += c;
        }
      }
    }
    return B;
  }
}

data {
  int<lower=2> J;
  array[J] real x;
}

transformed data {
  vector[J - 1] w;
  for (k in 1:(J - 1)) w[k] = inv_sqrt(k * (k + 1.0));
}

parameters {
  vector[J - 1] eta;
  real<lower=0> tau;
  real<lower=0> rho;
}

transformed parameters {
  vector[J - 1] y;
  vector[J] z;
  {
    matrix[J - 1, J - 1] L_C
      = cholesky_decompose(quad_form_helmert_schur(gp_exp_quad_cov(x, 1.0, rho), w));
    y = tau * (L_C * eta);
    z = sum_to_zero_constrain(y);   // z = H y, Cov(z) = tau^2 (K - K 1 1' K / 1'K1)
  }
}

model {
  eta ~ std_normal();
  tau ~ std_normal();
  rho ~ lognormal(0, 1);

  // Add the same likelihood involving z here.
}

generated quantities {
  real sum_z = sum(z);
}

This is all getting rather complicated for me, so sorry if I can’t engage with the math very well. One thing to add is that matrix normals are really convenient in the non-centered form, e.g. \bf{LzU}, where \bf{L} and \bf{U} are lower/upper triangular Cholesky factors of row- and column covariances, respectively. Especially for things like hierarchical GPs, where we might estimate a “mean row effect” and “mean column effect”, the sum_to_zero_matrix type is really convenient for this and I don’t really know how we’d implement it without the non-centered parameterisation in Stan. But my main concern was that both approaches are a valid way to identify the parameters, which Bob has confirmed!

Hey Sean,

Thanks so much for this, I’m implementing this in a model now. To test my understanding, would the approach generalise to matrix normal like this? Sorry for not deriving any math, this is very much monkey-see-monkey-do.

functions {
  // H' (K - K 1 1' K / 1'K1) H for symmetric K: the covariance of a GP
  // conditioned on 1'f = 0, compressed to the sum-to-zero subspace.
  // Helmert basis h_k = w_k (1_k, -k, 0), w_k = 1 / sqrt(k (k + 1)).
  // Two scalar loops, no J x J intermediates.
  matrix quad_form_helmert_schur(matrix K, data vector w) {
     // contents as above
  }
}

data {
  int<lower=2> J, M;  // rows and columns
  array[J] real x;
}

transformed data {
  vector[J - 1] w;
  for (k in 1:(J - 1)) w[k] = inv_sqrt(k * (k + 1.0));
  vector[M - 1] u;
  for (k in 1:(M - 1)) u[k] = inv_sqrt(k * (k + 1.0));
}

parameters {
  matrix[J - 1, M - 1] eta;
  real<lower=0> tau;
  real<lower=0> rho;
  corr_matrix[M] Omega;
}

transformed parameters {
  matrix[J - 1, M - 1] y;
  matrix[J, M] z;
  {
    matrix[J - 1, J - 1] L_C
      = cholesky_decompose(quad_form_helmert_schur(gp_exp_quad_cov(x, 1.0, rho), w));
    matrix[M - 1, M - 1] U_C 
      = cholesky_decompose(quad_form_helmert_schur(Omega, u))';
    y = tau * (L_C * eta) * U_C;
    z = sum_to_zero_constrain(y); 
  }
}