Deep neural networks remain sensitive to distribution shifts, motivating regularisation informed by data geometry.
We introduce TASER (Task-Aware Stein Regularisation), a compact training-time framework that couples model derivatives to density scores by penalising the Langevin Stein operator.
We establish a weighted Sobolev interpretation and a hierarchy of stability bounds under progressively stronger distributional assumptions.
Moreover we introduce TASERLight, which is based on a first-order Stein operator.
TASER and TASERLight require neither explicit tangent-space estimation nor an inner perturbation search.
Controlled synthetic examples illustrate their effects on extrapolation and directional sensitivity.
Experiments on spectroscopy, skeleton motion, ECG, MNIST, and CIFAR-10 demonstrate gains under distribution shift and complementary benefits when combined with isotropic regularisation.