Hello everyone and welcome to episode 76 of the Stanford MLSys seminar series.

Today of course we're or this year we're very excited to be partnered with CS324 advances in Foundation models.

Today I'm joined by Michael say hi and Ivonica.

And today our guest is Jack Rae from OpenAI and he's got a very exciting talk prep for us about compression and AGI.

So so we're very excited to listen to him as always if if you have questions you can post them in YouTube chat or if you're in the class there's that Discord Channel.

So so to keep the questions coming and after his talk we will we'll have a great discussion.

Um so with that Jack take it away.

Okay fantastic thanks a lot.

Um today I'm going to talk about compression for AGI and the theme of this talk is that I want people to kind of think deeply about uh Foundation models and their training objective and think deeply about kind of what are we doing why does it make sense what are the limitations.

This is quite a important topic at present I think there's a huge amount of interest in this area in Foundation models large language models their applications.

And a lot of it is driven very reasonably just from this principle that it works and it works so it's interesting.

But if we just kind of sit within the kind of it works realm it's hard to necessarily predict or have a good intuition of why it might work or where it might go.

So some takeaways that I want so I hope people like people hopefully to take from this tour car some of them are quite pragmatic.

So I'm going to talk about some background on the minimum description length and why it's seeking the minimum description length of our data may be an important role in solving perception.

Uh I want to make a particular point that generative models are actually lossless compressors and specifically large language models are actually state of the art lossless compressors.

Which may be a counter-intuitive point to many people given that they are very large and use a lot of space and I'm going to unpack that in detail.

And then I'm also going to kind of end on some notes of limitations of the approach of compression.

So let's start with this background minimum description length and why it relates to perception.

So even going right back to the kind of ultimate goal of learning from data we may have some set of observations that we've collected some set of data that we want to learn about which we consider this small red circle.

And we actually have a kind of a two-pronged goal we want to learn like uh how to kind of predict and understand our observed data with the goal of understanding and generalizing to a much larger set of Universe of possible observations.

So we can think of this as if we wanted to learn from dialogue data for example we may have a collection of dialogue transcripts but we don't actually care about only learning about those particular dialogue transcripts.

We want to then be able to generalize to the superset of all possible valid conversations that a model may come across right.

So what is an approach what is a very like rigorous approach to trying to learn to generalize well I mean this has been a philosophical question for multiple thousands of years.

And even actually kind of full Century BC uh there's like some pretty good um principles that philosophers are thinking about.

So Aristotle had this notion of um assuming the super superiority of the demonstration which derives from fewer postulates or hypotheses.

So this notion of uh we have some um simple set of hypotheses um then this is probably going to be a superior description of a demonstration.

Now this kind of General kind of simpler is better um theme is more recently attributed to William 14th century or Cam's Razer this is something many people may have encountered during a machine learning or computer science class.

He is essentially continuing on this kind of philosophical theme the simplest of several competing explanations is always likely likely to be the correct one.

Um now I think we can go even further than this within machine learning.

I think right now Occam's razor is almost used to defend almost every possible angle of research but I think one actually very rigorous incarnation of what comes Razer is from race Island's theory of inductive inference 1964.

So we're almost at the present day and he says something quite concrete and actually mathematically proven.

Which is that if you have a universe of data which is generated by an algorithm and observations of that universe so this is the small red circle encoded as a data set are best predicted by the smallest executable Archive of that data set.

So that says the smallest lossless prediction or otherwise known as the minimum description length.

So I feel like that final one is actually putting into mathematical and quite concrete terms um these kind of Notions that existed through timing velocity.

And it kind of we could even relate this to a pretty I feel like that is a quite a concrete and actionable retort to this kind of um quite um murky original philosophical question.

But if we even apply this to a well-known philosophical problem cells Chinese room 4 experiment.

Where there's this notion of a computer program or even a person kind of with it within a room that is going to perform translation from English English to Chinese.

And they're going to specifically use a complete rulebook of all possible inputs or possible say English phrases they receive and then and then the corresponding say Chinese translation.

And the original question is does this person kind of understand how to perform translation.

Uh and I think actually this compression argument this race on this compression argument is going to give us something quite concrete here.

