CodeArt: Better Code Models by Attention Regularization When Symbols Are Lacking

TL;DR

This work enhances the vanilla tokenization and model architecture of a BERT model, construct and utilize attention masks, and introduces a new pre-training algorithm that substantially outperforms other general pre-training techniques of code understanding models.

Abstract

Transformer based code models have impressive performance in many software engineering tasks. However, their effectiveness degrades when symbols are missing or not informative. The reason is that the model may not learn to pay attention to the right correlations/contexts without the help of symbols. We propose a new method to pre-train general code models when symbols are lacking. We observe that in such cases, programs degenerate to something written in a very primitive language. We hence propose to use program analysis to extract contexts a priori (instead of relying on symbols and masked language modeling as in vanilla models). We then leverage a novel attention masking method to only allow the model attending to these contexts, e.g., bi-directional program dependence transitive closures and token co-occurrences. In the meantime, the inherent self-attention mechanism is utilized to learn which of the allowed attentions are more important compared to others. To realize the idea, we enhance the vanilla tokenization and model architecture of a BERT model, construct and utilize attention masks, and introduce a new pre-training algorithm. We pre-train this BERT-like model from scratch, using a dataset of 26 million stripped binary functions with explicit program dependence information extracted by our tool. We apply the model in three downstream tasks: binary similarity, type inference, and malware family classification. Our pre-trained model can improve the SOTAs in these tasks from 53% to 64%, 49% to 60%, and 74% to 94%, respectively. It also substantially outperforms other general pre-training techniques of code understanding models.

Figures (11)

• Figure 1: Code examples calculating statistics: (a) and (b) are equivalent but have different statement orders; (c) is buggy at line 11 with the wrong operation.
• Figure 2: Attention maps of variable mean at line 11 of Fig. \ref{['fig:moti-ex']}(a) (in red), by a source code model (a) and by a binary code model (b). In (a), we show attention between source code tokens, corresponding to lines 9--11 in Fig. \ref{['fig:moti-ex']}(a). In (b), we show attention between binary code tokens. For readability, we also include the corresponding source code and part of the assembly. The variable mean is stored in the register rax at line B0 (in red). A line with a darker color denotes a larger attention value. In an instruction, the first operand is destination and the second the source, and comment " rax<-i" means that register rax stores variable i.
• Figure 3: Attention maps by CodeArt for variable mean in the code snippets in Fig. \ref{['fig:moti-ex']} (a) and (b). The instructions in bold are in the dependence context of the variable and the blue arrows denote dependences.
• Figure 4: CodeArt encoder pipeline
• Figure 5: Regularized Multi-head Attention (with (b) illustrating a zoom-in view of part of (a))