PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Jaeseung Heo, J Rosser, Dongwoo Kim

arXiv:2609.37842v1cs.LGcs.AI

TL;DR

Influence estimation at LLM scale requires reusable training gradients, but full-gradient storage is costly and future queries are unknown. The paper characterizes an optimal fixed-dimensional linear representation and approximates it with EOGP, combining EK-FAC projection, PCA correction, and one-bit storage. EOGP outperforms tested compression baselines on GPT-2 and remains competitive with much larger-storage baselines on OLMo 2 SFT models.

  • Problem

    Full training gradients are costly to store at LLM scale, while compression must preserve influence estimates for future queries that are unknown when representations are built.

    Source: §1 · p. 1

  • Method

    EOGP approximates the worst-case-optimal fixed-dimensional representation by applying PCA to half-whitened gradients within an EK-FAC subspace and storing the resulting coordinates with one-bit quantization.

    Source: §1 · p. 2

  • Results

    EOGP outperforms tested compression baselines on GPT-2 retraining evaluations and remains competitive on 1B–32B OLMo 2 SFT models against baselines allocated over 100 times as much storage per example.

    Source: §1 · p. 2

  • Takeaways & Limitations

    EOGP preserves useful attribution information with substantially less storage than existing compression methods, making reusable gradient stores more practical at LLM scale.

    Source: §6 · p. 10

    Candidate-subspace selection uses the Fisher approximation as a surrogate for the empirical Fisher, although their label distributions differ.

    Source: §3.2 · p. 5

Abstract

from arXiv · show

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.

1 INTRODUCTION

The introduction frames gradient storage as a bottleneck for repeated influence queries whose future targets are unknown. It presents EOGP as a way to preserve useful influence information compactly and reports stronger results than tested compression baselines across the evaluated settings.

1 INTRODUCTION

  • A single training-example gradient for an 8B-parameter model requires 16 GB in half precision, making full-gradient storage costly at LLM scale.The motivation is to support repeated attribution queries with compact, reusable gradient representations.Source: §1 · p. 1
  • Future attribution queries are unknown when representations are constructed, so compression must preserve information useful across different subsequent influence calculations.Stored training-example gradients can be reused across queries, avoiding repeated backward passes over the training data.Source: §1 · pp. 1–2
  • Under a normalized worst-case influence-error criterion, top-k PCA coordinates of half-whitened training gradients are optimal among k-dimensional linear representations.The objective measures discrepancies between compressed and uncompressed influence scores across possible future queries.Source: §1 · p. 2
  • EOGP first uses EK-FAC directions to reduce gradient dimensionality, then applies PCA within that subspace and one-bit quantization to retain more coordinates under a fixed budget.The stages address which directions to retain and how much precision to allocate to each coordinate.Source: §1 · p. 2
  • On GPT-2, EOGP outperforms tested compression baselines in linear datamodeling score and counterfactual retraining; on 1B–32B OLMo 2 SFT models, it more accurately reproduces reference rankings.With less than 16 KB per example, it remains competitive with tested baselines allocated over 100 times as much storage.Source: §1 · p. 2

2 PRELIMINARIES

Influence functions estimate how removing a training example and retraining would change model behavior, using query and training gradients with curvature information. The preliminaries describe damped curvature for practical scoring and structured Fisher approximations that support scalable influence estimation.

2.1 INFLUENCE FUNCTIONS

  • Influence functions estimate how a model’s behavior would change if a training example were removed and the model retrained.For a scalar target measurement, its query gradient is evaluated at the trained parameters; a target-document loss is one example.Source: §2.1 · p. 2
  • Because a neural-network loss Hessian can be singular, practical influence estimation uses a damped generalized Gauss–Newton matrix, which is positive definite.The passage defines the damped matrix as H_λ := G + λI_d for λ > 0.Source: §2.1 · p. 2
  • For a fixed model, training-example gradients do not depend on the query and can be stored for later influence calculations without repeated backward passes over training data.The stored gradients are reused across subsequent queries.Source: §2.1 · p. 2

