Efficient Kaczmarz

Unlike many of the other solvers provided by RegularizedLeastSquares.jl, the Kaczmarz method does not utilize a matrix-vector product with the operator $\mathbf{A}$ nor the normal operator $\mathbf{A*A}$. Instead, it uses the rows of $\mathbf{A}$ to update the solution iteratively. Efficient Kaczmarz implementation therefore require very efficient dot products with the rows of $\mathbf{A}$. In RegularizedLeastSquares.jl, this is achieved with the dot_with_matrix_row function.

using RegularizedLeastSquares
A = randn(256, 256)
x = randn(256)
b = A*x;

The dot_with_matrix_row function calculates the dot product between a row of A and the current approximate solution of x:

row = 1
isapprox(RegularizedLeastSquares.dot_with_matrix_row(A, x, row), sum(A[row, :] .* x))
true

Since in Julia, dense arrays are stored in column-major order, such a row-based operation is quite inefficient. A workaround is to transpose the matrix then pass it to a Kaczmarz solver.

At = collect(transpose(A))
A_eff = transpose(At)
256×256 transpose(::Matrix{Float64}) with eltype Float64:
 -1.31939   -0.668121  -0.325437   …  -0.697852    1.65017   -2.85298
  1.1136     0.304449   0.0496443      1.80761     0.762572   0.630775
 -0.560789  -1.21529    0.240076       1.43571     0.825138   0.335362
  0.905591  -0.710701   0.212489      -0.363416   -0.905077  -0.698773
 -0.698746   0.595845  -1.13295       -0.171527    1.14178   -0.166611
 -0.719597  -0.712582   0.588626   …   0.190637   -0.195595   0.377398
  0.79238    0.517067  -0.0127144      0.371449   -0.683738   1.28015
  1.42655   -1.01096   -0.416147       0.336105    0.496571  -0.843958
 -1.18432   -0.173788   1.23993       -1.58768     0.404443  -1.61528
  0.255828  -0.434313   0.610202      -1.07511     0.170909   0.132338
  ⋮                                ⋱                          ⋮
  0.293314   1.58986    1.16769        0.837321    0.195783   0.529521
  1.52816   -1.57829   -0.793271       0.652954    0.533092   0.376558
  1.42038    0.599648  -0.277254      -0.340933    0.909792  -0.0615635
  0.185308   1.3955    -1.30105    …  -0.258308    0.63135    0.551557
 -0.333761   0.303626   0.729463      -0.0197115  -0.285008  -0.387721
  0.959776  -0.391049   1.71564       -0.553084    1.38939    1.75479
  1.06114    0.70267    0.372533       2.229       0.848709  -0.780591
  1.6332     1.23879   -0.838119       0.28662     0.292435   0.715708
 -1.15939    0.964324   0.264709   …  -2.0395     -1.39512   -0.829786

Note that the transpose function can return a lazy transpose object, so we first collect the transpose into a dense matrix. Then we transpose it again to get the efficient representation of the matrix.

We can compare the performance using the BenchmarkTools.jl package. First for the original matrix:

using BenchmarkTools
solver = createLinearSolver(Kaczmarz, A; reg = L2Regularization(0.0001), iterations=100)
@benchmark solve!(solver, b) samples = 100
BenchmarkTools.Trial: 100 samples with 1 evaluation per sample.
 Range (minmax):  28.034 ms 28.500 ms   GC (min … max): 0.00% … 0.00%
 Time  (median):     28.249 ms                GC (median):    0.00%
 Time  (mean ± σ):   28.259 ms ± 109.105 μs   GC (mean ± σ):  0.00% ± 0.00%

                            ▂▂█    ▂▄ ▂▂                        
  ▄▄▁▄▆▆▁▄▄▄▄▄▆▄█▄▁▁█▄▄▆██▄████▄▁██▆███▁▁▄▁▆▁▆▆▆▄█▁▁▁▄▁▆█▄▁▄ ▄
  28 ms           Histogram: frequency by time         28.5 ms <

 Memory estimate: 17.31 KiB, allocs estimate: 505.

And then for the efficient matrix:

solver_eff = createLinearSolver(Kaczmarz, A_eff; reg = L2Regularization(0.0001), iterations=100)
@benchmark solve!(solver_eff, b) samples = 100
BenchmarkTools.Trial: 100 samples with 1 evaluation per sample.
 Range (minmax):  1.480 ms 1.684 ms   GC (min … max): 0.00% … 0.00%
 Time  (median):     1.491 ms               GC (median):    0.00%
 Time  (mean ± σ):   1.499 ms ± 30.446 μs   GC (mean ± σ):  0.00% ± 0.00%

   ▅▇█                                                     
  █████▅▆▃▄▃▁▁▃▃▁▁▃▁▃▁▁▃▁▁▁▃▃▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▃▁▁▁▃ ▃
  1.48 ms        Histogram: frequency by time        1.64 ms <

 Memory estimate: 17.34 KiB, allocs estimate: 507.

We can also combine the efficient matrix with a weighting matrix, as is shown in the Weighting example.

Custom operators need to implement the dot_with_matrix_row function to be used with the Kaczmarz solver. Ideally, such an implementation is allocation free.


This page was generated using Literate.jl.