UNIT: Backdoor Mitigation via Automated Neural Distribution Tightening
Siyuan Cheng, Guangyu Shen, Kaiyuan Zhang, Guanhong Tao, Shengwei An, Hanxi Guo, Shiqing Ma, Xiangyu Zhang
TL;DR
Backdoor attacks enable targeted misclassification by injected triggers. UNIT provides a post-training defense that learns a unique, tight activation distribution for each neuron from a small clean set and clips activations exceeding the learned boundary via an optimization-guided tightening of the per-neuron thresholds $σ^l_k$, producing clipped activations $\hat{F}^l_k(x)$. It outperforms 7 baselines against 14 backdoor attacks (including 2 advanced) using only $5\%$ clean data, with modest runtime overhead and broad generalization to multiple datasets, architectures, and even transformer models. The approach offers robust defense against adaptive attacks while preserving benign accuracy, making it a practical, scalable post-training backdoor mitigation technique.
Abstract
Deep neural networks (DNNs) have demonstrated effectiveness in various fields. However, DNNs are vulnerable to backdoor attacks, which inject a unique pattern, called trigger, into the input to cause misclassification to an attack-chosen target label. While existing works have proposed various methods to mitigate backdoor effects in poisoned models, they tend to be less effective against recent advanced attacks. In this paper, we introduce a novel post-training defense technique UNIT that can effectively eliminate backdoor effects for a variety of attacks. In specific, UNIT approximates a unique and tight activation distribution for each neuron in the model. It then proactively dispels substantially large activation values that exceed the approximated boundaries. Our experimental results demonstrate that UNIT outperforms 7 popular defense methods against 14 existing backdoor attacks, including 2 advanced attacks, using only 5\% of clean training data. UNIT is also cost efficient. The code is accessible at https://github.com/Megum1/UNIT.