So uh this is kind of going back to the small red circle large white circle if if we have all possible translations and then we're just following the rule book this is kind of the least possible understanding we can have of translation.

If we have such a giant book of all possible translations and it's quite intuitive if we all we have to do is coin a new word or have a new phrase or anything which just doesn't actually fit in the original book this system will completely fail to translate.

Because it has the least possible understanding of translation and it has the least understandable version of translation because that's the largest possible representation of the the task the data set.

However if we could make this smaller maybe we kind of distill sorry we distill this to a smaller set of rules some grammar some basic vocabulary and then we can execute this program maybe such a system has a better understanding of translation.

So we can kind of grade it based on how compressed this rulebook is.

And actually if we could kind of compress it down to the kind of minimum description like the most compressed format the task we may even argue such a system has the best possible understanding of translation.

Um now for foundation models we typically are in the realm where we're talking about generator model one that places probability on natural data.

And what is quite nice is we can actually characterize the lossless compression of a data set using a generator model in a very precise mathematical format.

So race on enough says we should try and find the minimum description length well we can actually try and do this practically with a generator model.

So the size the lossless compression of our data set D can be characterized as the negative log likelihood from a genetic model evaluated over D plus the description length of this generator model.

So for a neural network we can think of this as the amount of code to initialize the neural network.

That might actually be quite small this is not actually something that would be influenced by the size of the neural network this would just be the code to actually instantiate it.

So it might be a couple hundred kilobytes to actually Implement a code base which trains a transformer for example and actually this is quite a surprising fact.

So what does this equation tell us does it tell us anything new well I think it tells us something quite profound.

The first thing is we want to minimize this general property and we can do it by two ways.

One is via having a generative model which has better and better performance of our data set that is a lower and lower negative log likelihood.

But also we are going to account for the prior information that we inject into F which is that we can't stuff F full of priors such that maybe it gets better performance but overall it does not get a bit of a compression.

Um so on that note yeah compression is a a cool way of thinking about how we should best model our data and it's actually kind of a non-gameable objective.

So contamination is a big problem within uh machine learning and trying to evaluate progress is often hampered by Notions of whether or not test sets are leaked into training sense.

Well with compression this is actually not not something we can game.

So imagine we pre-trained F on a whole data set D such that it perfectly memorizes the data set.

AKA such that the probability of D is one log probability is zero.

In such a case if we go back to this formula the first term will zip to zero.

However now essentially by doing that by injecting and pre-training our model on this whole data set we have to add that to the description length of our generative model.

So now F not only contains the code to train it Etc but it also contains essentially a description length of d.

So in this setting essentially a pre-contaminating f it does not help us optimize the compression.

And this contrasts to regular test set benchmarking where we may be just measuring test set performance and hoping that measures generalization and is essentially a proxy for compression.

And it can be but also we can find lots and lots of scenarios where we essentially have variations of the test set that have slipped through the net in our training set.

Okay so we've talked about philosophical backing of the minimum description length and maybe why it's a sensible objective.

And now I'm going to talk about it concretely for large language models and we can kind of map this to any uh generative model but I'm just going to kind of ground it specifically in the marsh language model.

So if we think about what is the log problem of our data D well it's the sum of our next token prediction of tokens over our data set.

Um so this is something that's essentially our training objective if we think of our data set D um and we have one Epoch then this is the sum of all of our training loss.

So it's pretty tangible term it's a real thing we can measure and F is the description length of our Transformer language model.

Uh and actually there are people that have implemented a Transformer and a training regime just without any external libraries in about I think 100 to 200 kilobytes so this is actually something that's very small.

Um and and as I said I just want to enunciate this this is something where it's not dependent on the size of our neural network.

So if a piece of code can instantiate a 10 layer Transformer the same piece of code you can just change a few numbers in the code it can instantiate a 1000 layer Transformer.

Actually the description length of our initial Transformer is unaffected really by how large the actual neural network is.

We're going to go through an example of actually using a language model to losslessly compress where we're going to see why this is the case.

Okay so let's just give like a specific example and try and ground this out further so okay llama it was a very cool paper that came out from fair just like late last week.

I was looking at the paper here's some training curves um now forgetting the smaller two models there are the two largest models are trained on one Epoch of their data set.

