Compression for AGI - Jack Rae | Stanford MLSys #76

hello everyone and welcome to episode 76

of the Stanford MLS seminar series

um today of course we're or this year

we're very excited to be partnered with

cs324 advances in Foundation models

um today I'm joined by Michael say hi

and ivonica

um and today our guest is Jack Ray from

openai and he's got a very exciting talk

uh prep for us about compression and AGI

um 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

um 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

and right

okay so

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

um

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

um

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

[Music]

um

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 and actually even right

now within Labs comparing large language

models this notion of contamination

affecting eval resources a continual

kind of phone in um in in the side of

kind of clarity

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

[Music]

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 this is like something that

exists in the real world it does require

technically infinite Precision to to use

exactly these number of bits and

otherwise you maybe you're going to pay

a small cost for implementation but it's

roughly approximately optimal in terms

of the encoding 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

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 so

he sends those two things what does that

equal in practice the size of f the size

of your generator model description plus

the size of your some of your negative

log likelihood of your data set so as

you can see it doesn't matter whether

the Transformer was one billion

parameters one trillion parameters

plus plus he's not actually sending the

neural network he's sending the

transcript of encoded logits plus the

code

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

um 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

um 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 he was thinking about

having a discrete set of data and how

can we better communicate to overload

over a low bandwidth Channel and

language models and entropy coding

essentially was the topic that he was

thinking about about labs

Okay so we've talked mechanically about

well we've talked about the philosophy

of kind of why do why why be interested

in description length relating it to

generalization talks about why

generative models are lossless

compressors talked about why our current

large language models are actually

state-of-the-art lossless compressors

than our providing some of the most

compressed representations of our source

data 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

allow it to retrieve over the past use

its own synthetic data to generate and

then learn better there's many many

angles we could think about that are

within the scope of a model

better better compressing it Source data

to generalize over the universe of

possible observations

I just want to remark at this point on a

very common point of confusion on this

topic which is about lossy compression

so I think it's a very reasonable

um

thought to maybe confuse what a neural

network is doing with glossy compression

especially because

information naturally seeps in from the

source training data into the weights of

a neural network and neural network can

often memorize it often does memorize

and can repeat many things that it's

seen but it doesn't repeat everything

perfectly so it's lossy and it's also

kind of a terrible lossy compression

algorithm so if in the velocity

compression case you would actually be

transmitting the weights of the

parameters of a neural network and they

can often actually be larger than your

Source data so 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 I think if we did then this article