[ANN] AcceleratedRPCholseky.jl: A fast way to find near-optimal low-rank approximations of kernel matrices

Package Announcements
, ,
1

AcceleratedRPCholesky.jl for fast, near-optimal low-rank kernel matrix approximations

I’m happy to announce AcceleratedRPCholesky.jl, which is an implementation of the Accelerated Randomly Pivoted Cholesky algorithm from Epperly et al. (I can’t include links, but you can find it at arXiv:2410.03969)

This is my first Julia package published to the repo, so I welcome any feedback/suggestions anyone might have, as this is a learning experience for me.

Also, a disclaimer: this package was written with the help of AI, which was immensely helpful in setting up my repo and all the GitHub actions, producing the readme, and writing tests. It also wrote the first pass of the algorithm itself, however that implementation was not very good, so I largely had to re-write the implementation myself by hand.

The Algorithm

  • Does a partial pivoted cholesky decomposition where the pivots are chosen randomly, proportional to the diagonal of the residual matrix
  • Generates a block of pivots to do all at once to leverage faster computer matrix multiplication
  • Subsamples the block of pivots using rejection sampling to ensure that the distribution of selected pivots is the same as were do you select and process pivots one at a time
  • Produces near-optimal approximations (as compared to truncated SVD/eigen decompositions) for moderately sized ranks
  • Read the original paper (which I had nothing to do with) for more details

What You Can Do With It

  • The algorithm produces a matrix B from a matrix A and a kernel function k such that k(A[i,:],A[j,:]) ≈ B[i,:] ⋅ B[j,:] where the number of columns of B is many fewer than the number of rows in A.
  • This lets your write many “kernelized” algorithms (e.g. Kernel Ridge Regression) in their linear form but operating on B.
  • Because B is much smaller than the naïve n × n size that would typically be required to express such algorithms, making it feasible to compute solutions for problems that would otherwise be intractable due to the number of training samples

Implementation Details

  • The algorithm is linear in memory and time with respect to the number of data points for a fixed rank and block size
    • This is kind of amazing, since it means you don’t actually have to compute the vast majority of kernel function values and yet can still obtain near-optimal low-rank approximations
  • This particular implementation does all the main update computations in place in a pre-allocated output matrix
  • Block sampling and subsampling is also all done in place in pre-allocated buffers
    • As long as your kernel function doesn’t allocate, an iteration of the algorithm will not allocate any new working memory
  • All of the linear algebra operations use Julia’s mul! and will leverage BLAS as appropriate to actually do the operations if possible (including dispatching to syrk! where appropriate for cholesky update steps)
    • The core operations are all coded untyped and will work with whatever numeric type your kernel function returns, including things like Float32s and custom numeric types
    • Runtime can improve if you use an optimized BLAS library for your platform (e.g. use AppleAccelerate.jl on Apple M series processors to take advantage of the matrix coprocessor)
  • You don’t have to generate (or allocate memory for) the whole kernel matrix. The algorithm will compute kernel values on demand, and typically each kernel value that is computed is only computed once (aside from some additional kernel evaluations required for block sub-sampling, which is proportional to block-size and rank, and does not scale whatsoever with number of data points)
    • Because of this, an inefficient kernel that doesn’t leverage simd, allocates unnecessary working space, etc. can be a big bottleneck. Use high-quality implementations from something like KernelFunctions.jl
  • The input doesn’t need to be a matrix, you just need to be able to index into it rectangularly, and it needs to support views. In particular, you can pass a DataFrame from DataFrames.jl into the data argument directly
  • The algorithm will terminate early (before it reaches rank) if kernel matrix has a lower effective rank though an absolute test controlled by the atol keyword arg, or if the trace mass in the residual matrix falls below rtol of the trace mass of the true kernel matrix.

What about RPCholesky.jl?

  • There is an existing package that implements the non-accelerate versions of the algorithm, including the block version
  • RPCholesky.jl doesn’t implement the accelerated version with block sub-sampling that AcceleratedRPCholesky.jl does
  • AcceleratedRPCholesky.jl only implements the sub-sampling version, if you don’t want that sub-sampling feature, use the existing package
  • The implementation details are different
    • RPCholesky.jl works a the pre-formed kernel matrix, not on the data and the kernel function. If you already have the kernel matrix but not the original data and/or kernel function, this interface will be easier to use
    • Since the calling interface is fundamentally different, I decided to write my own package when implementing the accelerated version
    • It has a verbose mode that gives additional debug information
    • It is written in a style to be much closer to the mathematical formulation of the algorithm, which makes it simpler and easier to follow/validate for correctness at the cost of introducing unnecessary intermediate allocations