Smart Contracts to Embeddings: Using off-the-shelf LLMs for Fun and Profit

DeFi Security Summit 2024 · 9 Nov 2024 · 24:11

Markella Gioka of Dedaub shows how pre-trained large language models can turn smart contract code into embeddings for similarity search. Finding contracts that resemble a vulnerable pattern, a known hacker's deployments or other MEV bots all become distance queries between vectors. She stresses matching the input to what the model was trained on and using approximate nearest-neighbour search instead of brute-force comparison, as part of Dedaub's smart contract analysis tooling.

Speakers

Topics

From DeFi Security Summit 2024, hosted by DeFi Security Summit.

Transcript

Uh, hello everyone. Uh, do I sound okay or do I have to speak louder? Good. Okay. So, today I'll be talking about smart contract to embeddings. Uh, what is that? Why do that? And how we did it in case you're interested. So, first I want to start with this code from Yiannis Malagardis' GitHub. You may know him as my boss. Uh, it was from uh, DSS actually last year and it goes something like Web3 is a perfect playground for security tooling. It's the only opportunity one ever gets to have all the code open, all the state, all the past executions permanently recorded. And I wanted to

include that because that is true and it's a unique feature of the Web3 domain. Like you have literally everything that's happening open and if you know how to use that effectively, all the uh, real-time data, you can do some pretty incredibly incredible security tooling. But that is also true for machine learning. You have all this information, all this data and that's the perfect conditions to do machine learning applications. That's uh, the main thing uh, machine learning uh, researchers are looking for and it's a shame not to use that in this domain.

So, I'll start with a little bit of background so you all don't think I'm speaking Greek. But I'll try to keep it high-level so uh, you don't have to understand uh, everything I'm talking about and don't worry, there won't be any math. Uh, so the main thing I'll be discussing today is how to do similarity, smart contract similarity. We won't be talking about any other applications. So, that's the reason we're talking about embeddings. Uh, the good thing about embeddings is that you can use them for similarity. So, if you have two code snippets, you and you can create some good embeddings out of those uh code snippets, and you get the distance uh

between those vectors, that's a good metric for similarity. That, of course, depends on the quality of your embeddings. You may ask, "Why do I care about similarity?" That's a good question. I'm sure you can think of uh other applications as well, but let's say you have a vulnerable contract with a specific code pattern, and you want to find similar ones. Uh you may want to find similar contracts from a specific hacker, uh identify other MEV bots, and so on. Uh this is an overview of what we'll talk about today. So, uh first of all, uh I'd like to say that LLMs uh I know, LLMs and blockchains uh bunch of buzzwords, but they're useful. That's

why they're buzzwords. So, there's a lot of uh embedded knowledge in pre-trained LLMs. So, if you get take something out of this talk, is that you can do a lot of applications with very little effort, and that's a good thing. Uh the second thing is that you need to do it correctly in order to get results. So, one uh very important thing is to adjust, for example, your input to match what a model may have been trained on. Otherwise, you're not going to get uh much out of it. And finally, okay, you get all those embeddings or whatever, uh and you may want to find, for example, which have the smallest distance.

Don't do it uh brute force. Don't do all of the uh calculations. That's going to take forever. Use something like an approximate uh closest neighbors algorithm. And the good news is that you can probably get this for free. So TLDR embeddings high-level is just a way to represent data uh as a vector of numbers, as a list of numbers. That's code in our case, but it can be pictures, can be whatever. The simplest form of embeddings, you may all have heard of it, is one hot one hot vectors. So, let's say we have 15 bits.

Whenever one bit is one, that's a different uh entity, that's a different thing. But, as you can tell, that's not a good embedding. Good embeddings have to be semantically meaningful. What that means is, for example, if we use the example of word embeddings uh that's a classic example in word to vec and glove, where if you get the vector for king and you add the vector for woman, you get the vector for queen. Uh that means that embeddings have the this uh feature that they embed relationships within those numbers.

They may not always be linear like this, but uh that's the idea. Finally uh it means that similar entities will be closer in the vector space. So, for example, if you take the distance between king and queen, uh that will be smaller than the distance between king and rat, for example. Next, uh you should know that not all embeddings uh Sorry, good embeddings are learned. That's a different uh slide. So, uh what that means is that you need to do training to get good embeddings. There's no magic algorithm that's going to give it to you. Uh you need the statistical

properties of a big corpus. Uh as it It's a special code in the NLP world, you shall know a word by the company it keeps. That means that words that appear in similar contexts will have similar meaning. So, however, uh not all uh embeddings are created equal. So, we were speaking so far about word embeddings. So, if you were to embed those two words, those two sen- those words in those two sentences, uh I went to the bank to deposit money and I went to the river bank, you will get uh the same embedding for the word bank, but it is clear that in those two

