Nova: Generative Language Models for Assembly Code with Hierarchical Attention and Contrastive Learning

We thank the reviewer for the insightful review and questions.

1. Improvement on BCSD

We acknowledge that Nova’s improvement on BCSD is not as significant as on BCD, especially given that Nova has a larger model size.

Yet, binary code similarity is a much better-explored domain than binary code decompilation. There has been more related work from software engineering, and security domain design models for it. The baselines we compare (jTrans, DiEmph, and CodeArt) have already achieved impressive performance. By contrast, end-to-end binary code decompilation is a very new task and LLM4Decompile is the only existing SOTA baseline we can find. Thus, we think there is less space to improve on binary code similarity than binary code decompilation, and that Nova outperforming existing SOTAs is valuable and hard.

Besides, those specialized BCSD models we compare with are encoder transformers. The bidirectional attention benefits their performance on the BCSD task as it is essentially an encoding task. Using decoder-only LLM for encoding tasks such as the BCSD task is less common and harder. Our ablation study (in Appendix Table 12-15) also shows that without Nova’s designs, simply training DeepSeek-Coder-1B cannot outperform existing SOTAs on BCSD.

Then why do we explore Nova (a decoder-only LLM) on BCSD? Given that decoder-only LLM is getting much more popular and promising, and the latest advanced LLMs are almost all decoder-only (Llama, StarCoder, etc), Nova shows the potential of building on top of the latest decoder-only LLMs for encoding tasks such as BCSD so that we can make good use of those pre-trained decoder-only LLMs.

2. Table 3

In Table 3, Nova’s fine-tuning means applying hierarchical attention and using contrastive learning objectives. We want to show that, our approach brings improvement over standard fine-tuning on multiple backbones (both DeepSeek-Coder 1.3B and 6,7B, as well as Granite-3B-Code).

Maybe it is clearer to call it “Granite + Nova’s Approaches”.

3. Token Length

In the HumanEval-Decompile benchmark, the longest assembly function has 3240 tokens, and the average is 698 tokens. We do not truncate during inference and we truncate during training mainly for speeding up and saving GPU memory.

Our point is not purely for longer input, since the challenge of assembly code is longer input + low information density. The model has to summarize the tokens at a higher level to extract functional semantics.

4. Details of Functional CL

The F in the Functional CL (L 264) is a set of functions in the training data. For example, in practice, we have a batch size of 64 during the training with Functional CL, then F is a set of 64 unique functions, and each function has a source code, O0, O1, O2, and O3 assembly format.

Then we can calculate the loss $L_{fcl}$ (L 274) among this batch and perform backward propagation.

5. Apply of Hierarchical Attention

In Appendix Table 10, we show that without hierarchical attention or applying hierarchical attention on all the attention heads are both not good. That’s why we choose to apply on half of the attention heads.

Although a more fine-grained study could be conducted (e.g., applying on a quarter of heads and so on), given that each training of Nova requires nearly a thousand GPU hours, we are unable to thoroughly search the design choices.

Intuitively, this is a trade-off between the effectiveness of hierarchical attention and the pre-training knowledge that is kept in the original attention heads.

6. Availability

Yes, we plan to release the model weights, and code (for hierarchical attention implementation) after the notification.

We thank the reviewer for the insightful review and questions.

1. Fairness of Comparison

In our comparison with other models on decompilation, except for LLM4Decompile and Nova (these two are fine-tuned for this task), the other models are prompted with three examples as few-shot learning. Especially for GPT models, few-shot learning is the common practice of applying them to downstream tasks. We show that Nova-1B outperforms the most powerful models (as these are the most powerful commercial or open-source LLMs at the time) by a significant margin, highlighting the contribution of building such a foundation model for assembly code.

Besides, our main focus is the comparison with LLM4Decompile. LLM4Decompile is a closely related baseline that also fine-tuned the base backbone (DeepSeek-Coder) using assembly code collected from AnghaBench. Nova-1B significantly outperforms LLM4Decompile showing that it is indeed Nova’s design (hierarchical attention + contrastive learning) that brings the improvement.

2. Contribution of Each Component of Hierarchical Attention

We are only able to design the hierarchical attention empirically, as the build of Nova models is expensive, requiring thousands of GPU hours.

