linear_operator - Structured Linear Algebra in PyTorchLinear Operator at PyTorch Conference 2022

Key Takeaways

hey I'm Max um and I'll be telling you about linear operator today a library for structured linear algebra built on Pi torch lots of other folks have contributed to linear operator our main brains here are Jeff and Jake who've worked on this for a while um what is linear operator it is a pytouch library for structural and algebra it's got a number of features that make working with structured tensors easy fast and efficient first linear operator abstracts tensor structure and encodes it in the linear operator object which allows it to automatically use specialized linear algebra routines under the hood without the user needing to reason about them um linear operators compose so potentially complex structured tensors are automatically generated from Primitives and the library implements both specialized exact linear algebra methods for structured tensors as well as fast approximate iterative methods and we'll choose the right one depending on what's needed great so how do you use it um it's easy linear operators look and largely behave like pytorch tensors in this example we create a linear operator a of size A Million by a million that is the sum of a low rank tensor and a diagonal tensor importantly the one by one million by one million Matrix is never actually instantiated yet we can do a bunch of interesting stuff with it namely we can use the standard torch apis on those linearator objects so we can sum across columns add a scalar perform Matrix Vector multiplication or use similar algebra methods such as eye gauge which means that we can use these objects just like we would any other tensors if we have existing code written in pi torch and so what enables us under the hood is that we register the operator's specialized methods to the standard torch namespace functions and then you know since the operator carries around its structure with it we can know what kind of simplifications we can do on the computational side so that allows us to perform computations much more efficiently both in terms of time and memory complexity for instance here in our example a from the previous slide if we call solve on this with some right hand side B then it would automatically use the Woodbury Matrix identity under the hood to compute the solvent o of n time rather than o of n cubed and for n equal a million here this means massive speed UPS not to mention not having to understand to that Matrix let's look at a couple of other examples the first one is on The Chronic obstructure tensors so if a is M by n and B is n by P then the chronicer product is NM by NP torch implements that but only eagerly so if you call torch cron you would instantiate the full tensor K here linear operatives that creates a chronic or product linear operator for lazy evaluation which comes with a number of efficient implementations for common linear algebra operations so in this example we look at a simple Matrix Vector product and we compute the speedups of the linear operator implementation with the naive approach of first forming the chronicle product and then Computing the NVM which is what you probably do if you are not a linear algebra expert the heat map here on the right shows the speed ups for some fixed m equal 10 for very small nnp we don't really see much of a speed up as expected but as nnp increase we do see speedups of more than an order of magnitude memory consumption improves similarly the second example I have here is so-called Monarch matrices these are matrices that are parameterized as the product of two block diagonal matrices plus some permutations they were proposed recently there's an icml paper from this year to make neural network training inference more efficient and so with linear operator implementing such algorithms here is just a few lines of code which illustrates how we can utilize this package to quickly and simply Implement new research ideas simply by composing the so finally let me show you how linear operators can serve as a drop in replacement for standard tensors in existing Pi torch code to exploit the structure so my example here is encoding block diagonal structure in the covariance Matrix of the multivariate normal basically we just make a block diagonlinear operator and stick that into the multivariate normal Constructor so now we can use the resulting distribution just as we would use any other one we can draw samples we can compute log props Etc and so here in this example we use a block dangle Matrix with num blocks blocks Each of which is size n by n so under the hood uh the code calls tortelling like colesky on The covariance Matrix which is a linear operator um so that call is dispatch to the specialized implementation for blocked angle matrices which essentially long story short computes a batch of num blocks classic decompositions and reassembles them um vis-a-vis the you know naive method of computing colesky on the full large Matrix so we get a nice quadratic speed up here we end up with o of num blocks rather than o of non-blocks cubed both for sampling and Computing log probes and I want to emphasize here that this is executing unmodified core Pi torch code the torch function dispatch and the structure awareness of linear operator allow us to get those speed ups and you can use that similarly in other existing code bases as a drop in replacement so what's the catch um well there's two main ones currently the first one is that linear operator is written in pure Python and so iFly involved tensors are very small then it may just be faster to use dense torque sensors the second one is that the main focus so far has been on symmetric positive definite matrices um they're used in kernel methods such as scouting processes where you know some of this this package has its Origins there's no real technical reason for this it's just a matter of doing it and it would be great to get some help from the community to do that that's it please check out our poster at the conference if you can I'd love to judge more and you can learn more about linear operator on its GitHub repo thanks