Efficient Estimation of Kernel Surrogate Models for Task Attribution

TL;DR

The paper tackles task attribution in multi-task training by linking leave-one-out and influence-function approaches through a second-order analysis and then overcoming linearity limitations with KernelSM, a kernel surrogate that captures nonlinear task interactions. KernelSM is learned via a gradient-based estimator that uses a first-order approximation around a pretrained model, achieving sub-$2\%$ relative error without retraining. Empirically, KernelSM improves correlation with leave-one-out ground truth by about $25\%$ and enhances downstream task selection performance by roughly $40\%$ across math reasoning, in-context learning, and RL benchmarks, demonstrating scalability to large task sets. The work provides a practical, interpretable, and scalable tool for understanding and leveraging task attribution in modern AI systems, with code released for reproducibility.

Abstract

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 to quantify how each individual training task influences 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 a target task's performance for any subset of training tasks has emerged in recent literature. Prior work focuses on linear surrogate models, which capture first-order relationships, but miss nonlinear interactions such as synergy, antagonism, or XOR-type effects. In this paper, we first consider a unified task weighting framework for analyzing task attribution methods, and show a new connection between linear surrogate models and influence functions through 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 estimates with less than $2\%$ relative error without repeated retraining. Experiments across multiple domains -- including math 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. When used for downstream task selection, kernel surrogate models yield a $40\%$ improvement in demonstration selection for in-context learning and multi-objective reinforcement learning benchmarks.

Figures (6)

• Figure 1: We compare influence functions (IF), leave-one-out (LOO) retraining, and linear surrogate models. Each point corresponds to the individual effect of removing one training example.
• Figure 2: Illustrate linear vs. kernel models on a decision boundary.
• Figure 3: We investigate the surrogate model performance with different kernels.
• Figure 4: Relative approximation error of $\epsilon_{W}(x)$, tested on four different tasks.
• Figure 5: We investigate both linear and kernel surrogate models' fit under different Hessian regularization during model training. The linear surrogate model does not fit model outcomes when using SGD, and only works when its Hessian is regularized. By contrast, kernel surrogate models remain robust when trained with various optimizers and regularization.

Theorems & Definitions (11)

• Definition 2.1
• Proposition 3.1
• Corollary 3.2
• Remark 3.3
• proof
• Corollary A.1
• proof
• proof
• Proposition A.2
• Proposition A.3: Expressivity of RBF kernels