Yet, we compare Nova’s hierarchical attention design with LongCoder’s attention, showing that Nova’s attention design works better on assembly code.

3. Contribution of Each CL Objective

We conduct additional ablation studies to study the impact of each CL objective (FCL for functional contrastive learning, OCL for optimization contrastive learning). Below is the result:

Nova$_{-CL-HA}$ means no CL and hierarchical attention (HA), basically standard fine-tuning.

Nova$_{-FCL-HA}$ means adding OCL.

Nova$_{-OCL-HA}$ means adding FCL.

Nova$_{-HA}$ means adding both FCL and OCL.

From this result, we can see that both FCL and OCL bring improvement. Yet, it is clear that FCL indeed has a bigger impact than OCL. We have added this result to the Appendix in our updated version of the PDF.

4. Normalization

The design of the normalization process is based on the following insights and requirements:

• Existing code LLMs’ tokenizers do not fit the assembly code very well. For instance: ’(%rax),%rbx’ is tokenized to ‘(‘ ‘%’ ‘ra’ ‘x’ ‘),’ ‘%’ ‘rb’ ‘x’, where ’),’ is considered as one token. This does not properly show the semantics, since in X86 Assembly, parentheses are explicitly used to denote memory addressing while commas are used to separate operands. As we cannot re-train the tokenizer, we normalize the assembly code to make the tokenization better reflect assembly code semantics. We add white space to separate ‘(‘, ‘)’ and ‘,’, so that they are always considered as separate tokens.

• We remove all the hexadecimal values as LLMs are very poor at understanding hexadecimal values.

• Instead, we use the [INST-i] tokens to index the instructions. In decoder-only generative LLM, we want them to summarize the semantics of each instruction and thus have to put them at the end of each instruction since decoder-only LLM only has single-direction attention.

5. Train on O0, O1, O2 and Test on O3

We acknowledge this could be a more challenging setting and interesting future work. Yet, most existing baselines (LLM4Decompile for decompilation, jTrans, CodeArt for similarity detection) we compare with use the same setting as us, so we think this research question is out of the scope of this paper.

6. Randomness of Sampling

The randomness of sampling during generation is a common issue for other tasks such as code generation. We think the metric Pass@K [1] already handles the problem since Pass@K measures the expectation/chance of getting correct decompilation if we let the model generate K times (e.g., Pass@1 means the expectation of having one correct decompilation if the model is generated once, and Pass@10 means the expectation of having one correct decompilation if the model is generated 10 times).

In our evaluation, we let each model sample 20 generations and then calculate Pass@1 (i.e., how many of the 20 generations are correct?) and Pass@10 (i.e., if I pick 10 generations from the 20, what is the chance of having at least one correct decompilation?). We think reporting Pass@K with sampling generation is a common practice and the result is statistical.

[1] https://arxiv.org/abs/2107.03374

We thank the reviewer for the insightful review and questions.

1. Optimization Level During Inference

We conduct an evaluation of decompilation where the optimization level is not provided to the model. Fortunately, Nova is able to keep the highest performance without the optimization information.

Without the optimization information, the averaged Pass@1 of Nova-1B drops slightly from 25.17 to 24.92, yet it still significantly outperforms other baselines. We also find that without optimization information, Nova-1B actually performs slightly better on decompiling O3 assembly, which could be the variance of experiments.

We think one possible reason why Nova is able to handle decompilation without optimization information, is that during the pre-training (language modeling, contrastive learning) stage, the assembly code is input without any optimization information, Nova is trained to predict assembly instructions and encoding assembly code without knowing the optimization level. This training enables Nova to perform well on decompilation without knowing the optimization information.

2. GPT Baseline

• We want to clarify that during the inference of binary code decompilation, except for LLM4Decompile and Nova (these two are fine-tuned for the task), we provide the other baseline models with three examples for few-shot learning. Thus, the GPT models have three-shot examples.

• We conduct a one-round self-debug experiment using GPT-4o, letting GPT-4o revise the decompilation if it fails the test cases. The Pass@1 result is shown below:

Although GPT-4o is able to revise some incorrect decompilation if we provide the execution feedback to it, the Pass@1 after such revision is still much worse than Nova-1B which is designed and fine-tuned for assembly code.

3. Related Works

