# Efficient Estimation of Kernel Surrogate Models for Task Attribution

arXiv: https://arxiv.org/abs/2602.03783 · HTML: https://mlsanity.com/paper/2602.03783/

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Placed in Reinforcement learning for reasoning, one of 17 directions in LLM reasoning.
- Builds on 31 of its 44 references that are in the map.
- First posted on arXiv 3 Feb 2026.

## Placed in

- LLM reasoning / [Reinforcement learning for reasoning](https://mlsanity.com/lr/reinforcement-learning-for-reasoning/)

## Builds on (most cited first, 30 of 31)

- ["Why Should I Trust You?": Explaining the Predictions of Any Classifier](https://mlsanity.com/paper/1602.04938/) (2016)
- [Understanding Black-box Predictions via Influence Functions](https://mlsanity.com/paper/1703.04730/) (2017)
- [Rethinking the Role of Demonstrations: What Makes In-Context Learning Work?](https://mlsanity.com/paper/2202.12837/) (2022)
- [Sharpness-Aware Minimization for Efficiently Improving Generalization](https://mlsanity.com/paper/2010.01412/) (2020)
- [Meta-World: A Benchmark and Evaluation for Multi-Task and Meta Reinforcement Learning](https://mlsanity.com/paper/1910.10897/) (2019)
- [What Can Transformers Learn In-Context? A Case Study of Simple Function Classes](https://mlsanity.com/paper/2208.01066/) (2022)
- [Estimating Training Data Influence by Tracing Gradient Descent](https://mlsanity.com/paper/2002.08484/) (2020)
- [Meta-Learning with Implicit Gradients](https://mlsanity.com/paper/1909.04630/) (2019)
- [Studying Large Language Model Generalization with Influence Functions](https://mlsanity.com/paper/2308.03296/) (2023)
- [TRAK: Attributing Model Behavior at Scale](https://mlsanity.com/paper/2303.14186/) (2023)
- [Influence Functions in Deep Learning Are Fragile](https://mlsanity.com/paper/2006.14651/) (2020)
- [Datamodels: Predicting Predictions from Training Data](https://mlsanity.com/paper/2202.00622/) (2022)
- [DataInf: Efficiently Estimating Data Influence in LoRA-tuned LLMs and Diffusion Models](https://mlsanity.com/paper/2310.00902/) (2023)
- [A Kernel-Based View of Language Model Fine-Tuning](https://mlsanity.com/paper/2210.05643/) (2022)
- [Understanding and Improving Information Transfer in Multi-Task Learning](https://mlsanity.com/paper/2005.00944/) (2020)
- [When is Memorization of Irrelevant Training Data Necessary for High-Accuracy Learning?](https://mlsanity.com/paper/2012.06421/) (2020)
- [What is Your Data Worth to GPT? LLM-Scale Data Valuation with Influence Functions](https://mlsanity.com/paper/2405.13954/) (2024)
- [Heat kernel based community detection](https://mlsanity.com/paper/1403.3148/) (2014)
- [Precise High-Dimensional Asymptotics for Quantifying Heterogeneous Transfers](https://mlsanity.com/paper/2010.11750/) (2020)
- [Optimizing ML Training with Metagradient Descent](https://mlsanity.com/paper/2503.13751/) (2025)
- [MAGIC: Near-Optimal Data Attribution for Deep Learning](https://mlsanity.com/paper/2504.16430/) (2025)
- [Identification of Negative Transfers in Multitask Learning Using Surrogate Models](https://mlsanity.com/paper/2303.14582/) (2023)
- [Scalable Multitask Learning Using Gradient-based Estimation of Task Affinity](https://mlsanity.com/paper/2409.06091/) (2024)
- [ALPINE: Unveiling the Planning Capability of Autoregressive Learning in Language Models](https://mlsanity.com/paper/2405.09220/) (2024)
- [Noise Stability Optimization for Finding Flat Minima: A Hessian-based Regularization Approach](https://mlsanity.com/paper/2306.08553/) (2023)
- [Scalable Fine-tuning from Multiple Data Sources: A First-Order Approximation Approach](https://mlsanity.com/paper/2409.19458/) (2024)
- [Multi-group Learning for Hierarchical Groups](https://mlsanity.com/paper/2402.00258/) (2024)
- [Linear-Time Demonstration Selection for In-Context Learning via Gradient Estimation](https://mlsanity.com/paper/2508.19999/) (2025)
- [Scalable Multi-Objective and Meta Reinforcement Learning via Gradient Estimation](https://mlsanity.com/paper/2511.12779/) (2025)
- [A Versatile Influence Function for Data Attribution with Non-Decomposable Loss](https://mlsanity.com/paper/2412.01335/) (2024)

## Built on by (most cited first, 4 of 4)

- [WinQ: Accelerating Quantization-Aware Training of Language Models Around Saddle Points](https://mlsanity.com/paper/2605.17471/) (2026)
- [Learning Dynamic Representations and Policies from Multimodal Clinical Time-Series with Informative Missingness](https://mlsanity.com/paper/2604.21235/) (2026)
- [Long-Context Demonstration Selection Using State Space Models](https://mlsanity.com/paper/2609.17888/) (2026)
- [Every Batch Is Its Own Validation Set: Leave-One-Out Gradient Matching for Online Data Selection in LLM Fine-Tuning](https://mlsanity.com/paper/2610.00436/) (2026)

## Abstract (arXiv)

Modern AI agents such as large language models are trained on diverse tasks -- translation, code generation, mathematical reasoning, and text prediction -- simultaneously. A key question is how to quantify the influence of each individual training task on performance on a target task, a problem we refer to as task attribution. The direct approach, leave-one-out retraining, measures the effect of removing each task, but is computationally infeasible at scale. An alternative approach that builds surrogate models to predict the performance on a target task for any subset of training tasks has emerged in the recent literature. Prior work focuses on linear surrogate models, which capture first-order relationships but miss nonlinear interactions such as XOR-type effects. In this paper, we first consider a unified task-weighting framework for analyzing task-attribution methods and establish a new connection between linear surrogate models and influence functions via a second-order analysis. Then, we introduce kernel surrogate models, which more effectively represent second-order task interactions. To efficiently learn the kernel surrogate, we develop a gradient-based estimation procedure that leverages a first-order approximation of pretrained models; empirically, this yields accurate surrogate estimates with less than 2% relative error without repeated retraining. Experiments across multiple settings -- including mathematical reasoning in transformers, in-context learning, and multi-objective reinforcement learning -- demonstrate the effectiveness of kernel surrogate models. They achieve a 25% higher correlation with the leave-one-out ground truth than linear surrogates and influence-function baselines, enabling more accurate and scalable task attribution. When used for downstream data selection, kernel surrogate models further yield a 40% improvement in the aforementioned settings.