How Codeium Breaks Through the Ceiling for Retrieval: Kevin Hou

Um, so my name is Kevin, and I'm going to be talking about how embeddings are stunting AI agents. Uh, so I'm going to let you in on some secrets about how we build the product, uh, and exactly what we're doing behind the scenes to improve your code gen experience.

So at Codeium, we are building AI developer tools, and we're starting with an IDE plugin. And as, uh, as mentioned before, we've been downloaded over a million and a half times. Uh, we're one of the top-rated extensions across the different marketplaces, and to reiterate, reiterate, we offer free, unlimited autocomplete, chat, and search across 70 different languages and 40 different IDEs.

So we plug into all the popular IDEs. Uh, we are the highest-rated developer tool, as voted in by developers in the most recent Stack Overflow survey. Uh, and you'll note that this is even higher than tools like ChatGPT and GitHub Copilot.

And importantly, we are trusted by Fortune 500s to deliver high-quality code that actually makes it into production. And we do this with top-grade security, licensing, attribution for some of the largest enterprises on the planet. Our goal at Codeium is to empower every developer to have superpowers both inside of the IDE and beyond.

And today, I'm going to let you in on some secrets about how we've been able to build a tool like this, and why, uh, users choose us over the other AI tools on the market. And the short answer is: context awareness.

So here's a quick overview about what context looks like today. Uh, we're all familiar, since we're at an AI conference, with the basics of retrieval augmented generation, the idea being that a user puts in a query, um, you accumulate context from a variety of different sources, you throw it into your LLM, and then you get a response, whether that be a code generation or a chat message.

Um, here's a concrete example about how retrieval can be used in code generation. So let's say we want to build a contact form in React. Um, now, you could go to ChatGPT, you could ask it to generate a contact form, but in reality, on a moderately large code base, this is really not going to work.

It's not going to give you, uh, things that are personalized to you. Uh, and this is really where context retrieval comes in. We need to build a contact form that, you know, is in line with our design system components.

Let's say you already have buttons, uh, and inputs. It has to be able to, uh, pattern match with, uh, with local, um, local instances of other forms inside of your code base. It has to ingest your style guide.

For example, if you're using Tailwind, you have to be able to detect and make the form look and feel like every other thing on your site. Uh, and then, of course, there's documentation both locally and externally, um, for packages and other dependencies.

So the question becomes: how do you collect and rank these items so that our code generation can be both fast and accurate for your use case? So to dive into a couple of different methods about how people are tackling this today, there's really three main pillars.

The first one is long context. So this is the idea that if you expand your prompt window in your LLM, it can read more input and therefore be a bit more personal to what you're trying to put, uh, what you're trying to generate.

This is very ergonomically easy to use,right? You just shove more items into your prompt. But this comes at the cost of latency, uh, latency and financial cost. So one of the most recent examples was Gemini. Um, Gemini actually takes 36 seconds to ingest 325k, uh, tokens.

To put this into perspective, a moderately sized or even small repo is easily over 1 million tokens. Uh, and that accounts for about 100k lines of code. So in this instance, most enterprises have over a billion tokens of code.

It's simply not feasible to be throwing everything into a long context model. The second method is fine-tuning. So for those that are familiar, fine-tuning is the idea of actually tweaking the weights of your model to reflect the distribution of the data that your consumer expects,right?

And so this requires continuous updates. It's rather expensive computationally. You have to have one model per customer, and it's honestly prohibitively expensive for most applications. And finally, we have embeddings. And for all of you, hopefully you're familiar, this is a relatively proven technology today.

Um, it's pretty inexpensive to compute and store. Uh, but the difficulty that we're about to dive into is that it is hard to reason over multiple items. It also has a low dimensional space, and I'll, I'll talk about that shortly.

So to dive deeper into embeddings, the whole concept is that you take your objects, you throw it through an embedding model, and then you end up with some sort of vector, some sort of array of numerical values, and this is in a fixed dimension.

And so by mapping and chunking code, we can map it to an embedding, and that allows us to quickly search over our functions, our documents, whatever you decide to chunk by, um, and this is what embedding search is called.

Uh, embedding search, like I said, is not a new concept. There's a bunch of models that have tried to optimize, and in this example, we're looking at, uh, one of the kind of north star, uh, eval benchmarks. Um, it's become increasingly popular, and the question becomes: how do we fit millions of lines of code into an LLM model so that we can actually generate useful results?

And so it's evident through the years that we're actually hitting a ceiling on what is possible using these traditional, uh, vector embeddings. And over time, even the biggest models, uh, are approximating to around the same level of performance.

As you can see, everything's kind of within plus or minus five. And, uh, at Codeium, we kind of believe that this is because fundamentally, we cannot distill all the, the dimension space of all possible questions, all possible English queries down into the embedding dimension space, uh, that our vectors are going to occupy.

And so at Codeium, we've thought very critically about what retrieval matters to us. Are we measuring theright things? And does semantic distance between these vectors really equate to things like function relevance in the concrete example that I showed earlier?

And so what we landed on is that benchmarks like the one that I showed you before heavily skew towards this idea of needle in a haystack. It's the idea that you can sift through a corpus of text and find some instance of something that is relevant to you.