We thank the reviewer for pointing out additional related work, and we have added them to the "Binary Models" section in the related work part in our updated version of the PDF.

We thank the reviewer for the insightful review and questions.

1. Design Choice of Preceding-Instruction Attention

There are two reasons for using preceding-instruction attention:

• Intuitively, we separate the preceding-instruction attention from the inter-instruction attention, since inter-instruction attention is more about higher-level semantics. For the example in Figure 1 (a) and (b). The five instructions from 10 to 1c are corresponding to the if-condition in the source code. Such functional semantics are different from the more fine-grained dependencies in the compiler’s view (allocation of registers, maintaining of stack).
• Another reason is that Nova is designed for decoder-only generative LLM, where the attention is single-directional. The preceding-instruction attention enables every token in an instruction capture the dependencies with previous instructions. While the inter-instruction attention is only among the [INST-i] tokens (due to the design that we want this to capture functional semantics across multiple instructions), and thus the other tokens in the instruction do not have access to the inter-instruction attention.

2. Improvement of BCSD

We agree the improvement on the binary code similarity task is less significant than that on the binary code decompilation task. Yet, binary code similarity is a much better-explored domain than binary code decompilation. There has been more related work from software engineering, and security domain design models for it. The baselines we compare (jTrans, DiEmph, and CodeArt) have already achieved impressive performance. By contrast, end-to-end binary code decompilation is a very new task and LLM4Decompile is the only existing SOTA baseline we can find. Thus, we think there is less space to improve on binary code similarity than binary code decompilation, and that Nova outperforming existing SOTAs is valuable and hard.

Besides, those specialized BCSD models we compare with are encoder transformers. The bidirectional attention benefits their performance on the BCSD task as it is essentially an encoding task. Using decoder-only LLM for encoding tasks such as the BCSD task is less common and harder. Our ablation study (in Appendix Table 12-15) also shows that without Nova’s designs, simply training DeepSeek-Coder-1B cannot outperform existing SOTAs on BCSD.

Then why do we explore Nova (a decoder-only LLM) on BCSD? Given that decoder-only LLM is getting much more popular and promising, and the latest advanced LLMs are almost all decoder-only (Llama, StarCoder, etc), Nova shows the potential of building on top of the latest decoder-only LLMs for encoding tasks such as BCSD so that we can make good use of those pre-trained decoder-only LLMs.

3. Difference Between Nova’s Hierarchical Attention and CodeArt’s Attention Regularization

There is one fundamental difference between Nova's and CodeArt’s attentions. CodeArt is designed for Encoder LLMs (e.g., BERT) with bidirectional attention, while Nova is designed for Decoder-only generative LLMs with single-directional attention. This design choice is that decoder-only LLM shows more promise in generation tasks, and the most advanced LLMs are typically decoder-only.

Thus, Nova’s attention has to use the proceeding-instruction attention to let every token in an instruction capture dependencies with previous instructions.

In addition, Nova’s design is compatible with the standard self-attention used for source code (Figure 3 (c)), as we want to build Nova’s approach as a plug-in that can work with any transformer-based, decoder-only LLM to make use of their pre-training knowledge. While CodeArt’s design only works for assembly code and cannot handle input mixes assembly and source code or natural language.

4. How Contrastive Learning Improves Decompilation

Our insight is that contrastive learning enables the model to produce similar embeddings (i.e., the hidden states from the last transformer layer for [INST-i] tokens) for assembly code from the same source code. As LLMs understand the O0 assembly code the best (O0 assembly is decompiled with the highest pass@k), such aligning enables the LLMs to link the O1-O3 assembly with the O0 assembly and potentially learn the O1-O3 assembly better.

Besides, during the decompilation, the LLM generates the decompilation given the assembly code as input, i.e., the generation of decompiled source code is conditioned on the hidden states/embeddings of the input assembly code. Thus, we think better embeddings can lead to better decompilation.

5. Generalizability

Nova’s approach is generalizable to different assembly languages such as ARM or MIPS, since they share a similar challenge as X86. Nova’s approach can also be used on any transformer-based, decoder-only generative LLM.

[1] https://aclanthology.org/2023.acl-long.355/

  jmp .label2
.label:
  mov rax, rbx
  cmp rax, 0
  je .label
.label2:
  ret

