Token pruning methods affect vision language model safety differently
Understanding and Mitigating Token-Pruning-Induced Vulnerabilities in VLMs
Cryptography and SecurityMachine Learning
Summary
Vision-Language Models use many visual tokens to understand images, but removing some tokens speeds up the process. The authors show that most ways of removing tokens actually make these models less safe, but one specific method called Query-based Compression can improve safety even when almost all tokens are removed. They discovered that pruning can accidentally make the model focus more on harmful parts of an image, making safety worse. To fix this, they created a new method that spots harmful tokens and adjusts the model’s attention to keep it safer without slowing it down.
What this means in practice
- •For ai safety engineers: Improve safety of vision-language models during real-time processing by reducing vulnerabilities caused by token pruning.
- •For machine learning infrastructure teams: Maintain acceleration benefits of token pruning in vision-language models without compromising their safe behavior in deployed systems.
Authors
Shuailong Wang, Xinyu Lyu, Shengming Yuan, Jingkuan Song, Heng Tao Shen, Lianli Gao
Abstract
Token-Pruning accelerates Vision-Language Models by removing redundant visual tokens, yet its safety implications remain underexplored. In this work, we present the first comprehensive safety evaluation of Token-Pruning mechanisms and find that: most pruning strategies significantly degrade safety as pruning ratios increase, whereas Query-based Compression shows the opposite, with extreme pruning (up to 99.8%), unexpectedly improves model safety. This sharp contrast prompts a key question: How do different Token-Pruning strategies reshape model safety behavior, and is it possible to enhance safety without sacrificing acceleration? To answer this, we identify an unrecognized mechanism, termed Pruning-Induced Malicious Amplification, where removal of background tokens triggers a side effect: forcing the model's attention to collapse onto a few retained malicious anchors within the foreground, inadvertently amplifying their toxic semantics under jailbreak. To address that, we propose an inference-time and plug-and-play Safety-Aware Pruning (SAP) mechanism that counteracts such dominance via three steps: (1) identifying malicious anchors, (2) restoring pruned benign tokens, and (3) reallocating excessive attention from malicious anchors to benign tokens. Extensive experiments across three safety and four utility benchmarks demonstrate that SAP mitigates pruning-induced vulnerabilities, i.e., reducing ASR by up to 62%, without compromising efficiency or utility.