PSA-VLM: Enhancing Vision-Language Model Safety through Progressive Concept-Bottleneck-Driven Alignment

TL;DR

This paper addresses safety vulnerabilities in vision-language models caused by the visual modality bypassing LLM safeguards. It proposes PSA-VLM, a progressive safety alignment method that embeds safety concepts as bottlenecks using a Concept Bottleneck Model (CBM) with a Safety Projector, Safety Tokens, and a Safety Head, trained in two stages to improve risk detection and response without sacrificing multimodal capability. Stage I trains concept classifiers for safety features with frozen LLM and vision encoders, while Stage II unfreezes the LLM to integrate concept-level safety into decision making, guided by losses $\mathcal{L}_{s}$, $\mathcal{L}_{l}$, and $\mathcal{L}_{LLM}$. Evaluation on RTVLM and additional risk datasets shows PSA-VLM achieves state-of-the-art safety performance, with LoRA-based fine-tuning further boosting results, while keeping competitive general multimodal benchmarks. The approach enhances explainability and controllability by tying outputs to high-level safety concepts, offering a practical path toward safer VLM deployments in real-world settings.

Abstract

Benefiting from the powerful capabilities of Large Language Models (LLMs), pre-trained visual encoder models connected to LLMs form Vision Language Models (VLMs). However, recent research shows that the visual modality in VLMs is highly vulnerable, allowing attackers to bypass safety alignment in LLMs through visually transmitted content, launching harmful attacks. To address this challenge, we propose a progressive concept-based alignment strategy, PSA-VLM, which incorporates safety modules as concept bottlenecks to enhance visual modality safety alignment. By aligning model predictions with specific safety concepts, we improve defenses against risky images, enhancing explainability and controllability while minimally impacting general performance. Our method is obtained through two-stage training. The low computational cost of the first stage brings very effective performance improvement, and the fine-tuning of the language model in the second stage further improves the safety performance. Our method achieves state-of-the-art results on popular VLM safety benchmark.

Figures (17)

• Figure 1: Selected examples of using unsafe images to generate. The content inside the red box is the generated unsafe answer by other VLMs, while the content inside the green box is the safe answer generated by our PSA-VLM.
• Figure 2: Example of 10 tasks under Politics, Illegal Risk, Insults and Bullying, Fairness, Privacy, and Misleading categories in the RTVLM benchmark and other risk datasets.
• Figure 3: The overview architecture of PSA-VLM, which is trained in two stages: (1) safety concept extraction by freezing the LLM and vision encoder while training safety modules, and (2) enhancing safety alignment by unfreezing the LLM to integrate concept-level safety features into the VLM’s decision-making process.
• Figure 4: (a) t-SNE visualizations depicting the separation of unsafe image features in two-dimensional space. Each subplot corresponds to a distinct combination of feature sets and labels, illustrating differences between original and safe features. After using the safe projector, the features of unsafe images are significantly divided into different clusters. (b) The classification performance of safety level and safety type, including accuracy and F1-score.
• Figure 5: The filtered data by PSA-VLM in the MME dataset, including the tasks of Code Reasoning, Text Translation, Celebrity, Numerical Calculation, Poster, and Artwork.