sentences, it has a very different uh meaning. So, what we want are contextualized embeddings. So, that's where transformers and LLMs that you've already heard of uh come into play. High-level transformers have the attention mechanism. That's probably what you've heard about them. And the high-level idea is that it gives you access to all of the input. It gives you context. So, if you get the embedding of a word from uh transformer, it will have the meaning of the sentence kind of embedded into it.

For example, uh in BERT, that's one of the uh initial uh transformer models, if you uh do the classic training task of uh masked language, which is uh you mask part of the input or a word of the input and you tell the model to guess what that word is, uh the embeddings that the transformer creates internally for each word uh will be forced to learn the contextual information of the sentence. So, that's how you get uh contextualized embeddings.

You may be wondering why I'm talking about word embeddings. Well, if we treat uh code as a set of words, that uh that's what we want to do. So, we have a set of embeddings. How do we embed a whole sentence, for example? One way to do that it was with BERT model, for example. Uh what they did was add an extra token in the beginning of a sentence which then later the embeddings of the of the model will get trained on a sentiment analysis task. So, if is the sentence positive or negative?

What that achieves is that it teaches kind of the model to encode in this embedding the meaning of the whole sentence. So, in a way that's how we get our first sentence embedding. That's not the best way to do that. That's just an initial approach. A better way to do that is with Siamese networks. Siamese networks is just you take the same model twice. You embed two sentences that you know are similar. You take the the embedding of those sentences.

And then you calculate the distance of those two and use the distance in the loss function. What that means is that the model learns how to minimize the distance. So, similar sentences will have more similar embeddings. And now we didn't use Siamese networks because of Thailand. Um so, what are some of the challenges in creating smart contracts? Creating embeddings out of smart contracts. Well, first of all, there are too many of them. So, you have to keep that in mind. Just in Ethereum, there are tens of thousands of contracts deployed every day to with in total to the millions of

thousands of contracts. So, your search has to to quick and your embedding creation also has to be quick. Another problem is that we don't have source. Maybe that's not entirely accurate. We have source for about 60% of contracts, but still a 40% is left without source. So, doing similarity on the source code is not very reliable. And most interesting applications like malicious contracts or maybe MEV bots, whatever, don't really have an incentive to publicize their source. So, there's even more uh

reason to do so. Some other challenges are on-chain data, which are bytecode and assembly code, is too low-level, so it won't have very um uh in- intricate information about the contract we want to embed. We have to take into account that if we're using a lamps, they have specific uh a specified size, so you can um use very big uh inputs anyway. And finally, uh even though we do have a lot of data available, we don't have a lot of labeled data. That takes a lot of manual effort. So, we either have to use

out of the box approaches from other domains, which means uh pre-trained models, or we have to creatively use what data we do have. Okay, so then what are some of the things that we can do? Well, first of all, the compiled uh code is a good uh middle ground approach. It only depends on um on-chain data on bytecode. It retrieves high-level information like uh array operations, uh function calls, loops, if-else statements, stuff like that. Uh

and it looks like other object-oriented languages. So, you may take advantage of models that are already trained on code. Another thing, another revelation we had is that a public function of a smart contract can be considered as a sentence of the contract. So, it encapsulates a part of the functionality of a smart contract. So, you can use that as your building block to do similarity. And if you combine that with a tree second algorithm, which just means include underneath all of the private private function calls uh the function definitions of the private

function calls of the public function, you have all of the information that's uh encoded in that function. Our first approach is to use that as your input for a sentence encoder. That's going to give you a pretty good encode embedding because it has been trained to produce that. It may not be the best, like if you do a more fine-tuned approach, if you have, for example, a data set of pairs or whatever, that's going to be much better.

But uh for example, in this case, because the code is transformed to something that the model has seen something similar before, so high-level code which most LLMs have been trained on some extent during their initial training, are you going to be able to benefit from transfer learning, which is when a model is able to perform well on a task that it hasn't been trained on, but is uh similar.

Another idea, uh what if we use some of the data that we do have available in the blockchain? So, one thing that we do know is that public functions are associated with a function selector. What that is is just a hash of the function signature. That function selector is then used in a function selector piece of code in the bytecode, which uses that hash to call the appropriate public function. In that way, you can build a dataset of public function uh selectors and names, and then reverse engineer the names for contracts that don't even have verified source. That

allows us allowed us, at least, to build a labeled dataset of decompiled code and name of uh this function in this decompiled code. As you can imagine, uh having the name of a function tells you a lot about a function. If I tell you these two functions are uh have the same name, you're going to guess they do similar things. So, we thought well, what if you train an LLM to uh predict uh what the name of a public function is going to be based on the decompiled code, which already has a lot of information about it.

And then use the internal state of that LLM as an embedding. Turns out that is a pretty good uh sentence encoding, sentence uh as in public function. And the the reason I mentioned earlier the CLS pooling is because it's kind of the similar uh intuition. It's asking the model to distill all this information to a descriptive function name. Finally, uh how you're going to actually run that.

