Research map / Architectures and efficiency / RNNs and normalization
Efficient Sketches for Training Data Attribution and Studying the Loss Landscape
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 scalable gradient and Hessian-vector-product sketching algorithms (AFFD, AFJL, QK) that avoid lookup bottlenecks on TPUs/GPUs, and uses them to study training data attribution, intrinsic dimension, and LLM Hessian spectra.
Problem
Storing many gradients or Hessian-vector products for training data attribution, intrinsic dimension estimation, and Hessian eigenvalue analysis is memory-prohibitive; existing sketches (dense projections, FJL, Fastfood) scale poorly on modern accelerators.
Approach
Replace lookup-based sparse projections and Fastfood with randomized Hadamard/Kronecker-structured preconditioners (AFFD, AFJL) and random Haar-orthogonal Kronecker factors (QK), with theoretical concentration guarantees and explicit (rather than implicit) sketching for faster wall time.
Key results
- AFJL reaches 98% correlation with full-gradient influence scores at D=2^13 on GPT-2, while layer-selected sketches rarely exceed 90%.
- AFJL achieves about 100x wall-time reduction over FJL on TPU, and AFFD about 64x over FFD on TPU (Table of design choices).
- Intrinsic dimension search on XSUM gives D_int around 2^25 for Rouge1 and close to full model dimension (~2^27) for Rouge2, while SNLI classification gives about 2^13.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace FJL with AFJL to remove lookup-based operations | GPT-2 gradient sketch on TPU (v2) · GPT-2 gradient sketch, TPU (v2) | wall time: 8997 -> 89 ms | Helped |
| Replace FFD with AFFD to remove lookup-based operations | GPT-2 gradient sketch on TPU (v2) · GPT-2 gradient sketch, TPU (v2) | wall time: 8694 -> 134 ms | Helped |
| Use explicit rather than implicit sketching (ablation) | gradient sketch algorithms · Gradient sketching, GPU (V100) | wall time speed-up ratio: 1.96; 1.81; 1.76; 1.64; 1.41 | Helped |
| Replace the Walsh-Hadamard pre-conditioner with FFT (ablation) | AFFD and AFJL · Gradient sketching, GPU (V100) | wall time speed-up ratio: 1.98; 1.64 | Helped |
| Replace the Walsh-Hadamard pre-conditioner with Q (ablation) | AFFD and AFJL · Gradient sketching, GPU (V100) | wall time speed-up ratio: 1.45; 1.44 | Helped |
| Remove lookup operations through accelerator-oriented sketch design (ablation) | FJL and FFD · Gradient sketching, GPU (V100) | peak memory: 2.4x for FJL to AFJL; 1.3x for FFD to AFFD | Helped |
| Use QK instead of AFJL | GPT-2 gradient sketch on GPU (V100) · GPT-2 gradient sketch, GPU (V100) | wall time: 116 -> 82 ms | Helped |
| Use QK instead of AFJL | GPT-2 HVP sketch on TPU (v2) · GPT-2 HVP sketch, TPU (v2) | wall time: 108 -> 82 ms | Helped |
| Restrict gradients to a selected layer instead of using full gradients (ablation) | GPT-2 training-data attribution · WikiText-103 TDA gradient inner products | Pearson correlation: 0.16; 0.75; 0.89; 0.90; 0.89; 0.78; 0.38 | Hurt |
| Use AFJL sketches rather than full gradients for TDA scores | GPT-2 training-data attribution · WikiText-103 TDA gradient inner products | Pearson correlation: 98% | Helped |
| Use a larger target dimension for QK than AFFD | gradient inner-product sketching · TDA gradient inner products | minimum log2 D for r > 0.95: 10 -> 14 | Hurt |
| Use a larger target dimension for QK than AFFD | Hessian top 10 eigenvalue reconstruction · Hessian eigenvalue estimation | minimum log2 D for relative error < 0.05: 10 -> 13 | Hurt |
Surprising findings
- Layer selection for training data attribution yields unreliable influence-score correlations, with most layers below 90%.
- Hessian spectra of pre-trained LLMs (RoBERTa, GPT-2L, BART) deviate from conjectures based on smaller networks, e.g., RNEG is inconsistent and outlier counts do not follow K-1 rule.
Limitations
- Results may vary for model architectures other than Transformers.
- Generative intrinsic dimension study is limited to one summarization task, with in-depth study left for future work.
- QK requires larger target dimension than AFFD and AFJL to reach comparable sketch quality.
Experimental setup
Datasets: WikiText-103 zsRE SNLI XSUM
Benchmarks: SNLI XSUM WikiText-103 zsRE
Metrics: Pearson correlation wall-time (ms) peak memory (GB) relative mean absolute error of eigenvalues Rouge1 Rouge2 accuracy
Baselines: FJL FFD TRAK dense random projection layer selection
Builds on: Fastfood transform Fast Johnson-Lindenstrauss transform TRAK intrinsic dimension of fine-tuning (Aghajanyan et al.)
Closest papers by meaning
- SketchINR: A First Look into Sketches as Implicit Neural Representations (CVPR 2024 · similarity 0.87)
- Sketching for Distributed Deep Learning: A Sharper Analysis (NeurIPS 2024 · similarity 0.86)
- Lego Sketch: A Scalable Memory-augmented Neural Network for Sketching Data Streams (ICML 2025 · similarity 0.85)
- CTSketch: Compositional Tensor Sketching for Scalable Neurosymbolic Learning (NeurIPS 2025 · similarity 0.85)
- Sketched Adaptive Distributed Deep Learning: A Sharp Convergence Analysis (NeurIPS 2025 · similarity 0.84)
- SketchDNN: Joint Continuous-Discrete Diffusion for CAD Sketch Generation (ICML 2025 · 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.