Note, it is only one single needle. So in reality, code search requires multiple different needles,right? We showed that slide earlier. When you're building a contact form, you need all these different things in order to actually have a good generation.

And these benchmarks really don't touch that. And so we decided to use a different metric, and it's called recall 50. The idea, and its definition, is that, um, it's what fraction of your ground truth is in the top 50 items retrieved.

So the idea being now we have multiple documents, and we're now looking at the top 50 documents that we retrieved, how many of those are part of our ground truth set? So this is really helpful for understanding document, multi-document context, especially, again, for those large, large code bases.

And now we actually have to build a data set around this. And so this is where we did a little bit of ma a little bit of magic. We wanted to make the eval as close as possible to our end user distribution.

So we had to compile our own data set. So what we did, this is a PR that I put out, um, a few months ago. We looked at PRs like this. It's broken down into commits. Those commits, we can extract and actually match them with the modified files,right?

So now we have this mapping from something in English to a list of files that are relevant to that change. And you can imagine we can hash this in many different ways. But ultimately, the point I'm trying to make is we are creating a eval set that mimics our production usage of something like a code gen product.

And so this message serves as the backing for this new type of eval, where now we can run at scale this idea of product-led benchmarks. It gets us closer to the ground truth of what our users are actually experiencing and what retrieval tweaks in retrieval actually mean to the end product.

And so we threw some of the, uh, currently publicly available models at this notion of retrieval, this idea of using commit messages. And we found that there is reduced performance. Um, they're unable to reason over specifically code, but then also specifically this kind of real-world notion of, of English and, and commits,right?

And so at Codeium, we've been able to actually break through this ceiling. This is something that we've worked very hard at. We have to, to, to redefine exactly how we are approaching retrieval in order to be kind of in our class of our own so that when you are typing in your IDE, when you're chatting with our assistant, when you're generating autocompletes, we're retrieving the most relevant things that are, are for your, your intents.

So now the question becomes: how do we actually get this kind of best-in-class retrieval? And so I'm here to give you the very short and sweet answer, which is we throw more compute at it,right? But of course, that can't come with, uh, absurd, absurd, uh, uh, cost,right?

Financial cost. Uh, so how do we do this actually? In, in production, how do we actually do this without recurring an unreasonable cost? And so this goes back to a little bit of Codeium secret sauce,right? We are vertically integrated.

And what this means is that we train our own models. So number one, we train our own models. This means that these are custom to our own workflows. So when you're using our product, you're touching Codeium's models. Number two, we build our own custom infrastructure.

This is actually a very important point and actually connects to the whole Exafunction, to Codeium Pivot that we discussed earlier. Exafunction was a ML infrastructure, uh, company. And so what we've been able to do is build our own custom infrastructure down to the metal.

This means that our speed and efficiency is unmatched by any other competitor on the market so that we can serve more completions at a cheaper cost. And finally, we are product-driven, not research-driven. Now, what this means is we look at things like actual end user results.

When we actually ship a feature, we're looking at real-world usage, and we're always thinking about how does this impact the end user experience, not just some local benchmark tweaking. And so we could spend all day talking about, you know, kind of why Codeium has done this and yada yada, but that's a talk for a different time.

So I'm going to talk about something that I find very cool, and this is the reason why we've taken this vertical integration approach and been able to turn it into something that we call MQuery. So MQuery is this way of taking your query, so similar, it's that idea of taking your retrieval query.

You have your code base, and let's just say you have n different items. And because we own our own infrastructure and train our own models, we're now making parallel calls to an LLM to actually reason over each one of those items.

We're not looking at, uh, vectors. We're not looking at small dimension space. We're literally taking models and running them on each one of those items so that you can ensure you can imagine, you know, you run ChatGPT and tell it to say yes or no on, on an item, for example.

That is going to give you the highest quality, highest dimension space of reasoning. This leads into very, very high confidence ranking that we can then take into account things like your active files, your neighboring directories, your most recent commits, um, you know, what, what, what is the ticket that you're working on currently?

And we can compile all this to give you, you know, the top n, uh, documents that are relevant for your generation so that we can start streaming in higher quality generations, higher quality chat messages, uh, things of that nature.

And the reason behind this is, again, it's that vertical integration. It's that idea that our computation is one one hundredth of the cost of the competitors. We are not using APIs. And as a result, our customers and our users actually get 100x the amount of compute that they would on another product.

And so we're willing to do that. We're willing to spend more compute per user because it leads to a better experience.

And so, like I mentioned earlier, I lead our product engineering team. So we always want to anchor ourselves around these three different things. One, that we have to build a performant product. It has to be really fast. For those of you that have used the product, you can probably attest to this.

MQuery runs thousands of LLMs in parallel so that the user can start streaming in code within seconds, not minutes, not hours, seconds, and oftentimes milliseconds. It has to be powerful,right? None of this matters if the actual quality and the actual generations that you're building are wrong,right?

