Coordinate Descent Non-negative Least Squares optimization

Question

In my fast implementation of Non-negative Matrix Factorization (Rcpp Machine Learning Library, RcppML), about 40% of the runtime is spent solving Non-negative Least Squares (NNLS) systems (the rest of the runtime is spent calculating the left and right sides of the NNLS systems, a and b, simple cross-products). I'd like to reduce NNLS runtime, if possible.

I'm using stochastic coordinate descent to solve least squares equations. Since I'm not expecting any theoretical improvements here, what computational approaches can I use to speed up this function? I am open to absolutely anything.

//[[Rcpp::export]]
Eigen::VectorXd nnls(const Eigen::MatrixXd& a, Eigen::VectorXd& b, double tol = 1e-8, size_t maxit = 100) {
  Eigen::VectorXd x = Eigen::VectorXd::Zero(b.size());
  tol *= b.size();
  while(maxit-- > 0) {
    double tol_ = 0;
    for (size_t i = 0; i < b.size(); ++i) {
      double diff = b(i) / a(i, i);
      if (-diff > x(i)) {
        if (x(i) != 0) {
          b -= a.col(i) * -x(i);
          tol_ = 1;
          x(i) = 0;
        }
      } else if (diff != 0) {
        x(i) += diff;
        b -= a.col(i) * diff;
        tol_ += std::abs(diff / x(i));
      }
    }
    if (tol_ < tol) break;
  }
  return x;
}

---

EDIT in response to Rainer P.

Thank you for taking the time to thoroughly understand the problem and for your great suggestions! I did a lot of benchmarking while incorporating your suggestions on a variety of test cases including different ranks and well-conditioned or random systems:

1. Your code "as is" was around 30% slower than my code "as is" for small ranks, but was nearly break-even for ranks > 100. For random systems (which rarely occur in NMF) your code was likely to break even with mine. For L1-regularized systems, your code was significantly slower (1.5-2x slower).
2. Removing the if(diff != 0) logical check: another 10% slower
3. Modularizing the code was helpful for benchmarking. nnls_iteration is the bottleneck, but not as much as it might seem. The percentage of time spent solving nnls_iteration/x += diff/nnls_mean_error for a 20-, 40-, or 100-variable system is 40/30/30%, 50/25/25%, and 70/15/15%, respectively.

I struggled to understand why your code was slightly slower for many test cases, until I realized that cases where x(i) is to be set to zero require first diff(i) = -x(i) and then x(i) += diff(i). Thus, there is assignment to a temporary, then floating point addition, instead of a logical check and assignment to zero if true.

Chunking up the solving step would make sense if ranks were very large, but typically an NMF model is rarely larger than a rank of 100 or so. Impressive work!

---

EDIT in response to G. Sliepen:

Thanks G. Sliepen for the thoughtful and sensible answer. I wish I could say your ideas inspired a faster version, but not so.

Logical branching/vectorization: Performance of vectorization and logical branching in coordinate descent NNLS is a bit noncanonical. NNLS can contain a significant proportion of zero-valued indices, particularly with L1 regularization. If I take out any one of the logical branches (e.g. diff == 0 or x(i) == 0) the processor does not know that a zero is a no-op, and so the code is executed regardless, taking time to do nothing. By my benchmarks, both logical checks provide speedups 10-80% in real life (k = 10 to 100, for real datasets). Ditto with vectorization, because if the gradient b is computed for all indices, it must then be corrected with a second calculation for indices where x == 0, which is highly inefficient. I'm not sure if some form of conditional vectorization could be used that would have an advantage over an element-wise loop.

Division: This is interesting, one I hadn't thought much about. But passing the const inverse diagonal vector of a to avoid division actually slowed down the function ~10% generally. This inspired an idea where a and b are "normalized" (similar to "anti-lop NNLS") such that the diagonal of a is all ones, with the caveat that x must then be "un-normalized". This was marginally faster for massively parallel NNLS where a is constant, but doesn't hold true for cases where a is variable (i.e. NMF with masking).

Warm-start initialization: Spot on, and I do this in my NMF implementation with great results.

Answer 1

If I understand your code correctly, you subtract the columns of the matrix a from vector b, each column multiplied by a different scalar diff. You also add diff to an element of vector x. Basically:

loop until convergence:
    for each column i:
        diff = ...
        b -= a.col(i) * diff
        x(i) += diff

