Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

GraSS: Scalable Data Attribution with Gradient Sparsification and Sparse Projection

Pingbang Hu, Joseph Melkonian, Weijing Tang, Han Zhao, Jiaqi W. Ma

TL;DR

The paper addresses the scalability of gradient-based data attribution by exploiting per-sample gradient sparsity to achieve sub-linear complexity. It introduces GraSS, a two-stage compression that sparsifies gradients and then applies a sparse projection (SJLT), and FactGraSS, which extends this idea to linear layers using gradient factorization (Kronecker structure) to avoid materializing full gradients. Together, these methods achieve substantial speedups over prior baselines while preserving attribution fidelity, with reports of up to 165% faster throughput on billion-scale models and successful application to GPT2-small and Llama-3.1-8B-Instruct, illustrating practical benefits for large-scale data attribution. The work broadens the feasibility of robust data attribution in real-world, large-model settings by reducing memory and compute bottlenecks without sacrificing accuracy.

Abstract

Gradient-based data attribution methods, such as influence functions, are critical for understanding the impact of individual training samples without requiring repeated model retraining. However, their scalability is often limited by the high computational and memory costs associated with per-sample gradient computation. In this work, we propose GraSS, a novel gradient compression algorithm and its variants FactGraSS for linear layers specifically, that explicitly leverage the inherent sparsity of per-sample gradients to achieve sub-linear space and time complexity. Extensive experiments demonstrate the effectiveness of our approach, achieving substantial speedups while preserving data influence fidelity. In particular, FactGraSS achieves up to 165% faster throughput on billion-scale models compared to the previous state-of-the-art baselines. Our code is publicly available at https://github.com/TRAIS-Lab/GraSS.

GraSS: Scalable Data Attribution with Gradient Sparsification and Sparse Projection

TL;DR

The paper addresses the scalability of gradient-based data attribution by exploiting per-sample gradient sparsity to achieve sub-linear complexity. It introduces GraSS, a two-stage compression that sparsifies gradients and then applies a sparse projection (SJLT), and FactGraSS, which extends this idea to linear layers using gradient factorization (Kronecker structure) to avoid materializing full gradients. Together, these methods achieve substantial speedups over prior baselines while preserving attribution fidelity, with reports of up to 165% faster throughput on billion-scale models and successful application to GPT2-small and Llama-3.1-8B-Instruct, illustrating practical benefits for large-scale data attribution. The work broadens the feasibility of robust data attribution in real-world, large-model settings by reducing memory and compute bottlenecks without sacrificing accuracy.

Abstract

Gradient-based data attribution methods, such as influence functions, are critical for understanding the impact of individual training samples without requiring repeated model retraining. However, their scalability is often limited by the high computational and memory costs associated with per-sample gradient computation. In this work, we propose GraSS, a novel gradient compression algorithm and its variants FactGraSS for linear layers specifically, that explicitly leverage the inherent sparsity of per-sample gradients to achieve sub-linear space and time complexity. Extensive experiments demonstrate the effectiveness of our approach, achieving substantial speedups while preserving data influence fidelity. In particular, FactGraSS achieves up to 165% faster throughput on billion-scale models compared to the previous state-of-the-art baselines. Our code is publicly available at https://github.com/TRAIS-Lab/GraSS.

Paper Structure

This paper contains 46 sections, 1 theorem, 6 equations, 9 figures, 3 tables.

Table of Contents

  1. Introduction
  2. Preliminary
  3. Influence function
  4. Random projection
  5. GraSS: Gradient Sparsification and Sparse projection
  6. Per-sample gradient sparsity
  7. Sparse Johnson-Lindenstrauss Transform.
  8. Implementation of SJLT.
  9. Effective parameter sparsity
  10. Random Mask.
  11. Selective Mask.
  12. GraSS & FactGraSS: Multi-stage compression
  13. GraSS: Sparsify first, sparse projection next
  14. FactGraSS: Exploiting layer-wise gradient factorization structure
  15. Recap on LoGra.
  16. ...and 31 more sections

Key Result

Theorem 3.2

There is a sub-linear compression algorithm with a complexity of $O(k^{\prime} )$ per sample, where $k < k^{\prime} \ll p$. Moreover, this extends to linear layers, where gradients are never materialized.

Figures (9)

  • Figure 1: Dense Rademacher projection.
  • Figure 2: Sparse Rademacher projection ($s = 1$).
  • Figure 3: SJLT with $s = 1$.
  • Figure 4: Benchmark of different projection methods with $p = 131,072$ under several sparsity levels (distinguished by opacity). Relative error is w.r.t. pair-wise distance preservation.
  • Figure 5: Mask.
  • ...and 4 more figures

Theorems & Definitions (2)

  • Remark 3.1
  • Theorem 3.2