Workflow:Norrrrrrr lyn WAInjectBench LLaVA Finetuning
| Knowledge Sources | |
|---|---|
| Domains | Prompt_Injection, Security, Fine_Tuning, Computer_Vision, VLM |
| Last Updated | 2026-02-14 16:00 GMT |
Overview
End-to-end process for fine-tuning a LLaVA vision-language model with LoRA adapters to perform binary prompt injection detection on images.
Description
This workflow fine-tunes the LLaVA-1.5-7B model to classify images as containing prompt injections (malicious) or not (benign). Rather than training the full 7B parameter model, it uses Low-Rank Adaptation (LoRA) to inject small trainable adapter matrices into the model's attention and feedforward layers, reducing memory requirements to fit on a single consumer GPU. The model is wrapped in a custom LlavaYesnoToken class that constrains output to binary Yes/No token logits, framing injection detection as a binary classification task. Training uses cross-entropy loss on the Yes/No logits with mixed-precision training, cosine learning rate scheduling with warmup, and automatic fallback to FP32 if numerical instability is detected. The best checkpoint is selected by validation TPR.
Usage
Execute this workflow when you need to train or adapt a LLaVA model for image-based prompt injection detection. You have labeled training and validation JSONL files containing image paths and binary labels, and you want to produce a LoRA checkpoint that can be loaded by the llava-1.5-7b-ft detector in the image detection pipeline.
Execution Steps
Step 1: Training Data Preparation
Prepare labeled training and validation datasets in JSONL format. Each record must contain a path field pointing to an image file and a label field (1 for malicious, 0 for benign). The training and validation files are separate JSONL files passed via CLI arguments.
Key considerations:
- Both train and validation JSONL files must follow the format: {"path": "...", "label": 1}
- Images are loaded as RGB via PIL at runtime
- No offline preprocessing is needed; images are processed by the LLaVA processor on-the-fly
Step 2: Model Initialization
Load the base LLaVA-1.5-7B model wrapped in the LlavaYesnoToken abstraction. This wrapper loads the model from HuggingFace (llava-hf/llava-1.5-7b-hf) with the specified precision dtype, initializes the processor and tokenizer, and identifies the single-token IDs for "Yes" and "No" in the tokenizer vocabulary. Gradient checkpointing is enabled to reduce memory usage.
What happens:
- The full LLaVA model (vision tower + language model) is loaded
- The tokenizer is probed to find single-token verbalizers for Yes/No
- The forward pass returns only the logits for [No, Yes] at the last token position
Step 3: LoRA Adapter Injection
Inject LoRA adapter matrices into the frozen base model using the PEFT library. Target modules include all attention projections (q_proj, k_proj, v_proj, o_proj), feedforward layers (gate_proj, up_proj, down_proj), and vision encoder projections (fc1, fc2). Only the LoRA parameters are set as trainable; all base model weights remain frozen.
Key considerations:
- Default LoRA rank (r) is 8 with alpha 32
- Typically less than 1% of total parameters become trainable
- LoRA dropout of 0.05 is applied for regularization
- The enable_input_require_grads() method is called to allow gradient flow through frozen layers
Step 4: Device Placement
Place the model on GPU with appropriate memory management. In single-GPU mode, the framework removes any accelerate dispatch hooks from the pretrained model, forces all parameters to the target device, and aligns LoRA child module devices with their parent weight tensors. If a single GPU runs out of memory, the framework can fall back to automatic multi-GPU dispatch via the accelerate library.
What happens:
- Accelerate hooks from model loading are removed
- All model parameters are moved to the target CUDA device
- LoRA adapter matrices are explicitly aligned to the same device as their parent weights
Step 5: Training Loop Execution
Train the LoRA-adapted model using AdamW optimizer with cosine learning rate scheduling and linear warmup. Each training step processes a batch of images through the LlavaYesnoToken forward pass to get [No, Yes] logits, computes cross-entropy loss against the binary labels, and updates only the LoRA parameters. Mixed-precision training (bf16 or fp16) is used by default with automatic gradient scaling for fp16.
Key considerations:
- Default learning rate is 2e-5 with 3% warmup ratio
- Gradient clipping at norm 1.0 prevents training instability
- If NaN or Inf loss is detected, the framework automatically falls back to FP32 with a learning rate backoff multiplier
- The GradScaler is only active for fp16 (not bf16)
Step 6: Validation and Checkpoint Selection
After each training epoch, the model is evaluated on the validation set. For each validation image, the model predicts Yes (malicious) or No (benign) based on which token logit is higher. True Positive Rate (TPR) and False Positive Rate (FPR) are computed from the confusion matrix. The checkpoint with the highest TPR across all epochs is saved as the best model.
What happens:
- Validation runs in no-grad mode with the same mixed-precision context
- The argmax of [No, Yes] logits determines the prediction
- The full confusion matrix (TP, TN, FP, FN) is computed
- Best checkpoint includes model state, optimizer state, epoch, TPR, and AMP configuration
Step 7: Checkpoint Export
The best checkpoint is saved as a PyTorch state dictionary (.pt file) containing the model weights (including LoRA adapters), optimizer state, epoch number, best TPR, and AMP configuration. This checkpoint can be loaded by the llava-1.5-7b-ft detector module for inference in the image detection pipeline.
Output artifacts:
- best_epoch{N}_tpr{X.XXXX}.pt — PyTorch checkpoint with best validation TPR