Transfer Safety Awareness for Cross-Modal Safety Drift

Official implementation of Transfer Safety Awareness for Cross-Modal Safety Drift in Multimodal Large Language Models (EMNLP 2026 Findings).

This repository contains the Qwen3-VL safety-awareness direction pipeline. It extracts a refusal-related activation direction from harmful and harmless multimodal examples, selects the direction that improves refusal behavior while limiting benign-task drift, and optionally trains a direction-specific safety-awareness vector.

The code currently supports Qwen/Qwen3-VL-8B-Instruct through Hugging Face Transformers.

Repository Structure

.
├── artifacts/directions/       # Precomputed selected direction and metadata
├── data/csv/                   # CSV manifests used by the pipeline
├── data/images/                # Local image roots (not included)
├── directions/                 # Direction extraction and selection
├── models/                     # Qwen3-VL model and processor utilities
├── training/                   # Safety-awareness direction training
├── utils/                      # Activation hook utilities
├── config.py
├── run_pipeline.py             # Direction generation and selection entry point
└── requirements.txt

Installation

Python 3.10+ and a CUDA-enabled PyTorch installation are recommended.

git clone https://github.com/cucu220123/transfer-safety-awareness.git
cd transfer-safety-awareness
python -m venv .venv
source .venv/bin/activate
pip install -r requirements.txt

The first run downloads Qwen/Qwen3-VL-8B-Instruct from Hugging Face. Make sure the machine has sufficient GPU memory and that you have accepted any model terms required by the model provider.

Data Preparation

The repository includes the CSV manifests used by the experiments, but does not redistribute the source image datasets. Obtain the images from their original dataset providers and place them under these project-relative roots:

data/images/vlsafe_images
data/images/vlsbench_imgs
data/images/mmvet_images

The mapping is:

Image roots and CSV paths can be overridden without editing code:

export VLM_HARMFUL_IMAGE_ROOT=/path/to/vlsafe_images
export VLM_HARMLESS_IMAGE_ROOT=/path/to/vlsbench_imgs
export VLM_KL_IMAGE_ROOT=/path/to/mmvet_images
export VLM_HARMFUL_CSV=/path/to/harmful.csv
export VLM_HARMLESS_CSV=/path/to/harmless.csv
export VLM_KL_CSV=/path/to/benign_vqa.csv

Please follow the licenses and terms of each source dataset when downloading or using the images and manifests.

Generate and Select a Direction

To run candidate direction generation and validation-based selection:

CUDA_VISIBLE_DEVICES=0 python run_pipeline.py --direction_only

The command uses 128 training and 32 validation examples per split by default. It writes intermediate results to artifacts/direction_runs/<model-name>/ and saves the selected direction to:

artifacts/directions/qwen3vl_refusal_direction.pt
artifacts/directions/qwen3vl_refusal_direction_metadata.json

Use --model_path to point to a local model directory or another compatible Hugging Face identifier. The --skip_filter flag disables the refusal-score filtering stage.

The repository already contains a precomputed direction artifact for the default Qwen3-VL checkpoint, so this stage can be skipped when that artifact is appropriate for your setup.

Train the Safety-Awareness Direction

Training updates the selected direction with the benign-task and refusal objectives used in the paper code. A four-GPU example is:

CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --nproc_per_node=4 --master_port=29501 \
  -m training.train_safety_awareness_direction \
  --epochs 8 \
  --batch_size 1 \
  --grad_accum_steps 1 \
  --lr 5e-3

Training outputs are written to outputs/qwen3vl_safety_awareness_train/ by default. Useful options include --model_path, --direction_pt, --direction_meta, --artifact_dir, --epochs, --batch_size, and --lr; run python -m training.train_safety_awareness_direction --help for the complete list.

Reproducibility Notes

• Set CUDA_VISIBLE_DEVICES and use the same model revision, input manifests, and image files for comparable results.
• The data loader assigns rows to train/validation/test using the CSV split column when present; otherwise it uses question_id % 10 deterministically.
• Model checkpoints and downloaded image datasets are intentionally excluded from this repository.
• Generated intermediate runs and training outputs are ignored by Git; copy any results you want to publish to a separate release or artifact store.

Citation

The paper citation will be added when the final bibliographic information and paper link are available.

License

The code in this repository is released under the MIT License. Third-party datasets, images, and the Qwen3-VL model remain subject to their respective licenses and terms.