2.2 KRONECKER-FACTORED CURVATURE APPROXIMATION

  • K-FAC approximates the Fisher with layerwise Kronecker factors from activation and output-gradient second moments, using smaller eigendecompositions to obtain its eigenbasis.Its block-diagonal form neglects cross-layer blocks.Source: §2.2 · p. 3
  • EK-FAC corrects K-FAC eigenvalues while retaining its layerwise eigenbases, and Figure 1 shows EOGP projecting and half-whitening gradients before PCA correction and one-bit storage.Query gradients undergo the same transformations without one-bit quantization; influence scores use scaled inner products with stored bits.Source: §2.2 · p. 3, Figure 1 · p. 3

3 METHOD

EOGP compresses reusable training gradients to preserve influence estimates for queries not known at storage time. It combines a worst-case-optimal representation with an EK-FAC-based approximation, PCA refinement, and one-bit storage.

3.1 WHAT TO STORE: PCA OF HALF-WHITENED GRADIENTS

  • Top-k PCA coordinates of half-whitened training gradients minimize worst-case influence-estimation error among k-dimensional linear representations.The criterion averages squared error over training examples and considers query gradients in a bounded set.Source: §3.1 · p. 4

3.2 TWO-STAGE PROJECTION: PCA WITHIN AN EK-FAC SUBSPACE

  • EK-FAC corrected eigenvectors select m candidate axes via a Fisher-based surrogate for training-gradient second moments; PCA on projected gradients then chooses k directions within their span.The surrogate is distinct from the half-whitening approximation: Fisher averages over predictive-label distributions, while the empirical second moment uses training labels.Source: §3.2 · p. 5
  • Within the retained span, PCA correction has no greater mean squared reconstruction error than the leading k EK-FAC axes, and matches full-space k-PCA when the span contains a top-k principal subspace.Retaining m > k candidate axes lets the final subspace be optimized while storing only k coordinates per example.Source: §3.2 · p. 5
  • For larger k, EOGP-R replaces dense PCA correction with an implicit SRHT, trading fixed-k reconstruction optimality for scalability while controlling influence distortion with high probability.The stated guarantee applies when k is sufficiently large and representations lie in a fixed low-dimensional subspace.Source: §3.2 · p. 6

3.3 ONE-BIT QUANTIZATION

  • One-bit storage retains projected-coordinate signs plus one scale per module, while query coordinates stay unquantized; influence is computed without recomputing training gradients.Each module stores sign bits and a half-precision scale per training example.Source: §3.3 · p. 6

4 EXPERIMENTS

The experiments test compact gradient stores against retraining on GPT-2, EK-FAC fidelity on OLMo 2 SFT, component ablations, and storage and computation costs. EOGP improves attribution quality over the evaluated compression baselines across these evaluations.

4.1 VALIDATION AGAINST RETRAINING

  • At 96 KB, EOGP exceeds every compression baseline at 1,536 KB on GPT-2 LDS; at 384 KB, its LDS is close to uncompressed K-FAC using over 400 times less storage.The counterfactual experiment also finds larger validation-perplexity increases for examples selected by EOGP than by compression baselines.Source: §4.1 · p. 7, Figure 2 · p. 7

4.2 FIDELITY TO EK-FAC AT LLM SCALE

  • On OLMo 2 SFT models from 1B to 32B, EOGP improves EK-FAC ranking fidelity over compression baselines, with gains at small budgets and across overall and top-ranked examples.On GPT-2, EK-FAC agreement correlates positively with LDS: ρ=0.95 for NDCG@20 and ρ=0.84 for overall ranking agreement.Source: §4.2 · pp. 7–8, Figure 3 · p. 7, Figure 4 · p. 8

4.3 ABLATION STUDIES

  • At matched storage budgets, EOGP’s PCA-refined basis has higher NDCG@20 than direct EK-FAC basis selection; at matched coordinate counts, it beats same-precision baselines under both FP16 and one-bit storage.Within EOGP, one-bit and FP16 representations achieve similar NDCG@20 at matched coordinate counts.Source: §4.3 · pp. 8–9, Figure 5 · p. 8

4.4 STORAGE AND COMPUTATIONAL COSTS

  • At 2,048 coordinates per module, the 32B model’s per-example payload is approximately 115.6 KB versus 62.4 GB for an uncompressed half-precision gradient.Shared PCA correction matrices require approximately 1 GB per module at the stated dimensions, but shared artifacts do not grow with the number of stored examples; EOGP store construction takes 0.247 seconds per example versus 0.278 seconds for LoGra on one NVIDIA B200.Source: §4.4 · p. 9

