Neural Network Research with Wolfram Language

Ian Wright presents DB nets, neural networks that train as differentiable softets then convert to 1-bit Boolean hardnets without performance loss. Using Wolfram Language, he demonstrates the approach on a car-quality dataset, achieving major compression for edge deployment while keeping soft and hard net performance nearly identical.
Questions this video answers
Ian Wright, machine learning engineer at GitHub.
January 27, 2024.
About 98% (three test examples wrong in the demo).
Hardnet classifier size ≈ 0.2 K (bits/bytes indicated by presenter); minimal MLP ≈ 16 K.
Paper: International Conference Machine Learning Workshop (this year); code: GitHub repository (speaker indicated a repo in the talk).
Chapters
Hello, everyone. Thanks for attending. My name is Ian Wright. I'm a machine learning engineer at GitHub. And today I want to talk about my experience using the Wolfram language for neural network research. In particular, I want to give an overview of how I use various language features to rapidly explore floor and prototype a new kind of neural network. So let's get straight into it. So, typically we train neural networks by iteratively adjusting their weights, by back propagating an error signal to minimize some measure of loss on the training data. And essentially this process performs gradient descent in a high dimensional weight space. Now, for back propagation to work, neural networks must be differentiable from functions, but this creates some problems. We can't learn discrete non differentiable functions, and sometimes discrete functions give better bias for certain kinds of learning problems. And the nets can be huge and very expensive to query because we store and operate on floating point numbers. Now, there's plenty of existing work that aims to reduce net sizes, such as weight quantization, which reduces the number of bits per weight. But quantization is a kind of lossy compression which degrades the network's predictive performance. So I wanted to see whether we could have our cake and eat it. I wanted to see whether we could define a new kind of neural network which is differentiable, so we can train it efficiently using backpropagation, but post training convert it to an equivalent Boolean function with 1 bit weights and without any loss of performance. So that was the aim. That's the idea behind DB net. So if you look at the diagram on the screen, a DB NET has two aspects. It has a softnet, which is differentiable, and a hard net, which is the Boolean function. We train the softnet as normal, convert its weights to 1bit Booleans and then bind those train weights with the Boolean function. And the result should be a Boolean function that behaves in exactly the same way as the trained differentiable net. So this is what I set out to do. I wasn't sure if it was possible, and I knew I had to rapidly explore and test lots and lots of different ideas, which is why I chose to prototype using the Wolfram language. And to explain how I use the Wolfram language to solve my problem, I need to start with some basics. I need to define two kinds of bits. So a hard bit is a Boolean value and a soft bit is a real number between 0 and 1. And, you know, obviously we can convert between soft bits to hard bits. So here's a conversion function. If the Soft bit is greater than a half, we count it as true, otherwise it is false. Now, what do we mean when we say that we want the hard net to be semantically equivalent to the softnet? What do we mean by that? So this equation at the bottom of the slide, the left hand side says essentially let's supply an input x of soft bits to the softnet and then get an output and then harden those soft bit outputs to hard bits. And then the right hand side says, let's harden the same input X and then supply the hard bits to the hard net to get another output. And if these two outputs are identical for all possible inputs, then the soft and hard nets are semantically equivalent. And this defines a hard equivalence requirement between the softnet and the hard net. And that's what DB nets have to satisfy. And this requirement turned out to be surprisingly hard to satisfy. I had to try lots and lots of different approaches. Over a period of a few months, I prototyped, no exaggeration, hundreds of ideas and the Wolfram language. Because it's a single unified system that provides everything immediately out of the box, it all works together. I essentially had a playground to try all kinds of wacky and wonderful ideas. And most of my ideas didn't work, but I could prototype really fast and discard them. So now I'll briefly explain how DBNETs in fact satisfy that hard equivalence requirement by stepping through a simple example. Now imagine that we want a network component that can learn to mask an input X to false if a learnable weight W is false. Otherwise we just pass through that X value. So essentially we want to learn a Boolean weight W for the Boolean function. And let me just slide forward and here is the and function. And let me just use a function from function repository, a truth table to show that and with a learnable weight W can learn to mask or pass through the input value X. You can see that in the truth table there now logical and of course is not a differentiable function. So we can't simply stick it inside a neural network and expect back propagation to work. So what can we do? There are real valued relaxations of Boolean logic. They already exist. And for example, product logic implements logical and by multiplying the two soft bits together. So how does this actually work? Well, we can use the manipulate functionality to explore how this function behaves. Now in this plot on the horizontal axis, axis is the input X and the vertical axis is the output of product. And and we can vary the other the input weight, like so. And here I Set the weight to one or two and so product and in this case should act as a pass through. And we can see that it does because when the input is low, the output is low. And when the input is high the output is high. Now let's set the learnable weight to zero. And you can see now it does act as a mask. No matter what the input value, the output is always zero. So product and is acting just like logical and at those two extremes of the weights. But let's look what happens at intermediate values. If we set the learnable way to say something like 0.6. Okay, The input here. Now the weight of a 0.6 hardens to true, so it should act as a pass through. But here the input is high, but the output is low in this region here. So it's a region where it doesn't act like product and at all. And therefore product and is not hard equivalent to logical. And that's a problem. Now the output varies continuously. We have gradients everywhere. So product and is a gradient rich function which is great for backpropagation, but it's not hard equivalent. And because it's not, if we used it at training time and then hardened the learned weights, in this case a single weight W and then use it with the logical and function, you would behave differently. We get different behavior. That means the softnet training wouldn't transfer to the hardnet. So that's no good at all. There are different relaxations. Here's another one. Girdle logic. Girdle logic defines logical and in terms of the min operator. And again we can use manipulate to dynamically look at how it behaves. Weight is one again. And yes, it's behaving as it should, like product hand weight is zero and yes, it's operating as a mask everything zero in the output. And what about that intermediate value? Now so you can see that unlike the previous product hand, girdle hand is hard equivalent even in the entire domain because it respects that boundary at 0.5 when a soft bit turns into from false to true, soft bit to hard bit. So that's good, it's hard equivalent. But there is a problem because if you think about it, the min operator essentially selects either the input value or the weight as its output. So any change in the unselected value won't change the value of the output. And that's why we can see a flat, a flat gradient here. And that means there's no gradient. And it means we can only back propagate error through one of its inputs. And that's actually bad for learning, that's bad for backpropagation. So girdle logic is hard equivalent, but it isn't gradient rich. So that is a problem that won't work for learning either. So we want DB nets to be composed of differentiable relaxations of logical operators that are both hard equivalent and gradient rich. So can we have our cake and eat it? And my eventual solution, after prototyping lots of different approaches, is something I call margin packing. And I can explain it with a dynamic visualization, basically an animation. I hope this will make it make it clear. Pretty easy to set up using the Wolfram language. Now here we're manipulating both the input value and the weight, and I'm plotting here the output value. And we can see that if I set the weight to one, and if we look at this vertical bold line here, the representative bit, the output is low, it's less than 0.5. And that's because the representative bit is just girdle and, and we know that is hard equivalent to logical. And so in fact if I make x high, then the whole, whole output of the representative bit turns high. Fine, Fine. Okay. So the next thing to notice is that there's always a margin between this representative bit and the hard threshold at 0.0.5. And any output values within this margin are guaranteed to be hard equivalent to logical. And so actually we, we have some wiggle room. We don't have to choose the representative bit. We can pick anything within this margin and still be hard equivalent. And so we can pack that margin with gradient rich information and we can actually compute an augmented bit, which you can hopefully see there, the augmented bit, which is still within the margin. But unlike the so unlike the representative bit which doesn't vary with the weight, the augmented bit does. And that means unlike girdle and this augmented bit, this new function does vary with all its inputs. So the mathematical details aren't important here, but we can simply margin pack girdle and to get a new function, which is this one here. And it's a piecewise continuous function, it's differentiable almost everywhere, which is sufficient for back propagation. And now we can actually visualize its behavior again. And you can see when the weight is high, it respects logicaland and as we vary remains hard equivalent, like girdle and but also like product hand, it is never flat anywhere, it always has a gradient. So we can back propagate through its inputs. So it's hard equivalent and gradient rich. And so we construct DB nets from margin packed differentiable analogs of Boolean operators. And then we can compose these new differentiable functions into multi input neurons and into multi input layers and then finally into complete neural network architectures. So in fact let's, let's do that. So I rapidly prototyped a DB NETS library building on top of Wolfram languages neural network support. And Wolfram's neural network functions, they're beautifully high level, they're very composable, they also give a lot of low level control. So I define layers that perform different kinds of logical operations on subsets of the inputs and all differentiable. But here, for the sake of time, I just want to look at one example which I think is particularly interesting because it demonstrates the power of Wolfram languages neural network support. So we'd like to be able to learn Boolean majority function. And the Boolean majority function is true if the majority of its inputs are true. So here's example of majority working now it's a discrete analog of a threshold function which turns out to be very useful for learning. And of course it could be implemented in terms of ands and ors, which we can already make differentiable. You know that we have those components which are hard equivalent to the logical operators. But so here's it is in DNF form a majority function with it with six inputs. But the number of terms grows exponentially with the number of inputs. And there's no general algorithm for, for finding the minimal Boolean representation of majority. And in machine learning applications we may want thousands and thousands of these things and so this becomes explosive to compute at training time. So what can we do about that? Okay, so notice that if we sorted the soft bit inputs in ascending size, then that the middle bit in that sorted array would be a representative bit and hard equivalent to Boolean majority. And why? Because if the majority of bits are high, then the middle of the sorted list is also guaranteed to be high. So remarkably, we can easily implement that idea using the Wolfram neural net network support. In fact we can use just create a net chain, just two lines of code. Essentially here's a function layer that takes built in sort function and applies it to the inputs and then a part there that picks out the middle bit, as I just explained. And then the Wolfram language takes care of everything, including compiling it down to differentiable code, which is pretty cool. And you might think, well, how on earth can we back propagate through a sorting algorithm? And the answer is, well, we're not actually, we're back propagating through the representative bit selected by the sorting algorithm. And then we can margin pack that representative bit with extra information from all the other inputs. So we can back propagate through all the inputs to the majority function. And then we get a hard equivalent and a gradient rich differentiable version of majority. So we avoid the expensive blow up by paying the cost of some runtime to sort the inputs. So let's in fact look at differentiable majority layer, which takes eight inputs. Now, every time we define a DB NET component, we get two things. We get the softnet and we get the corresponding hard net. Let's have a look at the softnet. The out of the box visualization is just really helpful here when we're wiring things together. And here we can see there's a component that does the sorting and there's a component that does the picking of the middle bit. Great. It's pretty incredible how the language just automatically wires the above two lines into a differentiable layer. We can immediately test this layer. Here's some. Some input. Eight soft bits. That's high. That's high. 1, 2, 3, 4, 5. Soft bits are high. So we expect the output of this differentiable majority layer to be high too. 0.52. It is high. And of course because it's actually picked out this as the representative bit, I haven't bothered margin packing here just for simplicity now. Okay, all well and good. What does the corresponding hardnet look like? Well, it's just a Wolfram function. There's some altercation because there's some technical issues there, but really just on the insides, all it boils down to is the built in Boolean majority function. What this means is that post training we can throw away all the sorting instead at query time we just actually use Boolean majority. And we can implement that incredibly cheaply, incredibly fast because it's just counting bits. So we trade off a bit more time at training time to gain a lot of time at query time. That's one of the ideas. Okay. And we can quickly check whether this Boolean function behaves the same as the softnet layer. So we're giving it the same inputs, we harden them and then we. Yep, it's true. 0.52 was true. It's all the same. It's true. So we got the hard equivalence. So the DB nets library exploits many different features of the Wolfram language. Neural network support. For example, I used a compiled layer to define special neurons that behave differently in their forward and backward passes. The details are important, but what is important, I think is that the Wolfram language allows you very deep control, which is essential for research purposes when you're exploring novel ideas. Okay, let's put this all together and demonstrate DBNETs working in practice. So let's pull down some data set from the Wolfram data repository. Basically, this is a data set describing features of cars. And the aim is to predict the four labels in this final column, which is one of unacceptable, acceptable, good and very good that defines whether the car is worth purchasing. So this is a multi class classification problem. I chose a small one just so I can demo it in real time. Okay, Let's again use a function from the function repository to quickly split that data into a train and test set. And then we can use the built in. Net encoder functionality to make it really easy to convert those input features into an indicator vector of boolean. Here, for every feature we have an encoder, it takes in some categorical values and outputs a vector of booleans. Very easy to do. And then we can compose those encoders into a single input encoding layer, which again we can visualize. Here are the input features and the output is a concatenation of those booleans into a single vector. So it's really cool how easy it is to do that and visualize what's going on. So that's our features into a Boolean representation. All right, let's just run this code. So now I'll define a complete VB NET architecture that learns this classification problem and let's have a look at it. So essentially it has two layers. One layer here is a differentiable learn to OR layer. And this layer here is a differentiable learn to not layer of those Boolean inputs. Each layer has about 64 logical neurons with these new kinds of margin packed activation functions. The NET has four output ports for each class. Each output port port is a vector of bits. We add up the bits in each port and interpret them as relative class probabilities. Which means we can then use standard cross entropy loss for this for training for this classification problem. And that is in fact ready to train. Compiles down to GPU code without me having to do anything. Runs on my local laptops. Gpu, not particularly powerful gpu. And off it goes straight out of the box. That's very, very convenient, very pleasant. And we can visualize it and you can see that the loss is going down. It's learning something on this problem. We could stop the training anytime. I'll give it a little bit longer to learn a bit more. And I'M just going to stop it some random position. So I think it's probably sufficient. Stop there. All right, so we've done some training. Pull out the train net and therefore we've got the trained weights. Let's do a quick evaluation of the performance of that net on the test data. And here I'm going to use Wolfram's classifier measurements function to evaluate the softnet. And it automatically gives lots of very useful information. And here we can see the confusion matrix. And obviously the softnet has basically solved the problem. I think it's got three, three examples wrong. Accuracy around 98%. If I left it longer, it would have done even better. All right, but what we really want is the Boolean function that predicts whether a car is acceptable or not. We want the hard net, and we can extract it by binding the hardened train weights from the softnet with the corresponding hard net. And here we go, we've done that. What we get is a Wolfram function, again, more complex, this time bound with the learned Boolean weights. Now, this function, apart from the obfuscation, reduces down to logical operations. And let me just. That's quite a large expression. Let's actually look into that. And because this is a Wolfram function, we can symbolically evaluate it to see precisely what Boolean function has been learned. And there it is. And here we're looking at the symbolic hard bits that are outputted by the net on one of its four ports. The Bs correspond to the 21 input bits which represent features of the cars. And each Boolean expression logically combines different feature values to give one bit of predictive information. And so we have in fact learned by backpropagation, a Boolean function that classifies cars. And how does it perform? Well, let's evaluate it on the test data and this time we can actually run that Wolfram function to see how it performs on the. On the test data. There's no compilation here. It's actually virus it kind of symbolically, so that could be sped up. Accuracy is 98%. It's identical to the trained softnet. That's because we do meet that requirement of hard equivalence. So we also get this classification performance with a much, much smaller net at query time. In fact, this classifier is only 0.2 K because each weight is a single bit. So it can be very compactly represented. And in comparison, a minimal size multi layer perceptron on the same classification problem, the weights take up about 16k, and Boolean and integer operations can be much Cheaper than floating point operations. So this makes DB nets potentially a good choice when deploying machine learning on edge devices. And I should point out that DB nets are not restricted to toy classification problems. We can use them for numerical regression and image recognition, and in fact, anything we want. Okay, let me come to a close. The Wolfram language is great for rapidly prototyping new ideas. Everything works out of the box, all works together without any problems. You don't have to think about choosing packages or versions or dependencies. You can immediately start working on your problem without any distractions or frustrations. The neural network support is very high level, yet it's also very flexible and configurable. So if you want to really push what neural networks can do, explore entirely new kinds of networks, things that haven't been thought of before, then you can. In this framework, you're not restricted in any way. And I had to rapidly implement and evaluate hundreds of ideas before solving my particular research problem. And without the Wolfram language, I think I would have actually given up. But it's such a joy to use. And the gap between ideas and reality was so small that I enjoyed every step. Even when my ideas weren't working, I was still enjoying myself. If you want to learn more about DB Nets, there's a paper published at the International Conference Machine Learning Workshop this year and there's a GitHub repo with all the code. Let me stop there. Thanks for listening. Let me just, just check if there's any questions to answer and I'll do my best to answer them. So there is a question. Would it be possible to implement differentiable ranking loss functions from these and then extract hard ranks too? Yes, it is. It absolutely is. In fact, it's like a little new ecological niche in the design space in your networks. But you can try and do all kinds of new things. And had some more exotic experiments that are not even in the paper, which I would like to explore. But we can even. We can have differentiable counting and other cool things like that. I think we can also have differentiable if and else statements. And at the back of my mind is perhaps the hope that we can make arbitrary code differentiable. Although probably the local minimums might not be very good. But yes, you can do all kinds of new things. It's kind of, kind of interesting and I'm kind of excited about it. So there is a. The notebook is in the handout form, so if you want to try it, please go ahead. And the code is all open source, so give it a swing. If you're interested, happy to talk about it with anyone else after this event. But for now, I think I should come to a close. So thank you very much, everyone.
More from Guy

Cheryl Strauss Einhorn | The Human Edge: Smarter Decisions in the Age of AI | Talks at Google

Rachel Glennerster | Innovation and the Center for Global Development | Talks at Google

Rebecca Winthrop | Rethinking the Purpose of Education in the Age of AI | Talks at Google

Dr. Robert Wachter | A Giant Leap: How AI Is Transforming Healthcare... | Talks at Google






