CoLA: Exploiting Compositional Structure for Automatic and Efficient Numerical Linear Algebra

Summary

Similarly to automatic differentiation which enabled to shift the focus from computing derivatives to deriving new algorithms, the authors propose to automate large scale linear algebra, with "structure-aware" linear algebra. Various algorithms leverage problem structure to expedite the evaluation of operations on linear operators. This toolbox aims to bridge the gap, allowing individuals without extensive knowledge of algorithms and tricks to still benefit from them. The successful accomplishment of this paper's objective could have a fair impact on the machine learning community by automating a significant portion of model optimization.

Strengths

This article is well written. The problem is well posed, and the proposed solution is convincing. Importantly, the library's well-thought-out and modular structure is crucial for its growth and positive impact on the community.

Weaknesses

The number of dispatch rule is relatively small. The authors mention 70 dispatch rules on `l. 224`, but in appendix A only a dozen such rules are shown (perhaps we reach 70 dispatch rules by counting all the possible combination?). If I had to write a new model, I'm wondering if (1) I would use this framework, with the (arguably small) associated overhead of using a new library or (2) simply scroll through appendix A to see what are the rules relevant for my implementation.

Questions

**Porting to Julia.** It is a nice use case of multiple dispatch, and I like the functional paradigm used. Are there any plans to port it in Julia? There already exists in the `LinearAlgebra` module "structure aware" types (`Symmetric`, `Hermitian`), and the multiple dispatch-functional paradigm are built in, so I believed it would certainly meet its public there. **Using SVRG.** `l. 245` and `248`, the complexity of SVRG seems bigger than Lanczos (at least in $\kappa$; what is $M$?). A reference for the convergence rate of CG and Lanczos would be welcomed. Why is an accelerated version of SVRG not considered, to speed up convergence and further improve the runtime performances? The cyanure toolbox holds some [benchmarks](http://thoth.inrialpes.fr/people/mairal/cyanure/benchmarks.html) along with some [references](http://thoth.inrialpes.fr/people/mairal/cyanure/references.html). Is it because of the memory limitations? **Arnoldi Iterations.** In Fig. 3, it is shown that `Arpack` performs better than the Python implementation of the Arnoldi operations. Given the modularity of your approach, wouldn't it be worth to add a dispatch rule which specifically uses `Arpack`? **Typos.** * `l. 82` : matrix vector **multiples** ? * Fig. 1 could benefit from having one legend for all 3 subplots. Contrary to what is asserted in the paper's checklist, there are no error bars. Given the moderate runtime ($10^2$s), using more than 3 repetitions (e.g 10) with error bar is necessary. * `245`: missing parenthesis

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

4 excellent

Presentation

4 excellent

Contribution

2 fair

Limitations

The article has no Limitations section but seems compared extensively to other libraries in sec. 3.4.

Summary

This work proposes a new library for solving linear algebra problems involving structured matrices. The library incorporates highly efficient linear algebra kernels tailored to handle specific types of structured matrices. Moreover, it uses compositional rules to tackle problems involving matrices with composition structures. By employing these rules, the library eliminates the need for manual implementation of many efficient algorithms applicable to such matrices.

Strengths

1. Designing efficient numerical linear algebra libraries is important to the machine learning community. 2. This work is well-written and the presentation is clear.

Weaknesses

Overall I believe the novelty of the work is very limited. The numerical linear algebra algorithms discussed in the paper do not appear to be novel, and the proposed library essentially consolidates these existing algorithms. While this library may be the first to utilize compositional rules for automated algorithmic development, the rules themselves are straightforward and uncomplicated, lacking significant novelty. In summary, this work represents a combination of various engineering efforts, but its scientific contributions are mostly incremental.

Questions

n/a

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

1 poor

Limitations

n/a

Summary

In this paper, the authors introduce CoLA, a library which streamlines the use of various linear algebra routines within frameworks of relevance for machine learning applications. The library revolves around the LinearOperator object, and provides numerous implementations of various numerical algorithms involving operations with such object, from inversion (system solution), to eigenvalue computation, to operators manipulation. Most notable features of the library are: - Structure-awareness and automatic dispatch: the CoLA framework is able to leverage information regarding the relevant structure of the LinearOperator considered (such as positive-definiteness, symmetry, sparsity, composition of operators), so to automatically identify and utilise the most apt algorithm specialisation for the target operation - Integration with existing frameworks: CoLA can interface with both JAX and PyTorch codebase, supports both GPU and TPU acceleration, and has some support for low-precision arithmetic and complex numbers - Autograd for iterative routines: CoLA implements some relevant iterative routines for linear algebra, and defines some ad-hoc rules for efficiently performing automatic differentiation through them The main reported results pertain the performance comparison of the CoLA-implemented routines versus alternative implementations. Overall, CoLA performance is shown to be comparable to that of the baselines; moreover, when there is gain to be had from leveraging the structure of the underlying linear operator, CoLA is shown to be able to do this effectively.

Strengths

- The authors propose a very useful framework. Most noticeably, it automates away the need for manually tuning the method choice depending on the properties of the operator, when numerical linear algebra algorithms are involved. Moreover, it effectively leverages some clever design choices (such as multiple dispatch). Overall, it can become a valid tool for numerous applications in the field of ML (and other fields as well) - The authors propose an interesting solution to efficiently performing autograd on iterative procedures, which are ubiquitous in linear algebra applications - The paper is reasonably well-written. Even though the breadth of applications considered in their work is indeed rather large, the authors still manage to present the key advantages of their framework in a clean manner, without being dispersive

Weaknesses

- The main weakness I see, is the lack of actual “novelty” in the work being proposed - at least in the classical sense of the word. Apart from the autograd rule in Appendix B2, in fact, the various methods proposed and implementation choices are not new. This notwithstanding, the main goal of the project consists in collecting available linear algebra routines into a unified, ready-to-use, efficient library which can be easily encapsulated within existing ML frameworks, and as such it is still valuable to the research community

Questions

Overall I’m quite satisfied with the paper. One minor doubt / curiosity I still have is: - When comparing CoLA with other existing baselines in Fig3, it seems like your implementation underperforms in both (a) PCA and (d) Spectral Clustering. You elaborate a bit for (d) in Appendix D4, but could you expand on this? In particular, is the difference in performance simply due to the lack of an optimised implementation on your side, or are there some structural causes? And what are the main implementation differences?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

4 excellent

Limitations

Limitations of CoLA are not explicitly commented upon, but at the same time the framework proposed seems very flexible and efficient (as showcased in the experiments). The main limitations are hence connected to what methods are readily implemented in the framework, rather than being structural.

Summary

This works presents CoLA, a framework for extending the linear algebra interface of modern numerical packages to take advantage of the structural properties of linear operators present in machine learning and other applications. By adding adaptive multi-type based dispatching CoLA adaptively exploits both dense and iterative methods to decrease the computational costs of performing a given operation. CoLA is also capable of exploiting the compositional structure that is present by combining different linear operator properties, creating a large collection of possible applications. Additionally, CoLA provides all these features in an extensible framework capable of backpropagation to ensure integration in modern machine learning and deep learning applications. Evaluations support the authors assertion that providing structural information during processing facilitates better performance on many benchmark applications.

Strengths

- The work clearly addresses an ongoing issue in many numerical linear algebra applications that require specialized structural operator structures to be individually implemented and exploited on a per-application basis. This process is not on tedious and time-consuming but also ripe with opportunities for error during the computation of the backward pass updates required for integration in a modern machine learning application. - Overall the writing and exposition of the problem, proposed solution and evaluations are clear and well articulated in the text. - The proposal dovetails naturally with the well-known LinearOperator interfaces that exist in Scipy and provide an easier route to define further extensions for developers to provide application-specific knowledge for further customization. - An interface for composing different linear operator properties for CoLA to exploit seems to be a novel and interesting extension over other existing implementations. - Performance results sufficiently demonstrate that having access to more structural information naturally supports more opportunities for application performance improvement while lowering the cost of the developer to exploit that structure through an intuitive and simple interface.

Weaknesses

- Though interesting the work presented may be a better fit in a venue that focuses on numerical methods and software. The intended audience is quite broad. - The core contributions and methods are relatively straightforward and are well-known to the numerical linear algebra community. This work seeks to make the process of exploiting the structural properties of linear operators easier for developers to use and integrate into a modern machine-learning application but it's not clear whether this contribution would have an appreciable impact on the developer community. Though this is my personal opinion it seems that problems encountered by developers are simple enough to be solved manually in most cases. - The evaluation section focuses on applications that exhibit basic structure but it's hard to see the value of the additional features mentioned: backprop and lower precision. - There seems to be a strong reliance on the appendix to fill in missing explanations due to size constraints. - Although the interface provides more flexibility regarding basic and/or compositional structure this still leaves a host of additional options to select the appropriate iterative approach to solve a sparse matrix.

Questions

- It probably doesn't help that one of the applications I know the most about, spectral clustering, is the application where CoLA seems to provide mixed results. Would closing the performance gap between CoLA and the sk (PyAMG) backend be a simple process of extending the multi-dispatch interface? - I found the added ability to efficiently backprop through iterative solve methods interesting, were evaluation results for this feature presented in any of the experiments?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

The limitations of the work are clear from the presentation and no issues require further acknowledgment, to the best of my knowledge.

Summary

The paper presents a library to automate the efficient execution of numerical linear algebra kernels commonly occurring in ML applications. The library recursively exploits compositional structure beyond versus standard packages which treat numerical matrix kernels as back boxes. The proposed framework provides memory efficient automatic, differentiation, low precision computation, GPU acceleration in both JAX and PyTorch.

Strengths

-) The quality of the manuscript is good. All concepts are communicated clearly and the paper is very well-written. -) The authors have put a lot of effort to provide automation for several important numerical kernels and structures. A long list of important applications are listed, and COLA can be of major significancy in several ML areas, essentially speeding up innovation. -) The numerical results indicate that COLA can be faster than baseline alternatives on a wide range of numerical tasks.

Weaknesses

-) One aspect I found confusing is the lack of information regarding the numerical algorithms featured in some of the results. For example, in Figure 1, it is not clear why the library is faster for the Kronecker problem. Likewise, the same is true for the Bi-Poisson problem. Can the underlying algorithms be found online? -) In similar spirit, it is not clear whether COLA is compared against state-of-the-art numerical approaches. For example, Multigrid is the de-factor choice for elliptic Poisson problems, is this what the authors list as 'iterative'? If not, what is the point of listing a comparison against a non-optimal iterative (or direct) algorithm? -) In Figure 3, scipy is competitive with COLA, if not better in some tasks. This furthers complicates the message of the paper. What is the main reason for publication? The fact that COLA includes a wide-range of solvers for composite tasks or that it can be faster or more memory efficient in general?

Questions

-) Is there any new numerical method involved in the library? My understanding is that every single numerical method used for linear systems, eigenvalue problems, etc, is known, and the main novelty of the paper is to wrap them together in an efficient manner for composite tasks. -) Judging COLA as a library requires some reasonable level of insight in the library itself. I understand that the authors must stay anonymous, but it is rather difficult to fully understand the benefits of COLA without looking at the code itself.

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

N/A