Thanks for the follow up. We generally agree with the reviewers, but we want to clarify a few disagreements.

1 & 3: CodeArt’s Limitation on Unidirectional Attention, which Motivates Preceding-Instruction Attention

CodeArt’s attention is indeed limited for the unidirectional attention (decoder-only) model. The reviewer says: “token level inter-instruction" fine-grained attention is not necessarily based on CodeArt’s Figure 6”, which seems not entirely true given their use of the <cls> token.

Considering two instructions, CodeArt puts <inst> at the beginning of every instruction and puts a <cls> token at the beginning of the entire binary program:

<cls>
<inst> mov ( rdi ) , ecx
<inst> mov ( rsi ) , edx

In CodeArt’s design, the <inst> tokens have access to all the other <inst> tokens, so they are responsible for cross-instruction dependencies, our inter-instruction attention is similar to this.

But, the other tokens such as mov, rsi do not have access to other instructions, so the embeddings of these tokens inside each instruction may not consider the context. To address this, CodeArt gives every token inside the instruction access to the <cls>, and <cls> has access to all tokens in the binary program. Thus the rsi in the 2nd instruction can capture the dependencies with rdi in the 1st instruction, through rsi -> <cls> -> rdi, and vice versa. The CodeArt authors call it “global context” (Figure 6b in CodeArt paper), where each token does have dependencies with every other token through this <cls> token.

However, this design is not applicable to the decoder-only models, because in decoder-only models, there cannot be such a <cls> token as a “hub” to transfer attention among tokens in different instructions:

• if we put <cls> at the beginning, the <cls> token won’t have access to any tokens.
• if we put <cls> at the end, the other tokens won’t have access to it.
• if we add <cls> in the middle, the <cls> won’t have access to tokens after it, and the tokens before <cls> won’t have access to <cls>, still lack of circular attention transfer.

CodeArt’s design actually shows we do want tokens inside instructions to have some attention to the context, although the intra-instruction attention within each individual instruction should be dominant. This benefits the hidden states of the non-<inst> tokens.

Nova achieves this with preceding-instruction attention. See the same example binary but in Nova’s format:

mov ( rdi ) , ecx [inst-1]
mov ( rsi ) , edx [inst-2]

The rsi could have access to rdi through rsi -> [inst-1] -> rdi. rdi cannot access rsi anyway due to the limitation of unidirectional attention.

The reviewer proposes using window attention with a window size equal to two instructions, which indeed could be a promising design. But this is conceptually similar to our preceding-instruction attention together with the intra-instruction attention.

We also agree that other mixed designs could be a choice. We would like to discuss them in the related work sections in the final version. However, we cannot explore these solutions in the scope of this work.

2: A Small Misunderstanding of Table 11

In Table 11, Nova−OCL−HA (the better one) is without OCL so it is actually using FCL. We did “components removal” in the ablation study design. Our conclusion is that FCL has a bigger impact than OCL.

FCL aligns the functionality. Given the same functionality, no matter with O0-O3, Os, Ofast optimization, the assembly should all be grouped together. This is a common method for BSCD, and this is generalizable to different optimizations.

OCL considers the optimization monotonicity, which has less impact and is less generalizable. But Table 11 still shows that adding OCL leads to better results, especially with O0-O3 binaries. So one should not ignore the contribution of OCL during training.

Except for this small misunderstanding, we agree with the reviewer's opinion on the BSCD results.

We appreciate the reviewer’s acknowledgment of Nova’s promising results on the decompilation task.

We thank the reviewer for the insightful review and questions.

1. Clarification

• The related work section has been moved to follow the introduction section.
• We have highlighted in the data collection section that Nova focuses on C source code and X86-64 Assembly code.
• We appreciate the careful review and suggestions. We revised the text accordingly and tried to solve the ambiguity (notations, citations, etc.).

2. L2 Distance versus Cosince SImilarity

The reason we use L2 distance during training is that LLMs (e.g., DeepSeek-Coder on which we build Nova) hidden states are not normalized. Thus, the embedding we obtained for each source code and assembly function is not normalized. So we use L2 distance instead of adding normalization for simplicity.