And finally, it has to be easy to use. We're building an end user product for people today that's in the IDE. Tomorrow, it might not be in the IDE. How do we actually build something that is intuitive to understand that people can grapple with and see exactly what my model is thinking?

And so because we have the benefit of distribution, uh, we were able to roll this out to a small percentage of our users. And by small percentage, we're dealing in the order of, you know, million plus downloads. This actually reached a surprising number of people.

And what we've been able to see is that, um, we were able to successfully reason over these thousands of files in people's monorepos and people's remote repos and select what was relevant,right? We can very accurately deem which files are relevant for the generation that you're trying to have.

And the result, as you can see, this is a real-time GIF, is both fast and accurate. So I'm asking for usage of an alert dialog. It's going through, and I think I've panned down here. Um, this is kind of a ShadCN component that I've modified internally.

We're, we're pulling in basically the source code of, of what is relevant for our generation. Um, and ultimately, the results of this experiment with that users were happy. They were thumbs they had more thumbs up on chat messages.

They were accepting more generations. And we were able to see that ultimately, we were writing more code for the user, which is the ultimate goal. It's that idea of how much value are we providing to our end users?

And so we built this, this context engine,right? This, this idea of MQuery, this idea of ingesting context and deciding what is relevant to your query to give you coding superpowers. And so our users will generate today, they're generating autocompletes.

They're generating chats, search messages. But in the future, they're going to generate documentation. They're going to generate commit messages, code reviews, uh, code scanning. They're going to take, you know, Figma artboards and convert them into compo uh, into UIs that were built by your own components.

The possibilities are endless. But what it starts with is this bedrock, this very hard problem of retrieval. And it brings us to, again, one of the reasons why Codeium is approaching this problem a little bit differently. Our iteration cycle starts with product-driven data and eval.

So we're starting with the end problem. We're building a product for millions of people. How do we start with what they're asking for? And how do we build a data set and eval system locally so that we can iterate on the metrics that matter?

Secondly, because we're vertically integrated, we're taking that massive amount of compute, and we're going to throw it at our users. You know, paying or not paying, we're going to throw it out at our users so that they can get the best product experience and the highest quality results.

And then finally, we're actually going to be able to push this out to our users in real time, overnight, and be able to get a pulse check on how this is going. You know, this is what we did for, for MQuery.

And when we evaluate in production, we can say, you know, thumbs up, thumbs down, and then hit the drawing board again, back to that same cycle, repetition. And so you can start seeing how these pieces of compounding technology come together,right?

We've alluded to some of them today, modeling, infrastructure, being able to retrieve. But then it also includes things like AST parsing, indexing massive amounts of repos, knowledge graphs, parsing documentation, looking at websites online. The list can go on and on and on.

But we're confident that we're solving these problems one piece at a time using that same iteration cycle, that same idea that we're going to take the, the distribution and knowledge that we have and that additional compute that we're willing to afford each user to solve each one of these puzzle pieces.

And, um, I want to leave you with, uh, uh, a parallel analogy. So in my past life, I had experience in the autonomous driving industry. So to bring over a metaphor from that industry, in 2015, TechCrunch boldly predicted that that was going to be the year of the self-driving vehicle.

Uh, it was largely, uh, you know, now we're in 2024, so we can look back in hindsight, largely untrue,right? We were doing things like sensor fusion. We were decreasing our polling rates. We were running offboard models, all this in the effort of making heuristics that would compensate for the lack of compute that was available because consumer graphics cards were not as popular or not as, uh, powerful as they are today.

Fast forward today, we're seeing 100x the amount of compute available to a vehicle. You can take a Waymo around San Francisco, which I encourage you to do. It's a wonderful experience. Um, but that means that we're actually able to throw larger models at these problems,right?

More sensors, higher frequency. And now, 2024, TechCrunch has released another article that said, "Will 2024 finally be the year of the self-driving vehicle?" And we can now look at this pattern and say, "Driving performance was substantially better by throwing larger models, being able to handle more and more data."

And so at Codeium, we believe that this embedding-based retrieval is the heuristic. We should be planning for AI-first products, throwing large models at these, at these problems so that AI is a first-class citizen. We're planning for the future.

And finally, we also believe that ideas are cheap. You know, I could sit up here and tell you all these different ideas about how, you know, we're going to transform coding and the way that the, the, the, the theory behind, uh, possible solutions.

But what we believe at Codeium is that actually shipping, actually showcasing this technology through a product is the best way to go. And so if you agree with these beliefs, you can come join our team. We're based in San Francisco, and you can download our extension.

It's free. I'm not obviously, uh, uh, what's it called? I'm not advertising, uh, the core product nearly as much. We're kind of talking about the technology. But you can experience this technology firsthand today by downloading our extension. It's available on all the different plugins, uh, VS Code, JetBrains, Vim, Emacs.

Uh, and you can see how this infrastructure and the way that we've approached product development has shaped the experience for you as a user. And then, of course, you can reach out to me on Twitter. Uh, I put my handle up there.

I'll be kind of floating around outside. So if you have other questions or are interested in what I had to say, um, but I hope that you learned something today. I hope that, you know, you use Codeium, you try it out, and see what the magic can do for yourself.

Thank you.