Scale By The Bay 2021 : Ryan Orendorff, Functional Programming + Dependent Types ≡ Verified Linear..
Hi everyone. Yeah, as mentioned, my name is Ryan and the talk I'll be giving is called functional programming plus dependent types is equivalent to verified linear algebra. So, a little bit about myself, I guess on top of what already has been said, I'm a research scientist at Facebook Reality Labs. By day, I do human-computer interaction research mostly at the moment and by night, I'm really passionate about theorem proving, programming languages, and dependently typed that languages such as Agda in particular. Um and you can find my repos and other talks at my GitHub page. I want to add a quick disclaimer that all the work that's presented here is done on my own personal time and equipment and is not related to any other employer I've had, past or present, and isn't related to any work at any of those employers. So, the goal of this talk is to be able to construct matrices in a functional way, that's part of the title, but we're going to do it in a way that we have guarantees that we've done the correct thing when we're constructing these matrices. And that we haven't used them improperly
So, what are a few of the ways that things can go wrong? One of the common ways that things can go wrong when you're dealing with linear algebra is the fact that the size of the matrix that you're dealing with is not congruent, it doesn't make sense. So, in this case, say we have a matrix defined in Haskell as just a list of list of A. And in that case, there's nothing that really specifies in the definition of list of list that the two lists, the N list that you pass along as the rows or say columns, depending on how you're thinking about it, of the matrix are actually the right size. So, you might have instead of a nice ordered table where it's a N by M matrix, you might have something where some of the columns are missing an element for no reason. You can also have in some programming languages improper data types. This is more common in say something like Python where if you use the very popular NumPy array, you can pass into NumPy basically anything and it will just cast it to the thing that it thinks it needs. So, in this case, you could make a NumPy array of an integer and also a string and it will just figure out that it needs to hold objects, um which is probably not what we're thinking of when we're thinking of linear algebra. And there's a few more surprising errors that we'll get into later
For this talk, we'll be using Agda specifically, so it'll be all Agda syntax. If you've used Haskell, it's very similar to Haskell and so you'll find yourself um very familiar in that territory, but ideally we'll be able to go through it such that the ideas work regardless of the language of your preference. So, we often think of a matrix as a list as a table of numbers. So, if we want to construct that, we can talk before as we said with the list of lists. So, in Agda, this would be defined as list of you can construct, say, a type of matrix of numbers using the constructor construct matrix of numbers, which takes in a list and list of A and produces that type. If you want to think about this in uh Haskell, it was you just could think of it as just a normal data constructor or you can think of the Agda syntax as maybe the GADT syntax structure for defining a data type. And in Python, there's not necessarily a canonical way to do this, but we're going to say you can use data classes and that data class just says that your class is really just a wrapper that holds that list of list of A. Now, uh given a constructor, this particular constructor, we can construct the following matrix, which is we just give it the numbers
It's 1 2 3 4 5 6, right? Easy peasy. Um and I want to just mention before we get too far any conventions that you might see here. So, I've introduced a few already, but basically anything that's a monotype with a capital is a type. Anything that's a fancy uh italics is a probably a matrix or a vector. M N P and Q are natural numbers and U V X and Y are vectors if you happen to see them around. So, what can we do with a matrix? We can apply a matrix or you can multiply it to a vector to produce a new vector. Um we can transpose a matrix. That means take all of the elements that are not on the diagonal of the matrix and flip them around
So, anything that's in that matrix position IJ is now matrix position JI. And then finally, we can combine two matrices through matrix matrix multiply. So, uh what does matrix vector multiplication look like? Well, really what that is is that's just taking a row of your matrix and multiplying it by a column of the input vector element wise and then storing that result. As we see here. So, that's one way to think about what a matrix vector multiply is. But potentially, another way to think about a matrix vector multiply is to think of M, the matrix, as a function from vectors of size three in this case to vectors of size two. Sometimes this may this function will be called a linear map um or a linear function or some other kinds of things, but uh in this case, we'll just think of it as just a function. So, what kind of functions, if we're trying to think of it in terms of functions instead of these numbers, can we generate as matrices? So, the most simple case of a matrix is the identity matrix
And what it is is it's basically ones along the diagonal of the matrix and zeros everywhere else. And what it does is that it takes in an input vector V and it returns it without any changes. Essentially, what it's doing is it's multiplying every element of the vector by one and then storing that in the same place that it came from. Now, we know what this looks like in terms of a functional sense. It's the identity function as well. The name is uh the same in both. And so, in order to generate an identity function that acts like the matrix of numbers that is the identity function, we really just have to take the input list in this case and return the list again, since we're encoding our vectors as lists at the moment. Now, say we have something slightly more complicated, the diagonal matrix
The diagonal matrix basically has on the diagonal a list of different numbers, not necessarily all ones, and everywhere else is zero. And what effectively this means is that you element-wise multiply the vector that represents those elements on the diagonal, which we'll call you, by the input vector V that we're multiplying by. So, a way that we can do this is we can use the uh matrix vector multiply that in this case is defined in a library, which just does says that we say U * V, and this is the matrix vector multiply. If you want to do it in terms of the more base case, and if you want to see how that's implemented, you could consider this to be a zip. So, basically, you're zipping with the multiplication operation across these two different lists. So, now that we have kind of a definition for some functional versions of a matrix, let's try to construct a functional matrix. We can construct it by basically wrapping a data type around this function that we're talking about. So, we're basically going to pass into our construct FM, which constructs a functional matrix for us, a forward function, which is the function that you would happen that corresponds to when you multiply a matrix by a vector
In this case, we're also going to include the transpose function, which is to say that it is the function that happens when you transpose the matrix and then multiply it by a vector. We'll need that later for other properties, um and it happens to be convenient to have, so we're storing them together. And these two pieces are all that you need to define a functional matrix. So, in order to construct that functional matrix, we just call our constructor construct FM on the list two list identities that we had before. Now, we want to get back to properties that we had before when we were talking about list of numbers. What we were talking about was a list of list of A, and then we were multiplying it by another list of A. We had matrix vector multiply, we had matrix transpose, and we had matrix matrix multiply. Well, these apply in the functional case as well
Matrix vector multiplication is simply taking out the forward function in our pair that we've defined as the forward and transpose functions, and applying it to the vector that we have. Now, the transpose is also something that's actually quite easy to define. What the transpose really means is it means taking these two functions that we have and switching them. So, now instead of having our forward function being applied, we're replacing with our transpose function, and instead of our transpose function, we're replacing with our forward function. And matrix vector matrix matrix multiply turns out to be something that's quite neat. It turns out to be function composition. So, when you perform a matrix matrix multiply, what this is effectively doing is that it is taking the two forward functions and it's composing them together, and then the two transpose functions and composing them together, and then that defines your new matrix. Now, one thing to note here that's just to make sure that all the numbers match up correctly, uh all the shapes match up correctly, is that the transpose function actually is flipped in its composition
So, instead of saying F1 uh F2 composed with F1 as you have in the forward, in the transpose case, it's one followed by two. And the reason why is because in linear algebra, when you transpose a matrix matrix multiply, the two operands around the matrix matrix multiply flip as well, and then you apply transposes to both sides. To both operators. So, uh we're going to be doing intuition checks throughout this talk just to make sure that we can define our matrices in terms of something that would be workable in terms of just a list of numbers or table of numbers. So, what we want to be able to show is that for our identity function, that it's equivalent to some other function such that when you uh that is only made up of multiplies and adds, that gives you the same answer. So, in this case, our identity function is you take the input, you return the output. But another way you could write this is that you could say that you would replicate for the given input vector, you would make a vector of si- of all ones and then element-wise multiply those together. That will also return to you the same input vector that you had, but now it's defined only in terms of multiplications and additions
Where the replicate really just happens to help with making sure the size is correct. So, is our functional matrix definition correct by construction? Well, uh one thing that we could do, since we can define any list of list of A to list of A functions, is that we could always return the empty list. And we can construct a function uh matrix this way. But, as you might guess, we can't really write this in terms of a list of numbers, because essentially what we're doing is we're taking the whatever the input vector is, and we're just discarding it and always returning the empty list. So, clearly something is missing. What do we need in order to get there? Well, what we need to get to this first set of constraints in order to get to correct by construction, we need to make sure that the sizes match. We need to make sure that whatever you pass in is either after the matrix has done its thing, has produced an output of the correct size. And what you can do in Agda is you can define instead of a list, you can define something called vector
And vector includes in the type the actual length of the list itself. This is sometimes called a dependent type, and is probably the most famous dependent type, the case of vectors. Um, but basically what this says is, say a back over natural numbers of size three means all vectors of natural numbers have to be exactly size three. You can't put a size four in there, you can't put a size two in there. That's not part of the type. If you try to do this, if you try to say like reassign a vector size three to a vector size two, Agda will complain to you. And specifically, it will say that it can't figure out that three is equal to two, which is a very nice thing that it knows. Um, and so it will basically prevent us from doing any cases where the shapes are end up mismatching from what we expected
So, now we can construct construct a functional matrix in a different way. We can now use the vector definition, and just replace all cases of list with cases of vector. And we can do what we we before with the identity function. And in fact, we don't even need the a different identity function to identi- create the identity matrix because it is the same function. You're just always taking an input and returning that output. In this case, that type is instead of a list is a vector of some size, but otherwise the identity function is the same. What this does is that it helps convert a what would otherwise be a runtime check into a compile time check. So basically before, if you made any errors when you were constructing these matrices functionally and accidentally made the wrong size, you would probably want to write a test suite such that you could test that the shapes on the outputs and the inputs were all made sense for your given application
Here, you can convert that into a compile time check and guarantee that you can't make an error in terms of the shape of the input and output vectors. So, are we at correct by construction based on the fact that we can write our our matrix as a list of numbers? Well, say instead of having a list of numbers, we have a list of cards, the suite of the cards. And so, we have diamonds, clubs, um and the others. So, in order to construct this, say we'd basically have our forward and uh transpose functions are basically just defined as replicating in this case the spade and the heart as a vector for whatever input we have. If we'd wanted to try to convert this into a list of numbers or in some way to act in terms of a matrix, we would find that we would have a problem. And the essential problem is that a card suite doesn't have a definition for multiplication and addition, or at least not one that like it comes to mind easily. So, elements have to be able to be added and multiplied. Uh in order to do that, we basically need to say that matrices need to be defined over some sort of type class like thing
So, what we need to say is that the type A in our vector can no longer be just any A. It has to implement some certain properties in terms of matrix multiplication or and multiplication and addition. So, where to define this, we it's called a field is the algebraic structure. And a field has the property the following properties: addition, multiplication, a negation or inverse for the addition function, a inverse function for multiplication, and some identities for those particular operations. So, a zero identity element such that when you add the zero identity element to some other element, you always get that same element back. And you also have a multiplicative identity such that when you multiply some value by the multiplicative identity, you always get that value back. And with that, we're able to define matrices in terms of the field properties that we need. So, in this case, this looks exactly the same as basically what we had before
The only difference here is we have a constraint that says that the type of A must implement the elements of field. Um if you're familiar with Haskell type classes, it's basically similar, but just in a different syntax in this case. Or traits or other types of things like that. And so, from this, we can generate the same identity matrix we've been talking about before. And in fact, nothing needs to change because since we're changing the elements that operate that the vector can be, the identity function still doesn't care what those elements are. It will just always return the right same thing back. So, we've added another check. Now, the question is, have we actually gotten to our correct by construction linear algebra? Have we remade basically a a table of numbers, but in a functional form? Well, one could also, and this has happened in prior examples, basically take some sort of uh input vector and always return a constant output
It will be the right shape in this case. It would be a vector of ones of whatever the size shape that you need in order to make the inputs and outputs match, but you'd always return ones. So, you can't do this with only multiplication and addition. Essentially, what you're doing here is that since you're producing a constant output, you're making some sort of decision to throw away your input and do something else entirely. Linear functions or linear maps don't have the ability to kind of just ignore the thing that's put into them. And so, clearly we're missing something. So, the properties we're missing in order to tackle some of this is basically what defines a linear function. There are two properties for linear function
Linearity, which is defined as multiple if you take a function and you multiply it by U plus V, you should be able to basically propagate that function over the addition. So, you can do M of U plus M of V instead. And similarly, there's another property called homogeneity, which says that if we apply a function to some constant multiplied by a vector, then we should be able to pull out the constant instead of doing it inside before we apply the function. Um when we had our prior definition, the problem was is that with the replicative one is that you don't get the right answers. When you do linearity, you find that if you say evaluate the left-hand side where the addition is done before the application of the function, you always return the the vector of ones. If you do it afterwards, you get the vector of twos. So, they're not equal. And if for homogeneity, it turns out that if you do the same thing, you either get the vector of ones on the left-hand case, or in the right-hand case where you pull out the constant, you get a vector of all that constant value
So, you don't quite There's That constraint is missing. So, in order to do this, we'll have to define something called a linear function. A linear function basically allows us to codify these concepts. So, what we'll need is a function like we had before, along with a proof that the linearity property holds, which is essentially what this says for all input vectors U and V that the linearity property holds, and some also prove that the homogeneity property holds for all inputs constant C and all vectors V that we can basically prove that homogeneity holds. And for this we'll define a little bit of a helper function that just says that we can extract out the uh the actual forward function that we want to use or the function in this case and that way we can ignore the proofs and just apply the thing. So, how do proofs work? This This triple equal sign. So, for a triple equals, what happens is that essentially you can take some base proofs or maybe some axiomatic proofs in this case for the case of a field and use it in order to prove other more larger proofs. So, for here if we want to say zero B plus zero times one is equivalent to B, what we can do is we can do it essentially the same thing we would do when we're writing it out on paper
We begin with the start of our definition on the left-hand side and then we'll start rewriting pieces. There's a property of a field that says multiply by one is the same. Um so, that is actually part of the definition of a field as well. We can apply that uh function to our input and basically remove the times one. Which is uh then we could do for the plus zero, we can apply the proof that says that plus zero is equivalent to doing nothing as well. Um and note that basically the way that this is written out is you have a term and then a rewrite step in the triple equals and then in the angle brackets. And then finally we get to our answer that we've proven B and then the little square is just says that we're done. So, how does this work for our uh identity function? For the identity function, the function itself definition is simple
The proof actually for linearity and homogeneity are both also simple cuz what happens is that when you define this function, you have to prove that the identity of the addition is equivalent to the identity applied to each element each vector independently and then added together. Well, it turns out Agda knows what to do with this because it knows what the definition of identity is and it just applies it for you. And so then what you do is you put this definition refl in there which is stands for reflexivity which says that these terms are the same by definition or in shape. Homogeneity works the same. It's also a a reflexivity proof. So, it turns out that linear functions can also be composed which is nice because then what this means is that just like we can compose smaller functions to make more useful functions, we can compose linear functions to build up larger functions that also uh carry the proofs with them. So, in this case what we're doing is we're taking G and H and we're composing them together. And what this does is that it basically is the same as composing together the functions underneath
We need to generate a proof for linearity and homogeneity in this case. And in order to do that, what we do is that we basically say we need to prove this top property that's in a comment that F of G of U plus V is the same as if you take F of G applied to U and then add F and G applied to V. Well, turns out this proof is actually pretty easy and it it's nice then for the sense that it uh uses a prior proof. So, what we do here is we start off with our base on the left-hand side. Then we apply basically our rewrite step in order to uh get to our next piece which is to take that H and apply it to U plus V. Now, here the cong just means inside basically inside of uh this operation, we're going to apply a proof here. Um and we're going to basically apply the proof that we have from H itself which is that H is linear. So, we can apply the fact that H is linear which is what this piece is to break apart H into H of U plus H of V
And now we can apply the linearity of G to this property. Now, we don't need cong in this case because we're actually doing at the top level. We don't need to kind of dig in and put the proof where we want it to. And that leads us to our proof. We have a proof that linearity composes, and we have a proof that the proofs of linearity also compose. Now, you can do this for the homogeneity example as well, but I'm going to leave that as homework. So, now we can finally define a matrix laid up of linear functions. What a matrix of linear function is is you just replace those vectors, the vector of n to vector of m, with linear functions that do the same
And the identity function again turns out to be very similar or the same as it was before. So, are we there? Are we at correct by construction algebra? Well, there's one other problem and ideally final problem that we deal with have to deal with, which is the fact that since we bundled the forward and transpose functions together, we actually need to make sure that they correspond to each other. So, in this case we can generate a function a matrix that is the identity function on one case and the for the transpose is the diagonal function, which is not the transpose of the identity function. The identity function is the identity function. So, we need a statement there that says that we can basically contain and uh that these two functions correspond in some way such that we can guarantee that they're doing the right thing. And what this is is the inner product proof. So, what we're saying is that for all input vectors x and y, we should be able to say x and then inner product with m times y is the same as y inner product with m transpose x, where the inner product is defined as the sum of the vector element-wise vector multiplication of the two input vectors. So, with that with that final piece, we can actually define a matrix fully uh in a way that is correct by construction
So, we're done in this case. We have reached our final uh case. So, what we need in order to construct a matrix is the forward function, which is a linear function, a transpose function, which is another linear function, and then this proof p which says that for all X and Y inputs that the inner product holds as we had mentioned on the last slide. Um So, here in this case to say generate a uh matrix that's the identity function if we want to generate it in our correct by construction case, we apply similarly to what we did before but we have to do one extra step step. We have an identity function where we've proven that it's linear. So, we have the linear version of the identity function. It's not the same as the regular one because there's the proofs that of linearity and homogeneity that come along with it. We need the transpose function
And then all we need a proof that the two transpose that forward and transpose functions actually match. Which is what this proof is here. So, we want to prove that X trans uh inner product with the identity of Y is the same as Y uh with the inner product of the identity of X. And the proof of this uh turns out to be pretty simple as well. So, what happens is we start off the with the left-hand case again. We apply a rewrite rule and in this case the rewrite rule you'll notice there's nothing in those brackets. And that's because Agda is basically it can figure out and apply the identity function. Essentially, what you can think of is in those brackets, even though it's empty, is reflexivity
Now, what we can do is that say we have a proof and I'm just going to say that we have the proof that the inner product the two elements of the inner product commute, so we can switch them around. And then we can rewrite the identity function again in reverse to get back uh to the identity of X. So, now we have a fully correct by construction uh identity function matrix. So, what can we do with this? We can multiply the matrix by a vector, we can transpose it, and we can combine it. Um the different pieces here are basically as you would expect before. In order to apply it, we pay take out the forward linear function then we use the forward the linear function applicator, which is the little dot with an L on it. Um and then for transpose, we again flip the forward and transpose functions. In this case, let's say cuz sometimes people say adjoint instead
And we need to actually also provide another proof in order to generate this transpose matrix. But in this case, the transpose matrix is we use the proof that we had at the other matrix and we basically just need to flip the operands. So we have from our original matrix a proof of X A Y is Y F of X. And we need basically the case where the uh two left and right sides of the triple equal sign are switched around. Sim gives us that. So essentially, it's the same proof. Just flip the left and right sides. Um and now we can derive define matrix matrix multiply
This is the same as it was before. It's we're composing in the same order. So the forward function is the two linear functions composed one composed with two. And then on the transpose, we have to basically account for the fact that the top part is also transposed when you or the matrix itself, the operations need to be transposed. And then you need to provide another proof. But I'll leave this as more homework. So we've gained some nice benefits from uh this construction. We can define a performant functional version of matrix algebra
We can guarantee that our implementation is correct. And we can use equational reasoning to prove two different implementations are correct. Or we can use equational reasoning to prove that any linear algebra statement in this case is correct, which is convenient um if you're operating on both the functional algorithmic definition of something and on the linear algebra side. Um but there is a cost. For one file in the library that implements this idea, uh out of 213 lines of code, 24 are actually defining functions and like what the function does and stuff like that. And everything else is a type definition, a proof, or an import or control statement. So about 90% of the work, mm minus the import statements and stuff like that, goes towards basically generating a lot of these proofs. So let's see an example, algorithms using linear algebra
So um we gain some benefits from using functions directly. We can write out a model for a process that happens in a more direct manner than potentially a list of numbers. And you get speed and time benefits. So a matrix can be extremely inefficient versus the case where we have the function version, which could be very efficient. And a good example, a base example, would be the identity function, which otherwise requires O of N time to do all that one multiplies by B. And in this case, uh requires constant time because you just always return the input vector. So I'm going to list a particular application from my work when I was a graduate student. So I worked in a field called magnetic particle imaging at Berkeley
Uh and what magnetic particle imaging is that you basically inject an iron tracer into uh an animal and soon to be uh humans. They're working on human scanners at a startup that came out of our lab called Magnetic Insight. And track where that iron goes in order to be able to detect uh anything basically that's challenging with the blood, whether or not it's you have blockages or you could try to detect cancer or where track cancer's going and stuff like that. And the machine on the left here is basically what does that imaging. It figures out where the iron is and produces basically figures that out by gathering some voltages as it changes some magnetic fields. And you get the image on the right. Which in this case is a rat that's been injected with some iron tracer. Completely safe iron tracer
Um So uh we get a sizable improvement in the performance with matrix of this reconstruction when we do it in in terms of functions instead of as a actual written out matrix of numbers. The amount of space the matrices themselves take up a ton of space. Some of them aren't sparse. So we went from a space size of 150 GB for some images to do some recon- 3D reconstructions down to just the bytes that it took to store the function pointers Uh and do the transformations. The time also had a 30-fold improvement, and the fact that we're using functional concepts is of course excellent. So, how does this work? Basically, our what we're going to do is that we have a matrix M, and we're have it as a model of what happens to some iron that we have, and that will give us the signal that we get out on the scanner. So, for some given input of iron or whatever we think is the distribution of iron, we apply M to that, and we should get the voltages that we actually get out of the system. And what we're going to want to do is that given an actual value, like actual voltages we get off of the scanner, what we're going to want to do is compare that to some input guesses
We're just going to keep guessing basically over and over again until in a somewhat smart way until we basically find a case where the input of the iron that we put into the function gives us the values that we actually got on the scanner itself. So, in order to do this, basically, we do what's called gradient descent, and gradient descent is basically where you take a function, you evaluate it at a point or the derivative at a point and use it to move farther along to a more optimal position on that function. So, if we have, say, this convex bowl, which is nice, and we start off at this point, we're finding a derivative, that's what the upside-down triangle sub means, that's a direction to basically move in, and then we go slightly down and find a new point. The uh we do this with this step function, which is just the way to encode it in Agda. And the gradient descent part is just basically applying the step function over and over again. We're finding the direction to go on the function, stepping slightly in that direction, and then repeating. So, what happens if we have our step function and we want to try to say optimize it? Are there other ways to rewrite it? Well, the simple case, assuming we know nothing about M, is that we can actually just distribute it over. We can say M transpose uh distributes over that M multiplied by the iron or applied to the iron minus the V, which is the voltages on the scanner
So, we can generate a proof of this. The proof of this is essentially that um you have There is a distribution a distributive law that allows you to basically take a matrix that's applied to U minus V, and then you can distribute it over that, similar to the addition. And what's convenient about that is that we can rewrite our programs in a way that while we preserve the correctness. So, say I write a program, but it's not particularly efficient, and I want to prove that it's equivalent to another program while maintaining correctness. Here, I can say is that in this case the two step functions are the same, and say one has more performance over the other, then I can say that my whole program is still correct, but potentially I have algorithmic advantages from doing this transformation. You can think of other transformations if you say new what M was. So, we've achieved our goal. We've agreed the correct bank construction linear algebra
We've eliminated wrong size results bugs, eliminated non-linear function bugs, and eliminated incorrect function pairings for the forward and transpose functions. Um I just want to make a brief mention that basically these ideas have been I've implemented them, and a lot of people have implemented parts of them um in three different ways. So, the base way is in Python is basically just the definition with functions with no checks on sizes. I wrote a very small and like probably not terribly well-written library while I was just investigating this called convex that does this with in Haskell, but with size vectors. And then finally, linear functions, which is the library that this is closely with this talk called functional linear algebra. And what I found is that not every library is a blast to use. And these three different appro- approaches kind of give you different spaces spots on this space of ease of use and correctness to trade off on. On the Python and basic side with just lists, you have really easy to use functions, but your connected correctness is pretty low
And you actually have to write a lot of tests to make sure that everything works. The Haskell size type variant basically increases your correctness by actually quite a bit while not deterring from the ease of use all that much. And then way off in the corner is basically this idea. So when you're doing these functional linear algebra proof by correct by construction with proofs, it's very difficult to generate stuff because you spend a lot of time in proof land, but the correctness is extremely high. So with that, that's the end of my talk. You can find more on GitHub including the links to this talk and all the definitions without the homework in them. This is a literate Agda file so you can actually see how it was written and it all type checks. And if you have any questions, I will be around for maybe some after this and otherwise on special chat.