How Can Training Gradients Be Compressed Without Losing Influence Estimates?

Researchers have proposed a method to shrink the stored gradients of large language models to a fraction of their original size while keeping the ability to trace how each training example shapes model behavior. The technique uses a combination of subspace projection, principal component analysis, and one-bit quantization to achieve compression ratios that were previously impractical.

Why Storing Full Gradients Fails At Scale

Influence functions measure how a single training example changes what a model does. Computing them repeatedly for different questions requires access to the gradients of every training example. A full gradient for a model with billions of parameters occupies gigabytes per example. Multiply that by millions of examples and the storage cost becomes impossible. The paper frames this as a compression problem: keep a fixed-size sketch of each gradient that still yields accurate influence estimates for queries that did not exist when the sketch was created.

The Core Idea

Compress the gradient. Keep the signal.

How EK-FAC Reduces Dimensionality

The proposed method is called eigenbasis-corrected one-bit gradient projection. It works in three stages. First, EK-FAC reduces the raw gradient to a lower-dimensional subspace that captures the most important curvature information. This step exploits the fact that gradient matrices in deep networks have structure. EK-FAC approximates the Fisher information matrix using a Kronecker product of smaller factors. The approximation lets the method project the high-dimensional gradient onto a much smaller subspace without losing the directions that matter most for influence estimation.

PCA Learns Compression Directions

Second, PCA learns the optimal compression directions inside that subspace from the distribution of training gradients themselves. The principal components are computed once from the stored gradients. They become the basis for all future compression. This means the compression adapts to the actual gradient distribution seen during training rather than relying on a generic assumption.

One-Bit Quantization Multiplies Capacity

Third, each coordinate is quantized to a single bit. One bit means positive or negative. This lets the system retain many more coordinates within the same storage budget than any multi-bit scheme could. The paper shows that the worst-case analysis characterizes the optimal fixed-dimensional linear representation. EOGP approximates that optimum at scale by combining the three stages.

Results On GPT-2

On GPT-2 the method predicts the effect of removing training examples more accurately than the evaluated baselines while using one-sixteenth of their per-example storage. The baselines include random projection and top-k gradient selection. EOGP outperforms them across multiple influence queries. The paper does not report wall-clock compression or decompression times.

OLMo 2 Experiments Scale Further

On OLMo 2 supervised fine-tuning models ranging from 1B to 32B parameters the approach stays competitive with baselines that are allocated over 100 times as much storage per example. The storage advantage grows with model size. The experiments cover two model families and a range of sizes but do not include models larger than 32B parameters.

What The Paper Does Not Say

It also does not say how the method behaves when the test-time query distribution shifts far from the training gradient distribution used to learn the PCA directions. Whether the same storage ratios hold for models with hundreds of billions of parameters remains an open question. The worst-case analysis characterizes the optimal fixed-dimensional linear representation but does not prove that the three-stage pipeline reaches that optimum for every model architecture.