TPV: Parameter Perturbations Through the Lens of Test Prediction Variance

TL;DR

This paper introduces test prediction variance (TPV), a label-free measure of model sensitivity to parameter perturbations, unifying various robustness analyses and enabling practical applications like pruning and model selection.

Contribution

The paper proposes TPV as a novel, unifying framework for analyzing model robustness to parameter perturbations, connecting theory with practical model selection and pruning methods.

Findings

TPV converges from training to test set in overparameterized models.

TPV correlates strongly with test loss across various settings.

TPV enables effective label-free pruning and model selection.

Abstract

We introduce test prediction variance (TPV)--the first-order sensitivity of a trained model's outputs to parameter perturbations--as a unifying framework for analyzing post-training robustness. TPV is a fully label-free object whose trace form separates the geometry of the trained model from the specific perturbation mechanism, placing SGD noise, label noise, quantization, and pruning under a single lens. The resulting expressions recover the wide-minima hypothesis for SGD and quantization noise, and yield a distinct Jacobian-spectral characterization for label noise connecting label-noise TPV with benign overfitting in nonlinear networks. Theoretically, we prove that training-set TPV converges to its test-set counterpart in the overparameterized limit, irrespective of generalization performance, providing the first result that prediction variance under local parameter perturbations can be inferred from training inputs alone. Empirically, this stability holds far more broadly, including at very low widths. Further, TPV correlates well with test loss, enabling practical applications: JBR, a label-free pruning criterion derived from TPV geometry matching state-of-the-art baselines; and training-set based model selection signal for in-distribution and transfer learning scenarios. Code available at github.com/devansharpit/TPV.