So actually we could sum their training losses uh AKA this quantity and we can also roughly approximate the size of of the um of the code base that was used to train them.

Um and therefore we can see like okay which of these two moles the 33b or the 65b is the better compressor and therefore which would we expect to be the better model at generalizing and having greater set of capabilities.

So it's pretty it's going to be pretty obvious at 65b I'll tell you why firstly just to drum this point home these models all have the same description length they have different number of parameters but the code that's used to generate them is actually of same of the same complexity.

However they don't have the same integral of the training loss 65b has a smaller integral Windows training loss.

And therefore if we plug if we sum these two terms we would find that 65b essentially creates the more concise description of its training data set.

Okay so that might seem a little bit weird I'm going to even plug some actual numbers in let's say we assume it's about one megabyte for the code to instantiate and train the Transformer.

And then if we actually just calculate this roughly it looks to be about say 400 gigabytes um you have some of your log loss converting into bits and then bytes it's going to be something like 400 gigabytes.

And this is from an original data set which is about 5.6 terabytes of rortex so 1.4 trillion tokens times four is about 5.6 terabytes so that's a compression rate of 14x.

Um the best text compressor on the Hudson prize is 8.7 X.

So the takeaway of this point is um actually as we're scaling up and we're creating more powerful models and we're training them on more data we're actually creating something which actually is providing a lower and lower lossless compression of our data even though the intermediate model itself may be very large.

Okay so now I've talked a bit about how large language models are state of the art lossless compressors but I just want to maybe go through the mechanics of how do we actually get a something like a generative model literally losslessly compress.

This may be something that's quite mysterious like what is happening like when you actually losslessly compress this thing is it the weights or is it something else.

So I'm going to give us a hypothetical kind of scenario we have two people sat here in Sundar Satya wants to send a data set of the world's knowledge encoded in D to send R.

They both have access to very powerful supercomputers but there's a low bandwidth connection.

We are going to use a trick called arithmetic encoding as a way of communicating the data set.

So say we have a token x a timestep t from of some vocab and a probability distribution p over tokens.

Arithmetic encoding without going into the nuts and bolts is a way of allowing us to map our token x given our probability distribution over tokens to some Z where Z is essentially our compressed transcripts of data.

And Z is going to use exactly minus log 2 p t x t bits so.

The point of this step is like arithmetic encoding actually Maps it to some kind of like floating Point number as it turns out and it's a real algorithm.

And we can use arithmetic decoding um to take this encrypted transcript and as long as we have our probability distribution of tokens we can then recover the original token.

So we can kind of think about probability probability distribution as kind of like a key it can allow us to kind of lock in a compressed copy of our token and then unlock it.

So if p is uniform so there's no information about our tokens then this would be this one over v p is just one over the size of V so we can use log 2 V bits of space uh that is just essentially the same as naively storing in binary uh our our XT token.

If p is an oracle so it knows like exactly what the token was going to be so P of x equals one then log 2p equals zero and this uses zero space.

So these are the two extremes and obviously what we want is a generative model which better and better molds our data and therefore it uses less space.

So what would actually happen in practice if Satya can take his data set of tokens trainer Transformer and get a subsequent set of probabilities uh over the tokens like so next token prediction.

And then use arithmetic encoding to map it to this list of transcripts and this is going to be of size sum of negative log likelihood of your Transformer over the data set.

And he's also going to send he's going to send that list of transcripts and some code that can deterministically train a larger Transformer.

And then on the other side Sundar can run this code which is deterministic and the mod is going to run the neural network.

It gives a probability distribution to the first token he's going to use arithmetic decoding with that to get his first token.

You can either train on that or whatever the code does so then continue on predict the next token etc etc and essentially iteratively go through and recover the whole data set.

So this is kind of like almost a fourth experiment because in practice to send this data at 14x compressed compression say if we're talking about the Llama model uh that's it's a bit more compressed than gzip but this is requiring a huge amount of intermediate compute switches to train a large language model which feels inhibitive.

But this thought experiment is really derived not because we actually might want to send data on a smaller and smaller bandwidth it's also just derived to kind of explain and prove why we can actually losslessly compress with language models and why that is their actual objective.

