Research map / Architectures and efficiency / RNNs and normalization
GraSS: Scalable Data Attribution with Gradient Sparsification and Sparse Projection
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes GraSS, a two-stage per-sample gradient compression (sparsify then sparse JL projection) and FactGraSS for linear layers, achieving O(k') compression cost for data attribution on influence functions and TRAK.
Problem
Gradient-based data attribution (influence functions, TRAK) is bottlenecked by the cost of computing and compressing high-dimensional per-sample gradients for large models.
Approach
Exploit the inherent sparsity of per-sample gradients and effective parameter sparsity: first sparsify gradients (random or selective mask) to k' dimensions, then apply a sparse Johnson-Lindenstrauss transform (SJLT, s=1), giving sub-linear O(k') cost; FactGraSS applies this factorized over the Kronecker structure of linear-layer gradients so the full gradient is never materialized.
Key results
- FactGraSS achieves up to 165% faster compression throughput than LoGra on Llama-3.1-8B-Instruct (72,218 vs 27,292 tokens/s at k_l=256).
- On GPT2-small/WikiText, FactGraSS keeps LDS close to SJLT while achieving up to 250% speedup over LoGra.
- On MLP/MNIST with TRAK, SJLT reaches LDS 0.4357 at k=8192 with compression time far below Gaussian projection (10.8s), and Selective Mask achieves LDS 0.4373 at k=8192.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use Selective Mask instead of Random Mask for gradient sparsification (ablation) | Trak on MLP · MNIST | LDS: 0.4318 -> 0.4373 | Helped |
| Use SJLT instead of FJLT for projection (ablation) | Trak on MLP · MNIST | compression time (s): 0.8997 -> 0.4919 | Helped |
| Add Random Mask before SJLT (ablation) | Trak on ResNet9 · CIFAR2 | compression time (s): 12.3590 -> 0.3652 | Helped |
| Add Selective Mask before SJLT (ablation) | Trak on ResNet9 · CIFAR2 | compression time (s): 12.3590 -> 0.3054 | Helped |
| Add Random Mask before SJLT (ablation) | Trak on Music Transformer · MAESTRO | LDS: 0.3861 -> 0.3755 | Hurt |
| Add Random Mask before SJLT (ablation) | Trak on Music Transformer · MAESTRO | compression time (s): 21.3192 -> 0.7433 | Helped |
| Use SJLT instead of Gaussian projection (ablation) | Influence function on GPT2-small · WikiText | compression time (s): 20.9835 -> 133.4029 | Hurt |
| Add factorized Random Mask before SJLT (ablation) | Influence function on GPT2-small · WikiText | compression time (s): 20.4839 -> 6.3075 | Helped |
| Replace LoGra with FactGraSS | Llama-3.1-8B-Instruct · OpenWebText, compression throughput | tokens per second: 27,292 -> 72,218; 160% faster throughput | Helped |
| Replace LoGra with FactGraSS | Llama-3.1-8B-Instruct · OpenWebText, cache throughput | tokens per second: 7,307 -> 8,584; around 17% | Helped |
Surprising findings
- Standalone random sparsification (Random Mask) achieves non-trivial LDS results on MLP/MNIST and GPT2-small, suggesting overparameterized models tolerate heavy parameter dropping.
- Replacing LoGra's Gaussian projection with SJLT alone degrades efficiency in the small-problem regime, so naive integration of sparse projection does not speed up attribution.
Limitations
- LDS relies on an additivity assumption known to be imperfect.
- Random Mask and Selective Mask lose information and sparsification alone falls short in LDS on complex models (ResNet9, Music Transformer).
- Qualitative evaluation of Llama-3.1-8B attribution is limited to a single prompt and manually filtered top-10 examples; the attribution dataset is only an assumed proxy for the unreleased pretraining data.
Experimental setup
Datasets: MNIST CIFAR2 MAESTRO WikiText OpenWebText
Benchmarks: MNIST CIFAR2 MAESTRO WikiText
Metrics: linear datamodeling score wall-clock time throughput (tokens per second)
Baselines: Gaussian random projection FJLT LoGra SJLT TRAK GradDot
Builds on: TRAK LoGra fast Johnson-Lindenstrauss transform linear datamodeling score
Closest papers by meaning
- LoRIF: Low-Rank Influence Functions for Scalable Training Data Attribution (NeurIPS 2026 · similarity 0.87)
- Final-Model-Only Data Attribution with a Unifying View of Gradient-Based Methods (NeurIPS 2025 · similarity 0.87)
- Fast Data Attribution for Text-to-Image Models (NeurIPS 2025 · similarity 0.86)
- Enhancing Training Data Attribution with Representational Optimization (NeurIPS 2025 · similarity 0.85)
- Leveraging partial stragglers within gradient coding (NeurIPS 2024 · similarity 0.85)
- Sparse maximal update parameterization: A holistic approach to sparse training dynamics (NeurIPS 2024 · similarity 0.84)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.