Method improves machine learning robustness by reducing overlap of probability distributions

Probabilistic Adversarial Training

Machine LearningArtificial Intelligence

Summary

Machine learning models can be tricked by adversarial examples—inputs designed to fool them. The authors think about these tricky inputs as coming from overlapping probability patterns. They propose a way to make the model more robust by pushing these patterns apart, reducing overlap and making it harder to fool the model. Their approach builds on a mathematical measure called KL divergence and connects it to existing adversarial training techniques. Experiments show it can improve robustness in practice.

What this means in practice

  • For machine learning engineers: Make classifiers tougher to trick by implementing probabilistic adversarial training to improve model resistance to adversarial inputs.
  • For security software developers: Enhance security systems by incorporating improved robustness methods based on probabilistic training to defend against adversarial attacks.

Authors

Andi Zhang, Xingyu Zhao, Siddartha Khastgir

Abstract

Building on a probabilistic perspective in which adversarial examples arise from the overlap between a distance-based distribution $p_{\mathrm{dis}}$ and a victim-classifier-induced distribution $p_{\mathrm{vic}}$, we start from a simple intuition: adversarial examples become harder to generate when these two distributions are pushed apart, as their overlap becomes smaller, thereby increasing robustness. This intuition naturally motivates a KL-based robustness objective. We then prove that $\mathrm{KL}(p_{\mathrm{dis}}\|p_{\mathrm{vic}})-\log Z_{\mathrm{vic}}$ is a lower bound on probabilistic robustness (PR), where $Z_{\mathrm{vic}}$ denotes the normalizing constant of $p_{\mathrm{vic}}$. Since PR is generally intractable to compute directly, maximizing this KL-based lower bound provides a tractable surrogate objective for improving PR. We further show that this objective recovers a scaled form of adversarial training, offering a probabilistic interpretation of adversarial training and a principled route to robustness improvement. We call the resulting method probabilistic adversarial training. Experiments show that it consistently improves PR, and ablation studies demonstrate that the induced scaling factor can even enhance the PR of non-probabilistic adversarial training methods.