# TASER: Task-Aware Stein Regularisation for Geometry-Driven Robustness

arXiv: https://arxiv.org/abs/2605.30601 · HTML: https://mlsanity.com/paper/2605.30601/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 20 of its 28 references that are in the map.
- First posted on arXiv 28 May 2026.

## Builds on (most cited first, 20 of 20)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Explaining and Harnessing Adversarial Examples](https://mlsanity.com/paper/1412.6572/) (2014)
- [Towards Deep Learning Models Resistant to Adversarial Attacks](https://mlsanity.com/paper/1706.06083/) (2017)
- [Benchmarking Neural Network Robustness to Common Corruptions and Perturbations](https://mlsanity.com/paper/1903.12261/) (2019)
- [Spectral Normalization for Generative Adversarial Networks](https://mlsanity.com/paper/1802.05957/) (2018)
- [Theoretically Principled Trade-off between Robustness and Accuracy](https://mlsanity.com/paper/1901.08573/) (2019)
- [Reliable evaluation of adversarial robustness with an ensemble of diverse parameter-free attacks](https://mlsanity.com/paper/2003.01690/) (2020)
- [Robustness May Be at Odds with Accuracy](https://mlsanity.com/paper/1805.12152/) (2018)
- [Diffusion Models for Adversarial Purification](https://mlsanity.com/paper/2205.07460/) (2022)
- [Parseval Networks: Improving Robustness to Adversarial Examples](https://mlsanity.com/paper/1704.08847/) (2017)
- [A Kernelized Stein Discrepancy for Goodness-of-fit Tests and Model Evaluation](https://mlsanity.com/paper/1602.03253/) (2016)
- [Improving Robustness using Generated Data](https://mlsanity.com/paper/2110.09468/) (2021)
- [Adversarial Risk and the Dangers of Evaluating Against Weak Attacks](https://mlsanity.com/paper/1802.05666/) (2018)
- [Measuring Sample Quality with Stein's Method](https://mlsanity.com/paper/1506.03039/) (2015)
- [Improving DNN Robustness to Adversarial Attacks using Jacobian Regularization](https://mlsanity.com/paper/1803.08680/) (2018)
- [Analysis of classifiers' robustness to adversarial perturbations](https://mlsanity.com/paper/1502.02590/) (2015)
- [Probabilistic Inference and Learning with Stein's Method](https://mlsanity.com/paper/2603.07467/) (2026)
- [TASTE: Task-Aware Out-of-Distribution Detection via Stein Operators](https://mlsanity.com/paper/2602.07640/) (2026)

## Abstract (arXiv)

Modern deep networks remain fragile under distribution shift and adversarial perturbations, often due to excessive or poorly structured input sensitivity. We introduce TASER (Task-Aware Stein Regularisation), a training-time regularisation framework derived from Langevin Stein operators. By penalising pointwise Stein residuals under the training distribution, TASER encourages geometric compatibility between predictors and data density, inducing anisotropic, data-aware smoothness. We provide theoretical links between Stein regularisation and reduced first-order shift sensitivity, develop scalable implementation variants compatible with modern architectures, and demonstrate improved robustness and stability across regression and vision benchmarks. Across CIFAR-10 experiments, TASER consistently improves the adversarial robustness of established training methods without incurring statistically significant clean-accuracy degradation.