CodeBERT and the AIs that read code

This is the fourth article in the series where we are going to teach a machine to find security flaws in code, all in .NET. Recapping the trip so far: machines learn from examples, not rules (article 1); transformers read text understanding context, and the BERT family is the one of the "readers", which interpret and classify (article 2); and a vulnerability is a repeating pattern, almost always born from too much trust in user input (article 3).

Now the two ends meet. There are transformers trained specifically to read source code, and in this article you will meet the main ones, understand the difference between them and find out which one will be the brain of our detector.

Pre-trained model: the professional who arrives fully educated

Before introducing the family, a concept that changes everything: the pre-trained model.

Training a transformer from scratch costs millions: weeks of processing on thousands of graphics cards, swallowing absurd amounts of data. Neither you nor I will ever do that, and the good news is nobody needs to. The research labs do that heavy training once and publish the result for anyone to use.

It works in two stages, and the analogy here is a doctor's education. The pre-training is medical school: six years learning the entire human body, with no focus on any specialty. The fine-tuning is the residency: it takes that already-educated professional and specializes them in cardiology in far less time, because the foundation already exists.

With models it is the same. Pre-training teaches the model to "speak code" in general. Fine-tuning takes that model and specializes it in a task: looking for defects, summarizing functions, whatever you want. And fine-tuning is radically cheaper: where pre-training costs millions, the fine adjustment costs hours on an ordinary machine. That division is what makes this series possible.

Where these AIs went to school

And what was the "college" of these models? The open source code of the world. There are millions of public projects on GitHub, and researchers built giant study collections out of them, like CodeSearchNet: millions of functions, in several languages, each accompanied by the documentation written by the programmer explaining what it does.

That is precious, because it gives the model both sides of the coin: the code and the human explanation of the code, side by side. It is by studying these pairs that the model learns to link "what the code does" with "how humans describe what it does".

CodeBERT: the pioneer

In 2020, Microsoft researchers introduced CodeBERT, the first big success of the idea "take BERT's recipe and apply it to code". It was pre-trained on six languages (Python, Java, JavaScript, PHP, Ruby and Go), reading code the same way BERT reads a sentence: all at once, with each piece paying attention to the pieces that matter to it.

The result is a model that "feels" code. It notices when a variable name doesn't match what it holds, when a function stands out from the pattern, when two snippets do the same thing written in different ways. And being from the family of readers, it is great at exactly what we need: reading a snippet and issuing a verdict.

One detail that will matter later: CodeBERT is small by today's standards. While the famous chat models have hundreds of billions of parameters, it has 125 million. It sounds like a lot, but in practice it means it runs on an ordinary machine, with no special graphics card, no expensive cloud. Keep that in mind.

GraphCodeBERT: the one that understands the flow

A year later came GraphCodeBERT, from the same group, attacking a limitation of its older brother: CodeBERT reads code as running text, without knowing that code has structure. It sees the words, but doesn't see that the variable on line 2 is the same one showing up on line 9.

GraphCodeBERT sees it. Besides the text, it receives the data flow of the code: a map of "who feeds whom", showing where each value comes from and where it goes. It is the difference between reading a recipe as any other text and understanding that the beaten egg in the small bowl is the same one that goes into the dough later.

For security, that is a feast, because many vulnerabilities are exactly a flow problem: the data that came from the user and reached the database without any cleaning along the way. On the other hand, preparing that flow map takes more work: the code needs to be analyzed and transformed before entering the model, which makes the whole process more complex.

CodeT5: the one that reads and also writes

Leaving Microsoft, Salesforce created CodeT5, which follows another recipe: instead of only reading (like BERT), it reads and writes. It is a two-part model, one that understands the input and another that generates an output, which makes it great for transformation tasks: summarizing a function in one sentence, translating code from one language to another, suggesting fixes.

It is a powerful family, but notice that its shape doesn't match our mission. We don't want to generate anything, we want a verdict: vulnerable or safe. Using a model that writes for a classification task is like hiring an award-winning writer to answer yes-or-no questions. He can do it, but he is not the tool designed for it.

UniXcoder: the Swiss army knife

Closing the family, UniXcoder (Microsoft again, 2022) tries to unify everything: a single model that can operate as a reader, as a writer, or as both, depending on how you call it. It also takes advantage of the code's structure information during its study phase. It is the most versatile of the group, and a strong candidate in any code project, at the cost of being the most recent and having less ready-made material in circulation.

What about the giant models, like ChatGPT's?

The question is inevitable, so I'll answer it quickly. The giant chat models read code very well, but remember the distinction from article 2: they are from the writers' family, trained to continue text. Besides, they are too heavy to run on your machine and, in general, they require sending your code to someone else's cloud, which not every company can or wants to do. For the task of "read and classify, fast, locally and for free", the small specialized readers remain unbeatable. And that is our task.

The chosen one

So, recapping the team: CodeBERT reads code as text, GraphCodeBERT reads code and the flow of its data, CodeT5 reads and writes, UniXcoder does a bit of everything.

For our detector, the chosen one is CodeBERT. And the decision follows the same architecture logic I always apply: it is not the most advanced model on the list, it is the one that best fits the problem. Our task is pure classification, so the readers' family is the right one. Within it, CodeBERT is the most established, the best documented, the one with the most material and community-published specialized models, and the simplest to get running. GraphCodeBERT is technically superior for security, and stays noted as a future evolution of the project, but starting with it would be adding complexity before having the basics working. That order never goes well.

And there is the detail I asked you to keep in mind: 125 million parameters run on any machine. Our detector will not need a graphics card, will not need an expensive server, will not send code to anyone's cloud. It will be a model file running inside a .NET program, on your machine.

Conclusion

Recapping:

A pre-trained model is the professional who already left college: the expensive training has already been paid for by the labs, we only specialize it (fine-tuning)
These AIs studied on the open source code of the world, reading millions of functions with their documentation
CodeBERT is the pioneering reader; GraphCodeBERT adds the data flow; CodeT5 reads and writes; UniXcoder does a bit of everything
The giant chat models are writers: heavy, costly and unnecessary for a local classification task
The chosen one is CodeBERT: the right fit for the problem, mature, light and ready to run anywhere

In the next article the theory ends for good: we are going to use CodeBERT, specialize it in recognizing vulnerable code and put it to work inside .NET with TorchSharp, all in C#. It is the hands-on moment this whole series has been preparing.

Blog

Series