5 RELATED WORK

Related work spans scalable influence estimation and gradient-compression methods for attribution. The paper situates EOGP among curvature-based influence methods and approaches that compress or project gradients.

Scalable influence estimation

  • Arnoldi iteration approximates a dominant Hessian eigenspace, while K-FAC and EK-FAC use structured curvature approximations for influence estimation.EK-FAC has enabled influence analysis at LLM scale (Grosse et al., 2023).Source: §5 · p. 9

Gradient compression for data attribution

  • Gradient attribution methods include random projections, curvature-corrected retrieval, Kronecker-factored projections, sparsification, and low-rank storage.The cited approaches include TRAK, TrackStar, LoGra, GraSS, and LoRIF; LESS and QLESS select instruction data using projected gradient similarities rather than inverse-curvature-based influence scores.Source: §5 · p. 9
  • The related methods differ in their compression or attribution strategy, including whether they estimate inverse-curvature-based influence scores or use projected gradient similarities.LESS and QLESS use projected gradient similarities, whereas the passage describes the studied scores as inverse-curvature-based.Source: §5 · p. 9

6 CONCLUSION

The paper presents EOGP as a compact reusable gradient store for influence estimation, approximating an optimal fixed-dimensional linear representation with EK-FAC subspace refinement and one-bit quantization. Its appendix establishes the half-whitened PCA optimum under the stated worst-case criterion.

6 CONCLUSION

  • EOGP approximates the optimal fixed-dimensional linear representation by applying PCA within a retained EK-FAC subspace and using scaled one-bit quantization to retain more coordinates per storage budget.The conclusion reports that EOGP preserves useful attribution information with substantially less storage than existing compression methods.Source: §6 · pp. 9–10
  • Under a normalized worst-case influence-error criterion, top-k PCA coordinates of half-whitened training gradients are optimal among k-dimensional linear representations.The conclusion frames this result as the target that EOGP approximates at LLM scale.Source: §6 · p. 9

A.1 OPTIMALITY OF HALF-WHITENED PCA

  • For the appendix’s fixed positive-definite Hλ setting, the optimal error is the sum of eigenvalues beyond k, attained by the proposed encoder–decoder pair.The proof states that the result uses the uncentered second moment and remains valid with repeated eigenvalues.Source: §A.1 · pp. 12–13

A.2 PROOF OF EIGENBASIS CORRECTION

The proof shows that PCA refinement within the retained EK-FAC subspace is no worse than using its leading axes, and derives the optimal scale for one-bit coordinates.

A.2 PROOF OF EIGENBASIS CORRECTION

  • PCA refinement minimizes reconstruction error among k-dimensional subspaces of the retained EK-FAC range, so the leading k EK-FAC axes are a feasible alternative.The proof decomposes the error into a term independent of the candidate basis and a term minimized by PCA.Source: §A.2 · p. 13
  • If the retained range contains a globally optimal k-dimensional principal subspace, its restricted minimum error equals the full-space minimum.The restricted feasible set makes its minimum no smaller, while containing the global optimum makes it no larger.Source: §A.2 · p. 13

A.3 SCALED ONE-BIT QUANTIZATION

  • For a fixed sign vector b = sign(a), the squared reconstruction error is minimized by scale s = ∥a∥1/k, including s = 0 when a = 0.The derivation uses b⊤a = ∥a∥1 and b⊤b = k.Source: §A.3 · p. 13

B ALGORITHMS AND PROJECTION IMPLEMENTATION

The procedure uses module-wise EK-FAC projections followed by PCA for EOGP or SRHT for EOGP-R, then stores quantized training coordinates for query scoring. PCA fitting and projection costs are also specified.

B.1 PROJECTION PATHS AND UNIFIED PROCEDURE

  • For each module, EOGP applies PCA after the first-stage map, while EOGP-R substitutes an SRHT; training examples store packed coordinate signs and FP16 scales.The algorithm fits EOGP projections from cached training-gradient coordinates and computes query coordinates before scoring stored examples.Source: §B.1 · p. 14

