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

2026-05-28

AI Summary

Deep neural networks exhibit vulnerability under distribution shifts and adversarial perturbations due to excessive or poorly structured sensitivity to inputs. This work proposes TASER, a training-time regularization framework based on the Langevin Stein operator, which enforces geometric alignment between the prediction function and the data density by penalizing pointwise Stein residuals under the training distribution, thereby achieving data-aware anisotropic smoothing. We establish a theoretical connection between Stein regularization and reduced first-order sensitivity to distribution shifts and develop a scalable implementation compatible with modern architectures. Experiments demonstrate that TASER significantly enhances adversarial robustness on visual and regression benchmarks such as CIFAR-10 without causing a statistically significant drop in clean accuracy.

Abstract

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.

Problem

• distribution shift
• adversarial perturbations
• input sensitivity
• robustness
• deep networks

Innovation

• Stein regularization
• distributional robustness
• geometry-aware learning
• adversarial robustness
• anisotropic smoothness