Actionable Cyber Threat Intelligence using Knowledge Graphs and Large Language Models

Prompt engineering significantly affects the output of LLMs, with elements such as prompt structure, spacing, line breaks, and punctuation playing a crucial role. The impact of the prompt design was evident when comparing outputs from different sizes of the Llama 2 model; a prompt that was highly effective with the 7B model yielded less satisfactory results with the 13B variant. For example, considering the one-shot prompt [INST] <<SYS>>\nExtractct [subject, predicate, object]-triples from the input text with the following predicates: isA, targets, uses, hasAuthor, hasAlias, indicates, discoveredIn, exploits, variantOf, and has.\n<<SYS>>\n\nInputut text: {ex_txt} [/INST] Extracted triples: {ex_triples} </s><s>[INST] Input text: {input_txt} [/INST] Extracted triples: . This prompt yielded the desired output as illustrated in Table V, where prompt engineering combined with greedy decoding produced the expected results for the 7B model. However, applying the same configuration to the 13B model resulted in merely listing entities from the text, e.g., The Adwind RAT, petroleum industry, US, new campaign, multi-layer obfuscation, evading detection,.... Consequently, this makes it a time-consuming process of trial and error. While adding a single example often enhanced the output, incorporating multiple examples, as the one mentioned in section 3.4.1, did not necessarily lead to better outcomes and, in some instances, even deteriorated the model’s performance. For instance, the Llama 2 7B chat model produced continuous text with over two examples, whereas other models struggled with inaccurately inventing relationships like hostedOn or usedFor from text verbs. Llama 2-Chat models presented notable challenges in generating outputs in a format suitable for further processing, in contrast to Mistral and Zephyr models. Llama models exhibited less consistency in output formatting, whereas Mistral, in particular, was able to produce the desired format reliably, even in zero-shot inference scenarios. The prompt [INST] <<SYS>>\nYouou are a helpful cyber-security assistant. Your task is to extract [entity,relationships,entity]-triples from the input text. Use only the relationships: ‘isA’, ‘targets’, ‘uses’, ‘hasAuthor’, ‘hasAlias’, ‘indicates’, ‘discoveredIn’, ‘exploits’, ‘variantOf’, and ‘has’. Your answers need to be in the format [entity,relationship,entity]. Reply only with the extracted triples.\n<</SYS>>\n\nWhatat are the triples in the following input text: {input_txt} [/INST] with Mistral greedy decoding already yielded triples in the format [Adwind RAT,isA,Remote Access Trojan (RAT)], [Adwind RAT,targets,petroleum industry], [Adwind RAT,targets,US].

To ensure the model maintained the desired relationship types in the triples, merely providing examples was insufficient; incorporating the ontology directly into the prompt proved to be crucial. The most effective results across all models were achieved by combining the ontology with one or two examples in the prompt, as illustrated in the prompt described in the first paragraph of this section.

Our comparative analysis of model performance with prompt engineering is summarized in Table III, which presents the ROUGE metric scores in the context of n-gram overlap for n=1,2,3,6 (ROUGE-N), as well as the longest common subsequence (ROUGE-L) for different prompts and models applied to the extraction of triples from the test data. It lists the results of the best-performing prompts for the models following the triple extracting from the LLM’s output. The ROUGE scores support the findings that Mistral and Zephyr outperform Llama 2 models when only prompt engineering is applied. The frequent failure of the Llama models to generate the desired triples led to a score of 0 for those specific test instances, which significantly reduced the overall ROUGE score. Noteworthy is the performance of the Llama 2 70B chat model, which, when prompted using a one-shot technique, achieved results that were the most comparable. However, its lower score can be attributed to its inability to fully comply with the specified ontology, particularly concerning the relationship types outlined in the prompt. This led to the generation of triples that diverged from the manually annotated reference, a discrepancy that became more pronounced in longer texts.

Guidance demonstrated remarkable efficiency, delivering impressive outcomes and proving easy to implement. It surpassed prompt engineering in achieving higher ROUGE scores, as illustrated in Table IV, which displays the scores for various models using guidance. For the smaller Llama 2-Chat models, optimal performance was achieved using the prompt Extract triples from the following input text. Your answers need to be in the format [entity1, relation, entity2]. with three examples, whereas the 70B chat model showed its best performance with a more extensive prompt that included the ontology, i.e., Extract [entity, relationships, entity]-triples from the following input text. Only use the relationships: ‘isA’, ‘targets’, ‘uses’, ‘hasAuthor’, ‘hasAlias’, ‘indicates’, ‘discoveredIn’, ‘exploits’, ‘variantOf’, and ‘has’. and just one example. Although the 13B chat model attained the highest ROUGE scores, manual human evaluation of the test data’s generated triples exposed its difficulty with complex syntactical structures such as negations, resulting in the majority of triples being incorrect. For instance, from the text First Twitter‑controlled Android botnet discovered Detected by ESET as Android/Twitoor, this malware is unique because of its resilience mechanism. Instead of being controlled by a traditional command-and-control server, it receives instructions via tweets. the model incorrectly extracts triples such as [Twitoor, isA, command-and-control server] and [Twitoor, isA, traditional command-and-control server], failing to correctly interpret the word instead.

