The five previous articles explained the five vulnerabilities. Now they become training material.
Before starting, a quick review of an idea from the series Teaching a machine to spot vulnerable code, because everything here depends on it.
Training is showing examples with the answer attached
A machine learning model does not receive rules. It receives examples where somebody already wrote the correct answer, and it adjusts its own numbers until it gets most of them right.
Each example has two parts: the content (here, a snippet of code) and the label (the correct answer, such as "this is SQL Injection"). A large collection of those pairs is what we call a dataset.
So the first practical question is: where do hundreds of code snippets with the answer already written come from?
The problem of getting examples
Writing five hundred snippets by hand, one at a time, would take weeks. And copying code from real projects has two problems: you do not know whether that code is vulnerable (which is exactly what you want to find out) and you would have to review all of it manually anyway.
The way out was to build the examples from moulds.
A mould is an example with holes
A mould is a snippet of code where some pieces were replaced by markers. Like this:
var sql = "SELECT * FROM {table} WHERE {field} = '" + {input} + "'";
using var command = new SqlCommand(sql, connection);The pieces between braces are the holes. Each mould comes with its label already decided, because whoever wrote the mould knows whether that pattern is safe.
Then you just fill the holes with different combinations. One combination fills in Users, Name and userName. Another fills in Orders, Reference and request.Email. Each combination generates a new example, with the same label as the mould.
In the project there are eight combinations per mould. One mould becomes eight examples.
Always in pairs: the right and the wrong side by side
Here is the most important decision in the whole dataset.
For every vulnerable mould there is a mould of the same feature done the right way. The concatenated query has the parameterized query next to it. The HTML without encoding has the HTML with encoding next to it.
The reason is to stop the model from learning the wrong thing. If you only showed vulnerable examples full of SELECT, the model would learn that the word SELECT is dangerous. And then it would flag every database query, including the correct ones.
Showing both sides, the word SELECT appears under both labels. It stops being a useful signal, and the model is forced to look at what actually differs: the data glued inside the text versus the data passed outside it.
Always in four shapes: the same pattern wrapped in different ways
There is a second decision as important as the pair, and far less obvious.
A mould, on its own, is a short fragment. Three lines of a database query, loose, with nothing around them. Except real code never shows up like that. It shows up inside a method, which is inside a class, with a route attribute on top and a return at the end.
That matters because the model does not look only at the dangerous line. It looks at the whole snippet, and the shape of the snippet is part of what it learns. If every training example is a short loose fragment, the model learns two things at once: the dangerous pattern, and that suspicious code looks like that. Take the second away and it no longer recognizes the first.
So every mould is rendered in four wrappings:
- loose, the fragment on its own
- inside a method
- inside an endpoint, with the
[HttpGet]on top - inside a full controller, with
[ApiController]and[Route]
The wrapping is the same for the vulnerable mould and for its safe pair, for the same reason as the SELECT that appears on both sides. If [HttpGet] only showed up in the vulnerable examples, it would become a danger signal. Showing up in both, it stops saying anything, and what is left for the model to look at is still the body.
I arrived at these four shapes by taking a careful look at what the detector actually receives in practice, with it already built and running. The next article shows that analysis and how much it changed the result.
The six labels
The detector has six possible answers, not five. They are the five vulnerabilities plus the Safe label, which means "I found nothing here".
And Safe receives two kinds of example. The first is the safe counterpart of each vulnerability, which I just described. The second is ordinary code, with nothing to do with any of the five categories: registering dependencies, a response record, an Entity Framework query, a simple controller.
That second group matters because, in a real pull request, the overwhelming majority of lines are ordinary code. If the model had never seen ordinary code in training, it would spend all day trying to fit every line into one of the five categories.
Adding it all up: 98 moulds, eight combinations each, in four shapes. That makes 3136 examples.
Now the training: what exactly gets adjusted
The detector is built on top of CodeBERT, which the previous series introduced. Worth recovering the idea in one sentence: CodeBERT is a model that was already trained for a long time, reading millions of snippets of code, and because of that it already knows something about how code is written.
That ready made model has two parts, and the difference between them decides the whole project.
The first is the body. That is where everything it learned by reading code lives. It is 125 million adjusted numbers, and a file of almost 500 megabytes.
The second is the classification head. It is a small layer placed on top of the body, whose job is to take what the body understood about the snippet and turn it into six scores, one per label. It is born blank and knows nothing.
Freezing the body: training only the head
There are two roads from here.
The first is adjusting the whole model, body and head. It gives more quality, and it has two costs: it needs a good graphics card to avoid taking hours, and the result is a 500 megabyte file.
That second cost is a concrete problem. GitHub refuses files above 100 megabytes. A 500 megabyte model simply does not go into the repository.
The second road is freezing the body. Freezing means marking those 125 million numbers as not adjustable: training reads the body, uses what it knows, and changes nothing in it. Only the head is adjusted.
The trained head gives a file of 2.3 megabytes. That fits in the repository with no effort. The body, which did not change, remains the original public file, and the tool downloads it on first run.
NeuralGuard takes the second road.
The trick that makes training take minutes
Freezing the body opens an optimization that changes the scale of the training.
Normally each training round passes every example through the whole model. Since the body is enormous, that is slow, and you repeat it every round.
But if the body never changes, what it produces for an example is always the same. Passing the same snippet through the body on round 1 and on round 300 gives exactly the same result.
So you can pass each example through the body once, store the result, and from then on train the head on top of those stored results.
In practice: the pass through the body takes a few minutes and happens once. After that, three hundred rounds of training the head take seconds, because the head is tiny.
In code, training the head is the same three step ritual from the previous series:
var optimizer = torch.optim.AdamW(model.HeadParameters(), lr: LearningRate);
var lossFunction = CrossEntropyLoss();
for (var epoch = 1; epoch <= Epochs; epoch++)
{
var logits = model.ClassifyFeatures(training.Features);
var loss = lossFunction.forward(logits, training.Labels);
optimizer.zero_grad();
loss.backward();
optimizer.step();
}The training.Features is precisely the stored result of the pass through the body. That is why the body does not appear in that loop.
Separating the exam
One piece is missing: how to know whether the model learned or just memorized.
The rule, which the recommender series covered in detail, is to always set aside a portion of the examples and not show that portion during training. At the end, the model answers those questions it never saw, and its score is measured there.
That is what I did: I set aside a portion of the examples for the evaluation and trained with the rest.
The first result
I trained the two models available in the project, CodeBERT and also the small transformer the previous series built from scratch, in order to compare.
CodeBERT got 80.9% right.
The from scratch transformer got 73.0%.
I wrote both down and moved on, because they make sense. The big model ahead of the small one, which is what you would expect, and both at a plausible level for a detector put together in an afternoon. There was nothing there asking for a second look.
Both were wrong.
The next article is the investigation of that number: where the defect was, why it is easy to commit without noticing, and what the real scores of the two models are after the repair.