Language model context attribution with fewer evaluations using sketching
Sketching Language Models for Context Attribution
Data Structures and AlgorithmsArtificial IntelligenceMachine Learning
Summary
Understanding how parts of a prompt affect a language model's response can be slow because checking each part separately needs many model runs. The authors created an algorithm that estimates the effect of removing each token with far fewer model calls, assuming the model behaves stably. They also proved that this efficiency is near the best possible and tested their method on a simple language task. The approach uses advanced math techniques applied in a new way to speed up calculations.
What this means in practice
- •For language model engineers: Reduce the number of model evaluations needed for context attribution analyses in transformer-based models.
- •For algorithm designers: Implement efficient sketching algorithms for sensitive analysis of arithmetic circuits and differential computations in machine learning.
Authors
Christopher Liaw, Aranyak Mehta, Grigoris Velegkas, Felix Zhou
Abstract
Context attribution asks how removing parts of a prompt changes a language model's output. We study leave-one-out (LOO) attribution for a prompt of $n$ tokens and a fixed measurement function. Computing all LOO values directly requires $n$ model evaluations. Under a stability assumption, our algorithm estimates all LOO values to error proportional to $\varepsilon$ with constant probability using $O(\sqrt{\log(1/\varepsilon)}/\varepsilon)$ evaluation calls. To complement our upper bound, we prove a matching $Ω(\sqrt{\log(1/\varepsilon)}/\varepsilon)$ lower bound on the number of evaluation calls in an oracle model. Empirically, we verify coefficient decay consistent with stability and evaluate the performance of our sketching algorithm on a small language-model retrieval task. Technically, our result is based on a sketching algorithm for arithmetic circuits whose Taylor coefficients satisfy the stability condition recently introduced by Gunn [Gun26]. A direct adaptation of their forward differentiation sketches would require $\widetilde{O}(\log(n)/\varepsilon^2)$ evaluation calls. Our key observation is that we can apply forward differentiation to the backpropagation circuit, with only constant overhead over a forward sketch of the same Taylor degree. This reduces the variance of the sketches and is key to our improved evaluation bound. Our lower bound builds on the techniques of Amsel, Chen, Duman Keles, Halikias, Musco, and Musco [ACDH+26, SIAM Journal on Matrix Analysis and Applications] for lower bounding the number of matrix-vector queries to estimate the diagonal elements of random matrices.