So, you have all your code, you create all the embeddings, and you want to run a query. How do you do that? Well, don't use uh classic indexing. Uh in indexing algorithms for classic databases are not meant to be used on uh vectors. Use a vector store. The reason vector stores are better is because embeddings can be seen as uh nodes in a space, and the distance between those nodes uh is are the edges of those nodes anyway.

So, you kind of create a graph that enables uh vector stores to use graph type of searches to do, for example, approximate approximate nearest neighbor, and find uh closest embeddings much quicker than just doing all the distances with all the embeddings in your data set. Also, the good news is that whatever data set you're use uh database you're using right now already supports some kind of uh vector extension. Uh all the major ones do anyway. And if you're building something from the ground up, you can maybe consider using uh vector focused uh database, which uh some These are some of the examples, but

there are many more. So, uh we did this. How did it work? We deployed it uh September this year. Uh we support most major uh chains that have EVM bytecode, and you can query any public function, and you can get results in on average less than 3 seconds. That's like the worst case I've uh seen. Funny story. Uh this past fall, on our company retreat uh a few like weeks or something after we deployed our feature, we had the chance to use it in real in real life.

Uh when we discovered the vulnerability, uh fun story, we were doing a competition between amongst our company uh for whoever finds a bug will win a PS5. Uh we ended up founding uh finding the Bedrock vulnerability, uh which as you might have heard was an infinite mint vulnerability, and that meant that 50 75 millions in dollars uh were at stake. The code that had this vulnerability was deployed on eight eight different chains, and at least in one of those chains the code wasn't even verified.

Uh that enabled us to try our tool uh in real time to find all the deployed code and ensure there weren't any other uh forks or pieces of code that use similar logic and had the similar vulnerability in the wild. If you want to try it, uh it's live now on our app on our site, app.dedaub.com. If you go on any contract, on any public function, we have a little button that says find similar, and if you click that, you will get something like this that shows you the similar functions and the addresses.

And let us know if you have any feedback. Ah, sorry, there's a last one. Almost forgot about that. So, uh what we talked about today, there's a lot of information encoded in pre-trained LLMs, so take advantage of that, it's free. Uh representation matters. You have to think about what your input to your model looks like. And finally, don't use uh brute force to do your comparisons. There are much, uh, better things that are for free to use.

And that is all. I think I'm a bit early. All right, any questions? So, uh, thanks for your talk. Uh, can you use this to, let's say you find a bug, you define it in some way, and then you can try to find a similar type of bug in different chains? Is this what you're doing?

Uh, so, that's a good question, and that's actually a feature that I would like us to implement at some point. It wouldn't be hard, it would just be you would have to submit that code, and it would be the same as doing a query as we have it right now, but at this point, no, we don't have that. You just, you can find similar code with already deployed code. Uh, our site is also a block explorer, so you can go to any, uh, deployed contract. So, you can only use it on deployed contracts at this point. Right, right, but like you, you, if I find a piece of code, and let's say it has a bug, and I want to find similar code, I can do that?

Uh, yeah, you can use the find similar. Uh, we support it on public functions. And if there is a public function that, uh, has a similar coding pattern, that's kind of, uh, clone, I guess, of this, uh, function, there's a good chance you will find it. And it queries Yeah, I've actually used it. It queries, uh, like the blockchains you support now, and maybe you're going to add support, or what, which blockchains does it query? Uh so I, at this point, I think we sup I don't remember all of the EVM-based uh chains we support. We can You can go on our site and see and like check. Uh we definitely have Ethereum, Binance, and Optimism.

Uh right now the feature doesn't work cross-chain, but that's a feature that's coming soon. Okay, thank you. Uh you said the database is available. Is it open source? Can we download it? Uh it's not ours. Uh most databases have extensions that are free, yeah? And uh if you have an open-source database, the extension's also open source.

I don't know if asked. Hey. Um thank you for the talk. Um I was wondering if you tried embedding not only the code snippet, but also maybe an LLM-based like LLM-generated summary of the code. Because I don't know how how LLMs work with code instead of like natural language. Um yeah. Uh that could be a good idea. Then there's a question of if an LLM could create a good uh summary of the code, maybe it can create a good embedding for the code.

So, it depends on the kind of code you will use. If it's uh something that model has been trained on, maybe it will do well. Uh but they seem kind of like similar tasks what you're explaining. So, if you can do one thing well, I don't see why it wouldn't be able to do the other one as well. And yeah, last one. Um you said you fine-tuned a model. Uh is that a like compared to using OpenAI embedding model, is there a Uh yeah. like a super good advantage or like it's like a slightly better?

Well, we don't use open I mean we have used open AI chat GPT. That's on projects I think. It creates a summary but we didn't use it for this. It's not super cost effective and also yeah, you can do a lot more things with your own model. You can create more lightweight models. You can do all sorts of things. So it's I'd say preferable to using that but of course it has its applications as well.

All right, I think that's all for questions then. I think there's a short break now and the talks are going to start again at 4:40. Okay.