Sketching the Readout of Large Language Models for Scalable Data Attribution and Valuation
| Authors | Yide Ran et al. |
| Year | 2026 |
| Field | Machine Learning |
| arXiv | 2604.16197 |
| Download | |
| Categories | cs.LG |
Abstract
Data attribution and valuation are critical for understanding data-model synergy for Large Language Models (LLMs), yet existing gradient-based methods suffer from scalability challenges on LLMs. Inspired by human cognition, where decision making relies on a focused readout of relevant memories rather than replaying all pathways, we introduce RISE (Readout Influence Sketching Estimator). Instead of computing and indexing gradients across the entire LLM, RISE focuses on influence hotspots at the output layer, where influence signals concentrate, and the gradient admits a decomposed outer-product form. This enables a dual-channel representation combining a lexical residual channel (RH) and a semantic projected-error channel (GH). Applying CountSketch projections to these channels achieves strong compression while maintaining accurate attribution. Across the OLMo (1B-32B) and Pythia (14M-6.9B) families, RISE reduces index storage by up to 112 compared to RapidIn and scales to 32B parameters LLM, where gradient-based baselines such as RapidIn and ZO-Inf become memory-infeasible. We evaluate RISE on two paradigms: (1) retrospective attribution, retrieving influential training examples for specific predictions, and (2) prospective valuation, scoring candidate data utility zero-shot. We validate RISE on three tasks: Howdy backdoor data detection, Finance-Medical domain separation, and Brain Rot high-quality data selection. In a closed-loop Brain Rot study, continued pretraining on RISE-selected data yields consistent downstream improvements. Overall, RISE provides a practical and scalable primitive for influence analysis and training-data selection in modern large language models.
Engineering Breakdown
Plain English
This paper addresses a fundamental problem in LLM interpretability: understanding which training data points most influenced a model's predictions. Existing gradient-based attribution methods don't scale to LLMs because they require computing and storing gradients for every training example, which is computationally prohibitive. RISE solves this by focusing computational effort only on influence hotspots at the output layer where signals concentrate, using a clever decomposition into lexical and semantic channels combined with CountSketch projections. This allows data attribution for LLMs with orders of magnitude better computational efficiency than prior methods.
Core Technical Contribution
The core novelty is identifying and exploiting that influence computations can be dramatically compressed by focusing only on the output layer rather than backpropagating through the entire network. The authors discovered that at the output layer, gradients admit a decomposed outer-product form that can be efficiently represented through a dual-channel architecture: a lexical residual channel (RH) capturing token-level information and a semantic projected-error channel (GH) capturing learned representations. By applying CountSketch random projections to these channels, they achieve sublinear memory and computation complexity compared to storing full gradients. This is fundamentally different from prior work that either computed full gradients throughout the network or used approximations that lost critical information.
How It Works
RISE operates by first recognizing that at an LLM's output layer, the influence signal concentrates in a structured form rather than scattered across the network. When computing how much a training example influences a test prediction, instead of backpropagating gradients through all layers, RISE computes two efficient representations: the lexical residual channel (RH) which directly operates on token embeddings and captures surface-level patterns, and the semantic projected-error channel (GH) which operates on the model's learned intermediate representations. Both channels are then projected into a lower-dimensional space using CountSketch, a probabilistic data structure that preserves inner products with high probability while using logarithmic space. The influence score for each training example is computed by querying these sketched representations, avoiding the need to store or index full gradient vectors for every training example.
Production Impact
For teams building LLM systems that need to understand model behavior, debug failures, or comply with data provenance requirements, RISE enables practical data attribution at production scale. Current gradient-based methods require storing dense gradient vectors proportional to model size × dataset size, which is infeasible for billion-parameter models and multi-billion token datasets; RISE reduces this to sketch vectors of logarithmic size. This enables use cases like identifying poisoned training data, understanding which pretraining examples caused hallucinations, and optimizing dataset composition without requiring GPUs with terabytes of memory or weeks of computation. The trade-off is that you get approximate rather than exact influence scores, and the approach assumes influence concentrates at the output layer—which holds well for supervised tasks but may be less accurate for multi-task or instruction-tuned models where different examples influence different output dimensions differently.
Limitations and When Not to Use This
The paper's core assumption—that influence concentrates at the output layer—may not hold equally well for all LLM architectures or fine-tuning scenarios. For models with sparse, multi-objective training objectives or where different examples heavily influence intermediate representations used in subsequent layers, focusing only on output-layer gradients could miss significant influence signals. The CountSketch approximation introduces error that accumulates with dataset size, and the paper doesn't fully characterize how approximation quality degrades with sketch size or provide principled guidance on choosing sketch dimensions for different scenarios. Additionally, the approach requires access to model gradients at inference time, which conflicts with the current industry trend toward serving LLMs as opaque APIs; it's unclear how RISE would apply to attribution tasks over proprietary or quantized models.
Research Context
This work builds on decades of influence function research starting from Koh & Liang's 2017 foundational work on data valuation through influence functions, but recognizes that naive gradient-based approaches simply don't scale to modern LLMs. It connects to the broader interpretability and mechanistic understanding literature investigating how neural networks make decisions, as well as emerging work on data quality and curation for LLMs. The paper advances the state of practical data attribution by making it computationally feasible for billion-parameter models, opening new research directions in adversarial example detection, training data debugging, and automated dataset optimization. This positions RISE as a key enabler for the emerging field of data-centric AI, where understanding and managing training data becomes as important as model architecture.
:::tip Subscribe Get weekly breakdowns of papers like this in AI Letters - the newsletter for engineers building production AI systems. :::