B.2 PCA FITTING AND RANDOM PROJECTION

  • EOGP fits uncentered PCA to first-stage training coordinates and can recover nonzero directions from an n × n Gram matrix instead of an m_u × m_u second-moment matrix.The PCA directions are the top-k_u right singular vectors of the coordinate matrix.Source: §B.2 · p. 14
  • EOGP-R maps first-stage coordinates to k_u dimensions with an SRHT, reusing the same realized transform for training storage and query projection without fitting a PCA correction.The transform zero-pads coordinates, applies random signs and a Walsh–Hadamard matrix, then selects rows.Source: §B.2 · p. 15

B.3 COMPUTATION AND COMPLEXITY

  • After query projection, the same query coordinates are reused to score all stored examples against their stored signs and scales.The PCA correction costs O(m_u k_u) operations per module; the implicit SRHT costs O(p_u log p_u).Source: §B.3 · p. 15

C EXPERIMENTAL CONFIGURATION

C.1 MODELS, DATA, AND GRADIENT COMPUTATION

The experiments use GPT-2 on WikiText-2 and OLMo 2 SFT checkpoints on Tülu 3 data, with specified gradient, curvature-fitting, projection, storage, and baseline configurations.

C.1 MODELS, DATA, AND GRADIENT COMPUTATION

  • GPT-2 is evaluated on WikiText-2, while released OLMo 2 SFT checkpoints use Tülu 3 data; WikiText-2 text is divided into 512-token blocks and OLMo conversations are examples.Training blocks and conversations supply attribution candidates, while validation blocks and separately selected conversations supply queries.Source: §C.1 · p. 15
  • Training and query gradients use token-summed next-token cross-entropy, with OLMo 2 gradients restricted to assistant responses and turn-ending tokens.Attribution covers attention and MLP linear layers; embeddings, the output head, and normalization parameters are excluded.Source: §C.1 · p. 15

C.2 CURVATURE FITTING AND PROJECTION DIMENSIONS

  • EK-FAC statistics are estimated with Kronfluence, and the resulting half-whitening transformation is applied to both training and query gradients.The Fisher statistics use model-sampled labels, and projection dimensions are chosen to match per-example storage budgets across methods.Source: §C.2 · p. 16
  • For OLMo 2, PCA uses 10,000 separate SFT conversations, and EOGP’s final dimensions range from 256 to 8,192 across modules.EOGP and EOGP-R use a first-stage width of 262,144; EOGP-R outputs 256 to 32,768 coordinates.Source: §C.2 · p. 16

C.3 BASELINE CONFIGURATIONS

  • Baseline configurations include LoGra, GraSS, LoRIF, and Gaussian random projection, with specified shared fitting settings and randomization procedures.The baselines share attributed modules, evaluation examples, queries, and loss definitions with EOGP.Source: §C.3 · pp. 16–17

D.1 GPT-2 RETRAINING EVALUATION

The section details GPT-2 retraining evaluation and influence-fidelity measures, then describes storage accounting. At small per-example budgets, EOGP maintains an advantage over evaluated compression baselines.

D.1 GPT-2 RETRAINING EVALUATION

  • GPT-2 evaluation uses LDS from shared subset-retraining losses and counterfactual retraining after removing top-ranked training blocks.LDS averages query-wise Spearman correlations; counterfactual runs remove 50 to 300 examples and evaluate perplexity across five matched seeds.Source: §D.1 · p. 17

D.2 FIDELITY TO EK-FAC INFLUENCE

  • Influence fidelity compares compressed scores with uncompressed EK-FAC scores using Spearman correlation and NDCG@20 over candidate examples.NDCG@20 treats membership in the reference top 20 as binary relevance.Source: §D.2 · p. 18

D.3 RELATIONSHIP BETWEEN FIDELITY AND RETRAINING

  • The study computes correlations across 28 configurations between query-averaged LDS, NDCG@20, and Spearman fidelity.Source: §D.3 · p. 18

E.1 PER-EXAMPLE PAYLOAD AND SHARED ARTIFACTS

  • One-bit representations store packed signs and one FP16 scale per module; at uniform dimension k, the payload is |U|(⌈k/8⌉ + 2) bytes.The reported accounting separates shared artifacts from per-example representation payloads.Source: §E.1 · pp. 18–19

E.2 STORAGE ESTIMATES BY MODEL

  • Table 2 compares per-example payloads for uncompressed FP16 gradients and one-bit representations at k_u = 2,048 per module.Both EOGP and EOGP-R use 258 bytes per module per example at that dimension.Source: Table 2 · p. 19, §E.2 · p. 19