The same issue was observed with the 7B model, positioning the 70B model as the top-performing guidance model according to human analysis.

Fine-tuning with fewer than the 718 examples in the dataset proved insufficient. The most effective approach utilized a constant learning scheduler with a learning rate of 2e-4, and LoRA parameters set with an alpha and rank of 16 for all linear layers. Significant variations in output were observed when different symbols were used to separate the input text from the prompt. Introducing line breaks into the input text for fine-tuning the model led to poor results, characterized by outputs that were either mere numbers or repetitive sections of the prompt. This underscores the critical importance of strict adherence to formatting, particularly regarding line breaks and prompt structure, to attain the intended output. Moreover, longer instructions tended to confuse the language model. For instance, the prompt [INST] Extract [subject, predicate, object]-triples from the following input text. Input text: ‘‘{input_txt}’’ [/INST] significantly outperformed more detailed prompts that included the ontology, as discussed in section 3.4.2. The latter failed to extract any triples, merely generating repetitive text regardless of the decoding strategy employed.

As illustrated in Table V, which presents a comparison of the different outputs generated by applying various decoding strategies in the context of Llama 2 7B chat prompt engineering versus fine-tuned model, the impact on the fine-tuned model is more pronounced. While greedy decoding restricted the output to only one word and multinomial sampling to only two triples with the fine-tuned model, beam-search multinomial sampling resulted in greater variance but highly repetitive output towards the end, even with a repetition penalty of 1.1. Generally, beam-search outperformed other strategies for fine-tuning, while prompt engineering achieved the desired output with greedy decoding and multinomial sampling, as evident in the table with the most relevant triples generated for those strategies. To avoid repetitive output, setting a repetition penalty of 1.1 for beam-search was necessary.

Fine-tuning the base models resulted in repetitive output that did not align with the instructions or the desired outcome, suggesting that the dataset size was insufficient for instruction-based tuning of the base Llama 2 models. Notably, loading the trained LoRA weights from the base models into the chat version for inference yielded good results, comparable to fine-tuning the chat versions directly. This suggests that chat-optimized models, being more adept at handling instructional tasks, can effectively capture essential patterns or knowledge relevant to this application through the LoRA weights, even if they were trained using a different prompt format.

The ROUGE scores for the Llama 2 7B and 13B chat versions, differentiated by the use of either <txt> and </txt> markers or quotes to separate the input text from the prompt, alongside the fine-tuned base models employing the alpaca prompt format, are presented in Table VI. The fine-tuned chat models demonstrate notable consistency across the metrics, achieving higher scores than any other techniques could attain, highlighting the effectiveness of fine-tuning in improving the quality of model outputs. Manual analysis of the generated triples from the test data supports the ROUGE scores. The fine-tuned Llama 2 7B chat model, utilizing input text marked by <txt> and </txt> , along with the fine-tuned Llama 2 7B and 13B models loaded onto their respective chat versions, exhibited the best performance. However, the latter models each failed to produce a response for different test samples.

Given that the fine-tuned 7B chat model demonstrated superior performance in both the ROUGE metrics and manual analysis, it was selected for the subsequent phase: generating a KG from CTI data.

Despite the model’s optimal performance on the test data, as evidenced by both ROUGE scores and manual human evaluations, its application to a large-scale dataset for KG generation produced very noisy triples. In the initial approach, which involved generating triples from the text without naming entity types, the fine-tuned 7B model extracted 188,547 triples from nearly 12,000 documents. Upon first manual inspection, it became evident that the model generated numerous nonsensical triples. It not only extracted entities that might be correct but did not correspond to a given entity type, such as [+44113320****, indicates, Phone Numbers], but also identified non-named entities, like [124, indicates, Android], and produced factually incorrect triples, for instance [/facebook, targets, Facebook]. A qualitative assessment of 100 randomly sampled triples revealed that only 39 of them were correct. Post-processing, for example, enforcing dates as the only valid subject for the discoveredIn relationship, or excluding numbers as subjects, was minimally effective and did not yield a usable KG. This scenario illustrates that while LLMs may perform well on specific, curated test datasets, they struggle with large-scale, primarily unprocessed data. The limitations observed could be attributed partly to the relatively small size of the LLM, having only 7 billion parameters, and possibly to insufficient training data.

To facilitate data analysis of the generated triples and perform post-processing steps in relation to the ontology, additional experiments were conducted not only to extract the triples but also to identify the entity types within them. For instance, triples were generated in formats such as [Adwind[Malware], targets, US[Location]] and [Adwind[Malware], isA, Trojan[MalwareType]]

