Experimenting with Fine-Tuning and PEFT (LoRA) for Threat Modelling for CII

With the growing importance of large language models (LLMs) in tasks such as cybersecurity threat modelling for Critical Information Infrastructure (CII), efficient model training techniques are essential. Recently, I conducted an experiment on AWS SageMaker to fine-tune the FLAN-T5 model—an instruction-tuned LLM—for enhanced cybersecurity threat modelling. The primary goal was to compare traditional fine-tuning approaches with Parameter Efficient Fine-Tuning (PEFT), specifically leveraging LoRA (Low-Rank Adaptation), and evaluate the trade-offs in performance versus computational resource usage.

Overview of the Experiment

Threat Model and Objective

The primary threat model focuses on identifying potential vulnerabilities in Critical Information Infrastructure (CII) by leveraging dialogue-based cybersecurity insights. The goal is to enable accurate, context-aware summarization of threat intelligence reports to support decision-making and incident response.

Initially, the experiment started with a Retrieval-Augmented Generation (RAG) approach. While RAG performed reasonably well for generating summaries by combining a retriever and generator, it lacked the domain-specific precision needed for CII threat modelling. This limitation prompted an evaluation of full fine-tuning and Parameter Efficient Fine-Tuning (PEFT) approaches to achieve:

1. Improved domain adaptation of the FLAN-T5 model.
2. Optimal balance between computational cost and performance.

Overview of the Experiment

Objective

The experiment aimed to fine-tune FLAN-T5 for cybersecurity threat modelling tasks using two approaches:

1. Full Fine-Tuning: Updating all model parameters to optimize for the threat modelling task.
2. Parameter Efficient Fine-Tuning (PEFT – LoRA): Updating a subset of parameters using Low-Rank Adaptation to achieve similar performance while using significantly fewer resources.

Tools and Frameworks

• AWS SageMaker: Managed service for scalable machine learning.
• Hugging Face Transformers Library: For model fine-tuning and evaluation.
• ROUGE Metrics: Evaluation of threat modelling output quality.

Original Model Performance

Before applying any fine-tuning techniques, the ROUGE metrics for the original FLAN-T5 model were as follows:

• ROUGE-1: 23.34
• ROUGE-2: 7.60
• ROUGE-L: 20.15
• ROUGE-Lsum: 20.15

These baseline results highlighted the need for fine-tuning to achieve better performance for the specific task of cybersecurity threat modelling.

Full Fine-Tuning of FLAN-T5

Training Data

For this experiment, the training data consisted of domain-specific threat intelligence reports and dialogue transcripts from cybersecurity incident analysis. These datasets were curated to include:

1. Threat Intelligence Reports: Detailed documents covering various cyber threat vectors, vulnerabilities, and mitigation strategies.
2. Dialogue Transcripts: Contextual discussions from incident response scenarios to help the model adapt to conversational and situational nuances.

The combination of structured reports and unstructured dialogue data was designed to improve the model’s ability to summarize diverse inputs effectively.

One significant challenge encountered during full fine-tuning is catastrophic forgetting, a phenomenon where the model loses previously learned knowledge as it adapts to a new task. For cybersecurity threat modelling, this can manifest as a degradation in the model’s ability to generalize or respond to non-specific inputs after fine-tuning. To mitigate this, careful attention was paid to preserving general capabilities while optimizing the model for domain-specific tasks.

In the first phase, I fine-tuned the FLAN-T5 model using traditional methods, updating all its parameters. While the model delivered high-quality outputs, the process required substantial computational power and memory, often necessitating multiple GPUs.

Results

The ROUGE metrics for the full fine-tuning approach were as follows:

• ROUGE-1: 42.16
• ROUGE-2: 18.03
• ROUGE-L: 33.84
• ROUGE-Lsum: 33.83

These results confirmed the potential of FLAN-T5 for cybersecurity threat modelling but highlighted the high resource costs of full fine-tuning.

When to Use Full Fine-Tuning

• High-Accuracy Applications: When the task demands the highest possible performance, such as critical decision-making systems.
• Availability of Resources: When sufficient computational resources (e.g., multiple GPUs) are available to support the training process.

Parameter Efficient Fine-Tuning (PEFT – LoRA)

PEFT techniques like LoRA inherently help mitigate catastrophic forgetting by freezing most of the model’s parameters and only updating a smaller subset. This approach ensures that the model retains its general-purpose capabilities while adapting to the specific requirements of cybersecurity threat modelling. By using LoRA, the trade-off between resource efficiency and preserving prior knowledge is more manageable.

In the second phase, I implemented PEFT using the Low-Rank Adaptation (LoRA) technique. Instead of fine-tuning all parameters, LoRA modifies only a small subset of parameters by adding low-rank matrices to freeze the rest of the model. This significantly reduces the computational requirements, making it feasible to train on a single GPU.

Results

The ROUGE metrics for the PEFT-LoRA approach were slightly lower but still competitive:

• ROUGE-1: 40.81
• ROUGE-2: 16.33
• ROUGE-L: 32.51
• ROUGE-Lsum: 32.49

While the performance dipped by a small percentage, the reduction in resource consumption made PEFT-LoRA a compelling alternative.

When to Use PEFT-LoRA

• Resource-Constrained Environments: When limited computational resources are available.
• Rapid Iteration: For quick experimentation and prototyping with minimal infrastructure.
• Production Deployment: When the slight performance trade-off is acceptable in exchange for cost and resource efficiency.

Trade-Offs Between Approaches

Approach	ROUGE-1	ROUGE-2	ROUGE-L	ROUGE-Lsum	Example Use Cases
Original Model	23.34	7.60	20.15	20.15	Baseline; useful for testing out-of-the-box capabilities
Full Fine-Tuning	42.16	18.03	33.84	33.83	High-stakes applications where accuracy is paramount
Parameter Efficient FT	40.81	16.33	32.51	32.49	Cost-effective production systems and rapid iterations

The table highlights the performance metrics and example scenarios for each approach. Full fine-tuning offers slightly better metrics at the cost of higher resource consumption, whereas PEFT-LoRA balances performance with efficiency.

Implementation in Jupyter Notebook

Below is a screenshot of my Jupyter Notebook on AWS SageMaker, showcasing the ROUGE metrics evaluation:

Key Takeaways

1. Efficiency Matters: PEFT-LoRA is an excellent approach for scenarios with limited computational resources.
2. Performance Trade-Offs: The marginal drop in performance with PEFT-LoRA is often acceptable for real-world applications.
3. AWS SageMaker Simplifies Experiments: SageMaker’s integration with Hugging Face models makes it easier to fine-tune and deploy LLMs at scale.

Next Steps

In future experiments, I plan to explore additional PEFT techniques and investigate their impact on performance and efficiency. Additionally, I aim to evaluate the deployment of these fine-tuned models for real-time threat modelling tasks.

By leveraging AWS SageMaker and Hugging Face’s tools, organizations can build efficient and scalable solutions for cybersecurity threat modelling and beyond.