Scaling influence functions in large language models with compressed gradients

Scaling Influence Functions in LLMs through Eigenbasis-Corrected One-Bit Gradient Projection

Machine LearningArtificial Intelligence

Summary

Understanding which training examples most affect a language model’s behavior can help improve or fix it, but checking this influence usually requires lots of storage and computing. The authors developed a way to store a compressed version of the training information that is much smaller but still accurate for predicting influence later on. They use mathematical techniques to shrink and simplify the stored data, allowing them to keep more useful details within a limited space. Their method worked better than other compression methods on models like GPT-2 and larger ones, using significantly less storage. This makes influence estimation more practical for very large language models.

What this means in practice

  • For machine learning engineers: Store compressed gradient data for large language models to efficiently compute influence of training examples on model behavior.
  • For model auditing teams: Analyze model behavior impacts from training data with less storage, enabling influence computations at scale for large models.

Authors

Jaeseung Heo, J Rosser, Dongwoo Kim

Abstract

Influence functions estimate how individual training examples affect the behavior of large language models (LLMs). Analyzing how training data influence different behaviors of an LLM involves repeated influence computation. Reusing stored training gradients reduces the computational cost, but storing full gradients is prohibitively expensive at LLM scale. We study how to compress these gradients while preserving influence estimates for future queries that are unknown at storage time. Through a worst-case analysis, we characterize the optimal fixed-dimensional linear representation and propose eigenbasis-corrected one-bit gradient projection (EOGP) to approximate it at scale. Specifically, EOGP uses EK-FAC to reduce gradient dimensionality, then applies PCA within the retained subspace to learn compression directions from the training gradients. We then apply one-bit quantization to the resulting coordinates, allowing more coordinates to be retained within a fixed storage budget. On GPT-2, EOGP predicts retraining outcomes more accurately than the evaluated compression baselines while using one-sixteenth of their per-example storage. On OLMo 2 SFT models from 1B to 32B parameters, EOGP remains competitive with the baselines allocated over 100 times as much storage per example.
View PDFOpen arXiv