As a first improvement, I suggest to simplify and modularize the code to make it more readable and to identify the bottleneck. You can get rid of if (-diff > x(i)) if you use diff = std::max(... instead. If you make diff a vector, you can separate the convergence check from the main loop. The result is more readable and maintainable, but we haven't improved performance yet. Here it is:

Eigen::VectorXd nnls(const Eigen::MatrixXd& a, Eigen::VectorXd& b, double tol = 1e-8, size_t maxit = 100) {
    Eigen::VectorXd x = Eigen::VectorXd::Zero(b.size());
    Eigen::VectorXd diff = Eigen::VectorXd::Zero(b.size());
    while(maxit-- > 0) {
        nnls_iteration(a, b, x, diff);
        x += diff;
        if (nnls_mean_error(x, diff) < tol) {
            break;
        }
    }
    return x;
}

void nnls_iteration(const Eigen::MatrixXd& a, Eigen::VectorXd& b, const Eigen::VectorXd& x, Eigen::VectorXd& diff) {
    for (size_t i = 0; i < diff.size(); i++) {
        diff(i) = std::max(b(i) / a(i, i), -x(i));
        if (diff(i) != 0) {          // optimization if zero is common
            b -= a.col(i) * diff(i); // critical inner loop
        }
    }
}

double nnls_mean_error(const Eigen::VectorXd& x, const Eigen::VectorXd& diff) {
    double error = 0;
    for (size_t = 0; i < diff.size(); i++) {
        if (x(i) != 0) {
            error += std::abs(diff(i) / x(i));
        }
    }
    return error / diff.size();
}

After modularization, it is obvious that nnls_iteration is the bottleneck and we cannot do much about it. We read each matrix element exactly once and this is a tight lower bound. The code in nnls_iteration is so simple that performance will be limited by memory bandwidth on most platforms. If the vector b is small enough to fit into L1 cache, the code is indeed optimal and cannot be improved further, assuming the compiler performs the vector-scale-and-add operation in a single pass over the vectors. With the Eigen library, this can be taken for granted (it's somewhere in the documentation).

A more interesting case arises when b is too large to fit into L1 cache. In this case, we perform two memory reads and one write per matrix element because we save and restore b once for each column of a. We can bring it back to roughly one read by processing the matrix in blocks, in this order (shown for 6 chunks of rows):

| 1|        12    |
| 2| 3|      13   |
|  4  | 5|    14  |
|   6    | 7|  15 |
|    8      | 9|16|
|     10       |11|

To make this work, we use the basic nnls_iteration for the on-diagonal blocks but we only update some rows of the b vector. We then apply the resulting diff vector to the bottom rows of b using a simplified nnls_iteration_off_diagonal before moving to the next on-diagonal block. In the end, we process the upper triangular matrix that we ignored until now and update the top rows of b that we didn't need for the on-diagonal blocks.

void nnls_iteration_off_diagonal(const Eigen::MatrixXd& a, Eigen::VectorXd& b, const Eigen::VectorXd& diff) {
    for (size_t i = 0; i < diff.size(); i++) {
        if (diff(i) != 0) {          // optimization if zero is common
            b -= a.col(i) * diff(i); // critical inner loop
        }
    }
}

void nnls_iteration_chunked(const Eigen::MatrixXd& a, Eigen::VectorXd& b, const Eigen::VectorXd& x, Eigen::VectorXd& diff) {
    size_t total_size = diff.size();
    size_t chunk_size = 100;
    // Process the below-diagonal block and then the on-diagonal block on the same row.
    for (size_t begin = 0; begin < total_size; begin += chunk_size) {
        size_t end = std::min(total_size, begin + chunk_size);
        auto left = Eigen::seq(0, begin - 1);
        auto middle = Eigen::seq(begin, end - 1);
        nnls_iteration_off_diagonal(a(middle, left), b(middle), diff(left));
        nnls_iteration(a(middle, middle), b(middle), x(middle), diff(middle));
    }
    // Process above-diagonal blocks after all on-diagonal blocks have been processed.
    for (size_t begin = 0; begin < total_size; begin += chunk_size) {
        size_t end = std::min(total_size, begin + chunk_size);
        auto middle = Eigen::seq(begin, end - 1);
        auto right = Eigen::seq(end, total_size - 1);
        nnls_iteration_off_diagonal(a(middle, right), b(middle), diff(right));
    }
}

Answer 2

Avoid branches

Branches in code, like those caused by if-statements, can sometimes be expensive on modern processors. Consider the case of diff == 0: all the code inside the if (diff != 0) block can be executed even if diff = 0, and it will be a no-op. The same is maybe even true for x(i) == 0?

As you mentioned in your edit, zeroes happen often, and in that case it is indeed worth it adding explicit checks to avoid doing unnecessary computations.

Divisions are expensive

Divisions are one of the more expensive operations for CPUs. Consider b(i) / a(i, i): if the matrix a doesn't change between calls to nnls(), it might be worthwile to pre-calculate the reciprocal of the diagonal of that matrix and pass it as an extra argument, so you can rewrite the aforementioned expression to b(i) * reciprocal_diagonal_of_a(i).

As mentioned in your edit, this slowed down your code somewhat, and perhaps as Rainer P. mentioned in his answer, memory bandwidth might be a limit, and by adding yet another vector to read in, this will slow things down.

Do more with whole vectors at a time

Instead of doing element-by-element operations in a loop, consider that Eigen has a lot of functions that can do element-wise operations on whole vectors or even matrices at a time. This simplifies the code, and might give a performance uplift. Consider:

while(maxit-- > 0) {
  auto diff = b / a.diagonal();                      // now a vector
  auto non_negative = (-diff < x).cast<double>(); // vector of 1s and 0s
  ...
  x += diff;
  x *= non_negative;                                 // resets negative components to zero
  double tol_ = (!non_negative + diff / x).sum();
  ...
}

(I left the hard part as an excercise for the reader.)

Use an estimate for x

You start with the vector x being all zeroes, but if you have some knowledge of what x should look like, it might be helpful to pass an estimate of x as a parameter. This might cause your algorithm to converge faster. Consider that if subsequent calls to nnls() would result in only slight differences in the return value, you can just pass the return value of the previous call to nnls() to the next one.