The Hidden Life of Embeddings: Linus Lee

Key Takeaways

[Music] everyone I'm L I'm here to talk about embeddings um grateful to be here at the inaugural a engineer conference who learned something new today yeah um before I talk about that a little about myself if you don't know me already I am lus I work on AI at notion uh for the last year or so before that I did a lot of independent work prototyping experimenting with trying out different things with language models with traditional LP things like tfidf bm25 to build interesting interfaces for reading and writing in particular I worked a lot with embedding models and Laten spaces of models which is what I'll be I'll be talking about today but before I do that I want to take a moment to say it's been almost year since notion launched notion AI our public beta was first announced in around November 2022 so as we get close to a year we we've uh been steadily launching new and interesting features inside notion AI from November we have ai autofill inside databases translation and things coming soon though not today so keep an eye on the space and obviously we're hiring just like everybody else here we're looking for AI Engineers product Engineers machine learning Engineers to tackle the full gamut of problems that people have been talking about today agents tool use evaluations data uh training and all the interface stuff that that we'll see today and tomorrow so if you're interested please grab me and uh and we'll have a little chat now it wouldn't be Al us talk without talking about lat in spaces so let's talk about it um the the one of the problems that I find always motivated by is the problem of steering language models and I always say that prompting language models feels a lot like you're steering a car from the back seat with a pool noodle like yes technically you have some control over the motion of the vehicle so like there's some connection but like you're not really in the driver's seat the control isn't really there it's not really direct there's like three layers of indirection between you and what the vehicle is doing and that to me trying to prompt a model especially smaller more efficient models that we can use for production with just tokens just prompts feels a lot like uh there's too many lers of indirection and as even though models are getting better at understanding prompts I think there's always going to be this fundamental barrier between indirect kind of control of models with just prompts and and getting with the model getting the model to do what what we want them to do and so perhaps we can get a closer layer of control a more direct layer of control by looking inside the model which is where we look at latent spaces latent spaces arise I think uh most famously inside embedding models if you embed some pieces of text that Vector of 1536 numbers or 1024 numbers is inside a high dimensional Vector space that's a latent space but also you can look at Vector at the latent spaces inside activation spaces of models inside token embeddings inside image models and then obviously other model architectures like Auto encoders today we're going to be looking at embeding embedding models but I think a lot of the general takeaways apply to other models and I think there's a lot of fascinating research work happening inside other models as well well when you look at an embedding you kind of see this right you see like rows and rows of numbers if you ever debugged some kind of an embedding Pipeline and you print out the embedding you can kind of tell it that's like a thousand numbers but it's just looking at like a matrix uh screen of numbers ring down but in theory there's a lot of information actually packed inside those embeddings if you get an embedding of a piece of text or image these latent spaces these embeddings represent in theory the most Salient features of the text or the image that the model is using to lower its loss or do its task and so maybe can disentangle some meaningful attributes or features out of these embeddings if we can look at them a little more closely and and and interpret them a little better maybe we can build more expressive interfaces that let them control the model by interfering or intervening inside the model another way to say that is that emedic show embeddings show us what the model sees in a sample of input so maybe we can read out what it sees and try to understand better what what the model's doing and maybe we can even control the embedding intermediate activations to see uh what the model can generate so let's uh let's see some of that so some of this some of you might have seen before but I promise there's some new stuff at the end so hang tight so here are some sentence that I have it's a sentence about this novel one of my favorite novels named diaspora it's a science fiction novel by Greg Egan that explores the evolution and existence of postum artificial intelligences uh something to do with alien civilizations and the questioning the nature of reality and Consciousness which you might be doing a lot given all the things that are happening and so I have I have I've trained this model that can uh generate some embeddings out of this text so if I hit this enter it's going to give us an embedding but it's a it's an embedding of length 248 and so it's it's quite large but it's just a row of numbers right but then I have a A decod or half of this model they can take this embedding and try to reconstruct the original input that may have produced this embedding so in this case it took the original sentence there's some Vari you can tell it's not exactly the same L maybe but it's mostly reconstructed the original sentence including specific details like the title of the book and so so we have an encoder that's going from text to embedding and a decoder that's going from embedding back to text and now we can start to do things with the embedding to vary it a little bit and and see what the decoder might see if we if we make some mod modifications to the embedding so here I've uh try to kind of blur the embedding and sample some points around the embedding with this this blur radius and you can see the text generated from those blurry embeddings they're a little off like this is not the correct title the title's kind of gone here um the the it still kept the name Greg but it's it's a different person and so there's kind of a blur a semantic blur that's happened here but this is kind of boring this is not really useful what's a little more useful is trying to actually manipulate things in more meaningful directions so now we have the same taste with text and now here I have a bunch of controls so maybe I want to find a direction in this embedding space here computed a direction where if you push an embedding in that direction that's going to represent a shorter piece of text of roughly the same topic and so I picked this direction and I hit go and uh it'll try to move push the embedding of this text in that direction and decode them out and you can tell they're a little bit shorter um if I push it a little bit further even so now I'm taking that shorter Direction and moving a little farther along it and and sampling generating text out of those embeddings again and they're they're even a little bit shorter but they've they've still kept the general kind of idea General topic and uh with that kind of building block you can build really interesting interfaces like for example I can I can plop this piece of text down here and maybe I want to generate a couple of a couple of sort of shorter versions um so this is like a little bit shorter this is even even more short but uh maybe maybe I like this version so I'm going to clone this over here and I'm going to make this the sentiment of the sentence a little more negative and you can start to explore the lat space of of latent space of this embedding model this language model by actually moving around in a kind of spal canvas interface which is is kind of interesting another thing you can do with this kind of end model is now that we have a vague sense that there are specific directions in this space that mean specific things we can start to more directly look at a text and ask the model hey where does this piece of text lie along your length direction or along your negative sentiment Direction so uh this is the original text that we've been playing with it's pretty objective like a Wikipedia style piece of text here I've asked CH PT to take the original text and uh make it sound a lot more pessimistic so uh things like the futile quest for meaning and plunging deeper into the abys of nihilism and uh and if I embed both of these what I'm asking the model to do here is embed both of these things in the embedding space of the model and then project those embeddings down onto each of these directions so one way to read this table is that um this default piece of text is at this point in this negative Direction which by itself doesn't mean anything but it's clearly less than this so this piece of text is much further along the negative sentiment axis inside this model when you look at other properties like how how much um of the artistic kind of topic does it talk about is roughly the same the length is roughly the same um the the maybe the the negative sentiment text is a bit more elaborate in its its vocabulary and so you can start to project these things into these meaningful directions and say what are the features that the model's what are the attributes that the model's um finding in the in the text that we're feeding it uh another way you could test out some of these ideas is by mixing embeddings and so here I'm going to embed both of these pieces of text this one's the one that we've been playing with this one is a beginning of a short story that I wrote once it's about this like town in the Mediterranean uh Coast that's calm and and a little bit old and um both of these have been embedded and so I'm going to say just this is a 2,000 dimensional embedding I'm going to say give me a new embedding that's just the first thousand or so Dimensions from the one embedding and then take the last thousand dimensions of the second embedding and just like slam them together and have this new embedding and navely you wouldn't really think that that would amount too much that would be kind of gibberish but actually if you generate some samples from it um you can tell you can see in a bit uh you get a sentence that's kind of a semantic mix of both you have structural similarities to to both of those things like you have this structure where there's a quoted kind of title of a book in the beginning there's topical similarities there structur there's uh punctuation similarities tone similarities and so this is an example of inter in lat space the last I have I have uh you may have seen on Twitter is um about okay I have this un edding model and I have kind of an un edding model that works pretty well can I use this un embedding model and somehow fine-tune it or otherwise adapt it so you can read out text from other kinds of embedding spaces so this is the same sentence we've been using but now when I hit this run button it's going to embed this text not using my embedding model but using open AI Ada text AA 2 uh and then there's a linear adapter that I've trained so that my decoder model can read out not from my embedding model but from open ai's embedding space so I'm going to embed it it's going to try to decode out the text from given just the open AI embedding and uh you can see okay it's not it's not as perfect but there's a surprising amount of detail that we've recovered out of just the embedding with no reference to the source text so you can see this proper noun diaspora it's like surprisingly still in there um this feature where there's a a quoted title of a book is in there it's roughly about the same topic things like the Rogue AI um sometimes when I when I rerun this there's also references to the author where the name is roughly correct so even surprising uh things features like proper nouns punctuation things like the quotes uh General structure and topic obviously those are recoverable given just the embedding because the the kind of deta amount of detail that these high high capacity embedding spaces have but not only can you do this in the text space you can also do this in image space so um here I have a I have a few prepared files um let's start with me and um for dumb technical reasons I have to put two of them in and then let's try to interpolate in this image space so this is now using Clip Clips embedding space I'm going to try to generate say like six images in between uh me and the the notion avatar version of me the cartoon version of me if that back end will warm up cold starting models is sometimes difficult uh there we go so now it's Shing six images bridging kind of interpolating between the photographic version of me and the cartoon version of me uh and uh again it's not perfect but you can see here it's on the left it's it's quite photographic and then as you move further down this interpolation you're you're seeing uh more kind of cartoony features appear here and uh it's actually quite a surprisingly smooth transition um another thing you can do on top of this is you can you can do text manipul as because aul text image model and can say time let's add some text I'm going to subtract the vector for a photo of a smiling man and instead I'm going to add the vector for a photo of a very sad crying man uh and then I'll I'll embid these pieces of text and I empirically I find that for text I have to be a little more careful so I'm going to dial down how much of those vectors I'm adding and subtracting and then generate six again and it's taking a bit okay I'm really sad and you can do you can do even more fun things like you can try to add like here's a photo of a beach I'm going to try to add some some beach in this this time maybe just generate four for a sake of time um or maybe there's a bug and it won't let me generate so in all these demos that I've done both in the text and image domain okay the beach didn't quite survive the uh the lat and space arithmetic but uh in all these all these demos the only thing I'm doing is calculating vectors calculating embeddings for examples and uh and embedding them and just adding them together and some with some normalization and it's surprising that just by doing that you can try to manipulate um interesting features in text and images and with this you can also do things like at style and subject at the same time you can uh this is a cool image that I thought I generated when I my first demo and then you can also do some pretty smooth transitions between landscape imagery so that's interesting um in all these prototypes one principle that I've tried to reiterate to myself is that often times when you're studying this very complex sophisticated um models you don't really necessarily have the ability to look inside and say okay what's what what's happening not even get an intuitive understand even getting an intuitive understanding of what is a model thinking what is a model looking at can be difficult and I think these are some of the ways that I've tried to render the these invisible parts of the model a little bit more visible to let you more a little bit more directly observe exactly what the model is um the representations the model is operating in and sometimes you can also take those and directly interact or let directly interact with the representations to explore what these faes represent um and I think there's there's a ton of interesting pretty groundbreaking research that's happening here um on the left here is the AO World model paper which is fascinating neurons on a hyack and then on the right is very very recent I had to add this in last minute because it's super relevant um in a lot of these examples I've calculated these feature Dimensions by just giving examples and calculating centroids between them but here anthropics new work along with other work from conjecture and other labs have found unsupervised ways to try to automatically discover these Dimensions inside models so that's super exciting and in general I'm really excited to see latent spaces that appear to encode you know by some definition interpretable controllable representations of the models input and output um I want to talk a little bit in the last few minutes about the models that I'm using the text model is a custom model um I won't talk I won't go into too much detail but it's fine-tuned from uh T5 checkpoint with as a d using Auto encoder it's an incoder decoder Transformer with some modifications that you can see in the code um so here's a general Transformer incoder on the left decoder on the right I have some pooling layers to get an embedding this is like a normal T5 embedding model stack and then on the right I have this special kind of gated layer that pulls from the embedding to uh to decode from the embedding you can look at the code it's a little more uh easy to understand but we take this model and we can adapt it to other models as well as you saw with the open AI embedding recovery and so on the left is the normal trading regime where you have an encoder you get an embedding and you try to reconstruct the text on the right we just just train this linear adapter layer to go from embedding of a different model to then reconstruct the text with a normal decoder um and today I'm excited to share that these models that I've been deing with uh that you may have asked about before are open on hogging face so you can go download them and try them out now uh these are the links on the left is the hug face models and then there's a collab notebook that lets you get started really quickly try to do things like interpolation and interpretation of these features and so if you find any interesting results with these please let me know and and if you have any questions also reach out and I'll be able to help you out the image model that I was using at the end wasow brain Carlo excited to see Korea um stepping out there in this model this model is an uncp model which is trained kind of like the way that Dolly 2 was trained as a diffusion model that's trained to invert clip embeddings so go from clip embeddings of images back to text and that lets us do similar things as the text model that that we used in all these prototyping I think a general principle if if you have one takeaway from this talk it's that when you're working with with these really complex models and kind of inscrutable pieces of data if you can get something into a thing that feels like it can fit in your hand that you can play with that you can concretely see and observe and interact with can be directly manipulated visualized all these things all the tools and prototypes that you can build around these things I think help us get a deeper understanding of how these models work and how we can improve them um and in that way I think models language models and image models generation generative models are a really interesting laboratory for knowledge for studying how these how these um different kinds of modalities can be represented and um Brett Victor said the purpose of a thinking medium is to bring thought outside the head to represent these Concepts in a form that can be seen with the senses and manipulate with the body in this way the medium is literally an extension of the mind and I think that's a great poetic way to kind of describe the the philosophy that I've approached a lot of my prototyping with so if you follow some of these principles and try to dig deeper in what the models are actually looking at build interfaces around them I think more human interfaces to knowledge are possible I'm really excited to see that future thank [Applause] you