Efficient Sketches for Training Data Attribution and Studying the Loss Landscape
This work addresses memory constraints in studying modern ML models, offering incremental improvements for researchers analyzing model behavior.
The authors tackled the problem of storing large quantities of gradients or Hessian vector products in machine learning by developing a scalable sketching framework, achieving efficient applications such as training data attribution and Hessian spectrum analysis for pre-trained language models.
The study of modern machine learning models often necessitates storing vast quantities of gradients or Hessian vector products (HVPs). Traditional sketching methods struggle to scale under these memory constraints. We present a novel framework for scalable gradient and HVP sketching, tailored for modern hardware. We provide theoretical guarantees and demonstrate the power of our methods in applications like training data attribution, Hessian spectrum analysis, and intrinsic dimension computation for pre-trained language models. Our work sheds new light on the behavior of pre-trained language models, challenging assumptions about their intrinsic dimensionality and Hessian properties.