E.1 PER-EXAMPLE PAYLOAD AND SHARED ARTIFACTS

  • For the 32B model, shared PCA correction matrices total approximately 481 GB, with their relative storage share decreasing as the example store grows.The stated share is approximately 4.2% for 100 million examples and 0.42% for one billion examples.Source: §E.1 · p. 19
  • At small per-example storage budgets, EOGP maintains an advantage over the evaluated compression baselines.Source: §E.1 · p. 19

E.3 IMPLEMENTATION AND RUNTIME

On OLMo 2 7B SFT, EOGP substantially outperforms LoGra and GraSS variants at matched storage budgets, while the original baseline configurations generally outperform their half-whitened variants.

E.3 IMPLEMENTATION AND RUNTIME

  • Table 3 reports PCA fitting and store-construction times on one NVIDIA B200; store-construction timing includes gradient computation, projection, disk I/O, and correction-matrix transfer.The setup excludes curvature fitting and model loading; PCA projection and storage process batches of approximately 11,000 documents.Source: §E.3 · p. 19, Table 3 · p. 20

F COMPARISON UNDER SHARED EK-FAC CURVATURE

E.3 IMPLEMENTATION AND RUNTIME

  • 18.4 minutes is the one-time PCA correction-matrix fitting cost for the 32B model, shared across examples and queries.Fitting starts from cached coordinates and includes reading them, computing correction matrices, and saving them in FP16.Source: §E.3 · p. 20
  • 0.604 seconds per query is the amortized cost for processing 100 queries against 1,000 stored examples on a single NVIDIA B200.For the 32B model, this includes correction-matrix loading and transfer, PCA projection, and scoring.Source: §E.3 · p. 20

F COMPARISON UNDER SHARED EK-FAC CURVATURE

  • Original LoGra and GraSS configurations generally outperform their half-whitened variants in this setting.The results suggest that curvature-approximation effectiveness depends on how it is combined with gradient compression.Source: F COMPARISON UNDER SHARED EK-FAC… · p. 20
  • Figure 7 tracks GPT-2 LDS for one-bit and FP16 storage at matched budgets; Table 4 reports median query-wise Pearson correlations over 481 queries.Figure 7 curves end where a method exceeds its supported projection dimension.Source: Figure 7 · p. 21, Table 4 · p. 21

G EFFECT OF ONE-BIT QUANTIZATION ON ATTRIBUTION QUALITY

At matched storage budgets, EOGP achieves the highest LDS among evaluated methods under one-bit storage, while LoGra and GraSS LDS decline as budgets grow. Fixed-dimension score comparisons also show strong FP16-to-one-bit correlation for EOGP, though Pearson correlation does not measure absolute score magnitudes.

Attribution quality under one-bit storage

  • EOGP achieves the highest LDS among evaluated compression methods at each displayed one-bit storage budget, with performance improving as storage increases.One-bit LoGra and GraSS LDS declines as budgets grow; LoRIF improves more gradually but remains below EOGP, and EOGP also beats other FP16 baselines at larger budgets.Source: G EFFECT OF ONE-BIT QUANTIZATION ON… · p. 21

Preservation of influence scores

  • EOGP’s FP16-to-one-bit influence-score correlations range from 0.890 to 0.919 across configurations, while LoGra and GraSS have substantially weaker agreement in several configurations.The comparison holds projection dimension fixed and reports median query-wise Pearson correlations.Source: G EFFECT OF ONE-BIT QUANTIZATION ON… · p. 21
  • Pearson correlation preserves score-pattern comparisons, not absolute magnitudes, because positive rescaling or an additive shift leaves it unchanged.Together with LDS results, the correlations provide complementary evidence that EOGP preserves useful attribution information under one-bit storage.Source: G EFFECT OF ONE-BIT QUANTIZATION ON… · pp. 21–22

A possible explanation

  • One possible explanation is that half-whitening reduces coordinate-magnitude imbalance, making representations more compatible with scaled sign quantization.The quantizer preserves signs and replaces coordinate magnitudes with a common scale based on their mean absolute value.Source: G EFFECT OF ONE-BIT QUANTIZATION ON… · p. 22

Related papers

Loading 2609.37842v1…