During the evaluation of the BCSD task. We reuse the framework provided in the baseline work CodeArt, which uses cosine similarity to rank the assembly functions. We normalize Nova’s embedding of assembly functions during the evaluation, and we know the L2 distance after normalization keeps the same order as cosine similarity (smaller L2 distance means higher cosine similarity). So using normalization and cosine similarity during evaluation will not affect Nova’s performance.

3. Pass@K

Pass@K is the most popular metric used to evaluate the correctness of code generation, defined as Equation 1 in the Codex’s paper [1].

To be brief, Pass@K measures how often the model produces functional correct code. Higher Pass@K means we are more likely to see correct decompilation if we let the model sample K decompilations.

4. Table 4 Clarification

In Table 4, the “DeepSeekCoder + Nova’s Attention” is essentially Nova-1B, and “DeepSeekCoder + LongCoder’s Attention” is Nova-1B but replacing our hierarchical attention with LongCoder’s Attention. The rest are the same: both models are trained with the same data, and the same procedure (including CL). We clarify this in the updated version.

5. Using objdump -S or gcc -S

gcc -S is an alternative way of collecting the Assembly data, and Nova can be generalized to the Assembly produced by using gcc -S.

However, a key difference is that the assembly generated by gcc -S does not undergo the linking step. In practical scenarios, binary decompilation and analysis are typically performed on executable or linked assembly code, as it includes linker modifications and reflects the final binary structure.

Besides, our data collection aligns with many existing baselines, such as LLM4Decompile and jTrans, that collect their assembly code from executable binaries.

6. Normalization

• Existing code LLMs’ tokenizers do not fit the assembly code very well. For instance: ’(%rax),%rbx’ is tokenized to ‘(‘ ‘%’ ‘ra’ ‘x’ ‘),’ ‘%’ ‘rb’ ‘x’, where ’),’ is considered as one token. This does not properly show the semantics, since in X86 Assembly, parentheses are explicitly used to denote memory addressing while commas are used to separate operands. As we cannot re-train the tokenizer, we normalize the assembly code to make the tokenization better reflect assembly code semantics. We add white space to separate ‘(‘, ‘)’ and ‘,’, so that they are always considered as separate tokens.
• We remove all the hexadecimal values as LLMs are very poor at understanding hexadecimal values.
• Instead, we use the [INST-i] tokens to index the instructions. In decoder-only generative LLM, we want them to summarize the semantics of each instruction and thus have to put them at the end of each instruction since decoder-only LLM only has single-direction attention.

7. Clarification on Typo in $L_{BCSD}$

Yes, thank you for pointing out the typo. The numerator $f^q_j$ should be $f^p_j$, $f^p_j$ is the positive candidate in the pool that comes from the same source code function as the query $f^q$. We have fixed it in the updated version.

8. Clarification on GPT Inference

The GPT models also sample 20 decompilation per Assembly code using the same hyper-parameter (temperature 0.2, top-p 0.95), we rephrase the text to make it clearer.

Reference: [1] https://arxiv.org/abs/2107.03374

We upload a new version of the PDF.

1. "Train constraints" is indeed confusing. We revise the text between Line 260--268. We means "train Nova so that its embeddings for functions satisfy the constraints". Line 275--277 explains that in practice, the loss is calculated among each batch of functions (F is a batch of functions, not the entire dataset)
2. This is added to the Appendix A.6
3. Yes, we run test cases to validate correctness and calculate Pass@K. This is added in line 355--356.
4. I think our initial naming of "DeepSeekCoder + Nova's attention" confuses you. The reason you find that the numbers in the second row in Table 4 is the same as Nova-1B in Table 2 is because they are indeed the same model. We are comparing "DeepSeekCoder + LongCoder Attention + CL" with "DeepSeekCoder + Nova Attention + CL" in Table 4. The only difference between the two rows is replacing Nova's attention from Nova-1B with LongCoder's attention.
5. This is added to the Appendix in Line 892--899.
6. I agree with your point. This design is more by intuition before running experiments, and we do not have enough resources to test on this design choices. Yes, LLMs may be capable to generalize directly. By removing, I mean removing the address/offset of each instruction, e.g., for instruction "4: push %rbp", "4" is the address/offset. We remove this and use [INST-i] to indexing each instructions. For other hexadecimal values used in instructions, we convert them to decimal values.

Let me know if you have additional concerns, we still have time before the end of Nov. 27th to update the PDF. Thank you.