| 4 min read
Table of contents
As we have stated over and over in the past, the most critical step in
our ongoing project of building a machine learning (ML) based code
classifier will be that of representing the code as vectors. In our
last article, we reviewed how this is done for
natural language. We looked at simple, though inconvenient methods, such
as one-hot and categorical encoding, which we actually used in our
first classifier attempt. We also took a
glimpse at the state of the art in vector representation of language,
which is based on neural networks, called word2vec.
Recall that this is a neural network which attempts to predict neighboring words from the central word. Its dataset is made up from a corpus, and the results will be very dependent on the sense of this text. While the task is certainly useful, we don’t care so much about the predictions, but rather about the weights which make up the network, which give us an intermediate representation. These are the vectors we are looking for to represent each word:
The vectors are implicit in the middle layer.
In turn, word2vec is based on a more simple task: predict a vector
from that same vector. Well, isn’t that just an identity function,
mapping every object to itself? As it turns out, no, given that the
weights in a neural network are randomly or arbitrarily initialized and
will be optimized to the task in an iterative process. For this task,
with a single hidden layer, we get something similar to word2vec,
called an autoencoder:
Autoencoder neural network via Carnegie Mellon.
Autoencoders turned out to be a foundational idea in neural network based dimension reduction, the task of representing high-dimensional objects (such as one-hot encoded words) as lower-dimensional vectors, while still retaining most of the useful information contained in it. Before that, most methods relied on matrix decompositions, but these tend to be computationally expensive and not scalable, thus not fit for representation of large codebases or natural language corpora.
Keeping these two neural network based ideas, namely that vector
representations can be obtained from middle layers in networks designed
for different, though seemingly frivolous tasks, one can begin to
understand how vector representations of code might come to be. What
will be the goal? Predicting the name of a function. What will be the
input data? Not too surprisingly, it will be the code of the function
whose name we would like to predict. In what form? That’s where the
waters get a little murky, since there are so many ways to structure
code, and so many representations to extract from it. Our readers might
already be familiar with the Abstract Syntax Tree, Version Control
(v.g., git) history and the Control flow and program dependence
graphs. One can even simply choose not to
represent anything: consider the code as sequence of words without
exploiting its syntax, and represent them as bags of words, as we did in
our previous run. One could also look at
the meta-code: metrics
such as modified lines per commit, code churn, cyclomatic complexity,
all of those could be thought of as possible candidates for inputs to a
neural code classifier.
One of the most apt of such representations is the Abstract Syntax Tree
(AST), which is universal in the sense that it can be taken out of
every language in existence, and could potentially be standardized so as
to eliminate the language barrier. Indeed, this is the input
representation chosen by the code2vec authors. More specifically, they
sample some paths in the AST at random. The labels in the training
phase, which would be the prediction targets later, are the function
names. The objective is to predict meaningful function names from the
function’s content. If the body is return input1 + input2, it would
seem obvious to a human to call that function add or even give_sum.
However, this is not the case, as there are developers who give strange
names to their identifiers and methods, such as
perform_binary_operation1. Maybe to make their code hilariously
unmaintainable
and keep their jobs forever, since nobody else would understand it; or
just make it sound like it does more than in reality.
The network itself to predict the function names from the randomly-taken
paths in their abstract syntax trees is a little more complicated than
an autoencoder or a single-layer word2vec network:
code2vec architecture. From Alon et al. (2018).
Notice that in this diagram, unlike most other seen here so far, the layers are the arrows, thus the objects are the intermediate products obtained after passing through the layers. Of course, we will be most interested in the green circles above, which are the code vectors we are looking for, and not so much in the final predictions. Not that the task of predicting function names is uninteresting, but it is not related to our interests. We only want the code vectors in order to pass them on to our classifier to determine the likelihood of containing security vulnerabilities.
Fortunately enough,
code2vec offers pre-trained models
and others that can be further trained.
So on my to-do list,
the top priority now is to figure out how this works in detail
and how to remove the final layer,
the one that gives the prediction,
and just keep the vectors.
If that sounds interesting,
stay tuned to our blog.
By the way, have you read our main blog post on secure code review?
References
-
U. Alon, M. Zilberstein, O. Levy, and E. Yahav. code2vec: Learning Distributed Representations of Code Proc. ACM Program. Lang., Vol. 3, No. POPL, Article 40. January 2019.
-
Z. Chen and M. Monperrus. A literature study of embeddings on source code. arXiv.
Table of contents
Share
Recommended blog posts
You might be interested in the following related posts.
