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:
- 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).
- Removing the
if(diff != 0)
logical check: another 10% slower - 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 solvingnnls_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.
x(i)
and two tiny tweaks to the logic, absolutely marginal gains stuff. Thanks! \$\endgroup\$