And if this kind of setup feels a little bit contrived well the fun fact is this is the exact setup that called Shannon was thinking about um when he kind of proposed language models in the 40s.

So let's just think about solving perception and moving towards AGI what's the recipe well it's kind of a two-step process.

One is collect all useful perceptual information that we want to understand and the second is learn to compress it as best as possible with a powerful Foundation model.

So the nice thing about this is it's not constrained to a particular angle for example you can use any research method that improves compression and I would posit that this will further Advance our capabilities towards perception based on this rigorous foundation.

So that might be a better architecture it may be scale further scaling of data and computes.

This is in fact something that's almost become a meme people say scale is all you need but truly I think scale is only going to benefit as long as it is continuing to significantly improve compression.

But you could any use any other technique and this doesn't have to be just a regular generative model it could even we could even maybe spend a few more bits on the description length of F and add in some tools add in things like a calculator allow it to make use of tools to better predict its data.

I just want to remark at this point on a very common point of confusion on this topic which is about lossy compression.

Um I think there's a very interesting New Yorker article about about this kind of Topic in general kind of thinking about you know what are what are language models doing what are Foundation models doing and I think there's a lot of confusion in this article specifically on this topic.

Where from the perspective of glossy compression and neural network feels very kind of sub-optimal it's losing information in Red so it doesn't even do reconstruction very well and it's potentially bloated and larger and has all these other properties.

I just wanted to take this kind of point to reflect on the original goal which is we really care about understanding and generalizing to the space of the universe of possible observations so we don't care and we don't train towards reconstructing our original data.

Um so I think there's one concern with this approach which is that it may be just the right thing to do or like an unbiased kind of attempt at solving perception but maybe it's just not very pragmatic and actually trying to kind of model everything and compress everything it may be kind of correct but very inefficient.

So I think Image level modeling is a good example of this where modeling a whole image at the pixel level has often kind of been prohibitively expensive to like work incredibly well.

If we turn this to video modeling every pixel of every frame it really feels preemptively crazy and expensive.

Um another very valid point is I think this is often framed to people that maybe are thinking that this is like the only ingredient for AGI is that crucially there's lots of just very useful information in the world that is not observable.

And therefore we can't just expect to compress all observable observations achieve AGI because there'll just be lots of things we're missing out.

Um so I think a good example of this would be something like Alpha zero so playing the game of Go um I think if you just observe the limited number of human games that have ever existed one thing that you're missing is all of the intermediate search trees of all of these expert players.

And one nice thing about something like Alpha zero with its kind of self-play mechanism is you essentially get to collect lots of data of intermediate search trees of many many different types of games.

Um so that kind of on policy behavior of like actually having an agent that can act and then Source out the kind of data that it needs I think is still very important.

Um I think people are a bit confused about um how is it possible for the receiver to decode the message and get the original data set back without having access to the train bottle.

The receiver does not receive the neural network it just receives the code to instantiate kind of the fresh neural network and run the identical training setup that it saw before.

So there's a fresh neural network uh that's going to give us like a probability distribution for the first token and so he's got this probability distribution for the first token and he's got the transcript um of what that token should be and you can use arithmetic decoding to actually recover that first token.

And then let's imagine for Simplicity we actually like train like one SGD step on one token at a time so we take our SGD step and then we have the model that's like was used to predict the next token so we can get that P2 we have Z2 and then we can recover X2.

So Michael is slacking me he wants me to ask if you follow the S4 line of work. Does that suggest you that like something like S4 or something with a different um you know encoding would uh would have these like implications for I don't know being more intelligent or or being a better compressor of these other modalities.

Yeah I think that's a really important architecture. So like on a broad brushstroke like S4 allows you to maybe have a much longer context uh than attention without paying the quadratic compute cost.

So one is we currently have a huge limitation in our architecture which is a Transformer or even just like a deep content and that is that the architecture does not adapt in any way to the information content of its inputs.

The Transformer will just spend four times more compute on the byte level sequence if it was fed it and it'll spend four times Less on the bpe sequence of this feather even though they have the same information content.

So we don't have some kind of algorithm which could like kind of fan out and then just like process the byte level sequence with the same amount of approximate compute. And I think that really hurts images.

I think S4 is very cool I think it could be could help in this direction for sure.