Table VII summarizes the ROUGE scores for various techniques and models that showed promise in earlier stages by producing triples with entity types for the test data. The 70B chat model failed to generate the triples in a format suitable for automatic processing in 19 of the 29 test examples, resulting in a low overall ROUGE score. Mistral and Zephyr achieved significantly better scores; however, manual analysis revealed that, although they extracted numerous triples, they frequently introduced incorrect relationships, such as hasCapability or hasFeature, and incorrect entity types, such as Command or Attack. Moreover, there were numerous errors in the entity types for the extracted triples, such as categorizing APK as Indicator and banking apps as AttackPattern.

The guidance models attained even higher ROUGE scores, but human evaluation of the generated indicated that they struggled with accurately extracting and classifying indicators, often misclassifying them as AttackPattern. Furthermore, when the extracted entity was correctly identified as Indicator, the relationship often did not adhere to the ontology, for example [FakeSpy[Malware], isA, ANDROIDOS_LOADGFISH.HRX[Indicator]]. The gudiance models also struggled with adhering strictly to predefined relationships.

The fine-tuned model was most accurate in adhering to the ontology, especially with AttackPattern and Indicator extraction and triple generation. However, it sometimes struggled to extract triples at all, merely repeating the input text, which resulted in lower ROUGE scores. Combining the fine-tuned model with the guidance framework during inference proved not to offer significant advantages.

Based on these results, both the fine-tuned Llama 2 7B chat model and the 70B guidance model equipped with entity type extraction were chosen for generating the knowledge graph from the 12,000 documents. Consistent with the performance observed on the test data, the fine-tuned extracted fewer triples, yielding a total of 77,307. The qualitative assessment revealed that 55 out of 100 randomly sampled triples were correct, indicating significant improvement compared to fine-tuning and extraction processes that did not utilize entity types. Post-processing efforts such as enforcing the specified ontology, excluding Time entities lacking a date, omitting triples where the Malware entity type was mentioned fewer than five times across all documents, and removing non-named entities from the Malware and ThreatActor categories, effectively reduced the number to 41,525 triples. Although the data still contained some noise, the qualitative assessment showed improvement with 77 out of 100 triples being correct. These refined triples were then used for link prediction. The guidance model extracted nearly 460,000 triples from the documents, but the qualitative assessment revealed that the majority were incorrect, with only 15 out of 100 random samples being correct. After post-processing, only 17,050 entities remained, demonstrating higher quality, as evidenced by 66 correct samples in the qualitative assessment. This significant reduction in triples was due to the model incorrectly connecting entity types with certain relationships and extracting triples with newly invented entity types or relationships. Given that the fine-tuned model yielded better-quality triples, the triples from the guidance model were not utilized for further link prediction tasks.

For evaluating the efficacy of the prediction task, especially within knowledge graph forecasting, a comparison with triples generated using LADDER, as presented by Alam et al. [3], based on the same corpus of 12,000 documents was conducted. Traditional metrics such as Mean Reciprocal Rank (MRR) and Hits@n were employed. These metrics derive from the rankings that the link prediction model assigns to all accurate test triples. MRR represents the mean of the inverse rankings of all correct test triples, while Hits@n indicates the proportion of instances where a true triple appears within the top n positions of the rankings. Higher values in these metrics indicate superior performance.

The evaluation utilized the TuckER [46] model within a transductive setting across two distinct test sets. TestSet1 comprises hand-annotated triples that are not included in the training triples and have not been mapped to MITRE ATT&CK techniques. Given TuckER’s reliance on pre-identified entities, an additional test set containing only entities that were extracted by both LADDER and the fine-tuned 7B chat model was constructed as TestSet2. TuckER was implemented using the pykeen library [63], adopting Alam et al.’s [3] parameters. The setup included 50 embedding dimensions, a batch size of 64, and an initial learning rate of 0.001, while the training was reduced to 500 epochs.

The results, as presented in Table VIII, demonstrate that for TestSet1, the TuckER model, when trained on triples extracted through the fine-tuned 7B chat model, achieved significantly higher scores, with a notable increase of 9.26% in Hits@30 compared to its counterpart. For TestSet2, the performance gap narrows, with the fine-tuned model exhibiting only marginal improvements in Hits@30 and MRR metrics. This suggests that the triples extracted using our methodology may be slightly more refined, thereby enhancing the model’s ability to predict links between known entities more effectively. However, when comparing our results to those reported by Alam et al. [3], it is evident that our overall scores are significantly lower. This discrepancy can be attributed to the absence of a mapping for MITRE ATT&CK techniques, which leads to a reduction in the number of unique entities, consequently simplifying both the training and prediction processes.

Further inductive link prediction with the NodePiece [48] implementation in pykeen library [63] was conducted on an inference dataset containing entities not included in the test data. The results of the two models, one trained with the extracted triples from LADDER and one with triples from the fine-tuned 7B chat model, are very close. This implies that the relationships within the triples are fairly clear, enabling both models to effectively generalize from observed to unseen data. For example, the trained NodePiece is capable of inferring a new relationship between the malware Android/AdDisplay.Ashas, which was not included in the training data, and the Android operating system. It accurately predicts this relationship when queried with the malware as the head entity and targets as the relationship. The second most likely prediction for this relationship is Russian, reflecting the significant number of users in Russia who were affected by this malware [64].