Add sum_to_zero_matrix

[spinkney]

Description

We can add a sum to zero constraint on a matrix where the rows and the columns both sum to zero. This is useful for random effect models where you have the random effect interaction coming from a matrix variate normal. In addition, this is useful to specify "cost" type matrices for optimal transport problems.

I've put up the constraining transform, the free version, and the derivative in c++ at https://godbolt.org/z/qnW7xTeP7.

[jachymb]

This would also be useful for identifying the beta parameter of categorical_logit_glm.

[WardBrian]

@spinkney if you have the pseudocode that produced this it might be helpful.

[spinkney]

MatrixXd sum_to_zero_matrix_constrain(const MatrixXd& x) {
    int N = x.rows();    
    int M = x.cols();    
    int s = std::max(N, M);

    MatrixXd Z(N + 1, M + 1);
    VectorXd beta(N);
    VectorXd a(s);
    VectorXd b(s);
    for (int i = 0; i < s; ++i) {
        a(i) = 1.0 / std::sqrt((i + 1.0) * (i + 2.0));
        b(i) = (i + 1.0) * a(i);
        if (i < N) {
            Z(i, M) = 0.0;
            beta(i) = 0.0;
        }
    }

    for (int j = M - 1; j >= 0; --j) {
        double ax_previous = 0.0;  
        Z(N, j) = 0.0;
        for (int i = N - 1; i >= 0; --i) {
            double alpha = b(j) * (b(i) * x(i, j) - ax_previous);

            Z(i, j) = alpha - beta(i);
            beta(i) += a(j) * (b(i) * x(i, j) - ax_previous);

            Z(N, j) -= Z(i, j);
            Z(i, M) -= Z(i, j);

            ax_previous += a(i) * x(i, j);  
        }
        Z(N, M) -= Z(N, j);
    }
    return Z;
}

No, beta is needed. The critical piece is that beta augments Z(i, j) and for each i,j the code accumulates the +1 dimension element for both the row and column.

   Z(i, j) = alpha - beta(i);
   beta(i) += a(j) * (b(i) * x(i, j) - ax_previous);

The +1 is at

  Z(N, j) -= Z(i, j);
  Z(i, M) -= Z(i, j);

I don't have pseudocode, I just wrote it up from that gnarly algebra.

[WardBrian]

I see now that it is used, I wasn't thinking about the multiple iterations in J

Weirdly, though, my copy of the code that deleted it entirely still summed to 0 along each row and column. I assume leaving it out would lose some of the nice properties of the transform, but it's interesting to me that it even still worked.

The lack of pseudocode makes writing the derivative in terms of adjoints a bit trickier. If you can help me with that it would certainly speed things up!

[SteveBronder]

I think the first thing we can do is just make a PR for prim that has sum_to_zero_matrix_constrain. @spinkney do you have the associated free function as well? Once that works we can just sit down by hand and work out the reverse pass for a specialization in rev. Since a, b, and beta will be doubles it actually shouldn't be that bad. It would be the last loop in reverse order with all the lines that accumulate in Z with x on the right hand side. We'd have to store the values of ax_previous during the forward pass as well.

[WardBrian]

@SteveBronder I have the prim implementations in https://github.com/stan-dev/math/tree/sum-to-zero-matrix

[SteveBronder]

Nice! I'll try to work out the reverse pass this week or next

[spinkney]

I think the first thing we can do is just make a PR for prim that has sum_to_zero_matrix_constrain. @spinkney do you have the associated free function as well? Once that works we can just sit down by hand and work out the reverse pass for a specialization in rev. Since a, b, and beta will be doubles it actually shouldn't be that bad. It would be the last loop in reverse order with all the lines that accumulate in Z with x on the right hand side. We'd have to store the values of ax_previous during the forward pass as well.

The free function is also in the godbolt. The derivative is there too but not the adjoint specialty.