RISE: Sketching the Readout of LLMs for Scalable Data Attribution
Sketching the Readout of Large Language Models for Scalable Data Attribution and Valuation
This paper introduces RISE (Readout Influence Sketching Estimator), a scalable data attribution and valuation method for Large Language Models (LLMs). By focusing on gradient energy hotspots at the LM-head and utilizing a dual-channel CountSketch compression, RISE achieves state-of-the-art performance, scaling to 32B models where traditional gradient-based baselines become memory-infeasible.
TL;DR
Understanding "why" an LLM makes a specific prediction usually requires tracing back through trillions of parameters—a task that is computationally bankrupt at scale. RISE (Readout Influence Sketching Estimator) changes the game by proving that we only need to look at the LM-head (the readout layer). By decomposing gradients into lexical and semantic channels and sketching them, RISE scales to 32B+ models, providing a 112× reduction in storage while outperforming full-gradient methods.
The Scalability Wall in Data Attribution
In the quest to understand data-model synergy, Influence Functions are the "Gold Standard." They theoretically tell us how a specific training point affected the prediction of query . However, classic methods require the Inverse Hessian—a matrix so large it's impossible to compute for LLMs.
Even first-order surrogates like TracIn require per-sample gradients. For a 7B model, storing these for 10k examples would eat 140TB of memory. This "Scalability Wall" has forced researchers to use smaller proxy models, which often don't align with the true dynamics of the large target model.
The Insight: "Readout" is the Influence Hotspot
The authors of RISE found a "physics-like" property in LLMs: Gradient energy is not distributed equally. As models grow deeper, the relative influence signal concentrates at the final readout layer (LM-head).
As shown in the figure, for an OLMo-32B model, the LM-head exhibits energy peaks up to 27.8× higher than internal layers. This allows us to ignore the model's "body" and focus purely on the "head" for attribution.
Methodology: The Dual-Channel Decompression
The core of RISE lies in the mathematical beauty of the LM-head gradient. It can be factored exactly as:
abla_{W_{lm\_head}} \ell = r_t \otimes h_t$$ where $r_t$ is the prediction residual and $h_t$ is the final hidden state. RISE exploits this via two channels: 1. **RH Channel (Lexical)**: Uses the raw vocabulary-space residual. It captures exact token-level matches. 2. **GH Channel (Semantic)**: Projects the residual back into the embedding space. This "smooths" the error, allowing the model to recognize that "kitten" and "cat" are semantically identical even if the tokens differ. To make this storage-efficient, RISE uses **CountSketch**. Unlike PCA, CountSketch is data-independent and hardware-friendly, allowing for "scatter-add" updates that are extremely fast on GPUs.  --- ## Experimental Results: Scaling to the Unfeasible RISE was tested against heavyweights like RapidIn and ZO-Inf. While the baselines crashed (Out of Memory) on OLMo-32B, RISE hummed along with a tiny memory footprint. | Model | Method | auPRC (Top-5) | Disk | GPU Mem | | :--- | :--- | :--- | :--- | :--- | | Pythia-1B | RapidIn | 0.190 | 1.6 GB | 14.2 GB | | **Pythia-1B** | **RISE** | **0.996** | **60.8 MB** | **6.7 GB** | | **OLMo-32B** | **RISE** | **0.993** | **41.9 MB** | **72.8 GB** | | OLMo-32B | RapidIn | N/A | OOM | OOM | ### The "Brain Rot" Study Beyond academic benchmarks, RISE was used for **Data Curation**. In the "Brain Rot" experiment, RISE was tasked with picking high-quality data from a pool of 90% "junk" (web-scraped noise). * **Purity**: RISE achieved 87.6% selection purity. * **Perplexity**: Training on RISE-selected data resulted in a 2.8× lower perplexity than the next best baseline. --- ## Critical Analysis & Conclusion RISE provides a powerful primitive for "data-centric AI." Its ability to perform **Prospective Valuation** (scoring new data before training) makes it a vital tool for the next generation of continuous pretraining. **Limitations**: While the readout layer captures most of the signal, it does ignore internal layer "hotspots" that might occur in specific architectures. Furthermore, the GH channel's fidelity depends on the quality of the LM-head embeddings themselves. **Future Work**: The authors suggest applying RISE to multi-modal models. If influence signals concentrate at the "readout" of a vision-language model similarly, we could radically optimize the way we curate massive image-text pairs. **Takeaway**: If you want to know what your LLM is "thinking" about its training data, stop looking at the 175B weights in the middle—just check the readout.