Efficient Estimation of Kernel Surrogate Models for Task Attribution

AI summary

Overview

Research area: Machine learning interpretability, specifically data and task attribution in multi-task training. Technical level: Advanced — the paper combines second-order Taylor expansions of the training objective, influence functions, reproducing kernel Hilbert spaces, and Johnson–Lindenstrauss arguments with large-scale empirical evaluations. Scope: The paper proposes kernel surrogate models (KernelSM) as a scalable substitute for leave-one-out retraining when measuring how individual training tasks influence a target task, and provides both a theoretical analysis of existing linear surrogates and an efficient gradient-based estimator.

What This Paper Is About

Modern AI agents are trained on many tasks at once — translation, code generation, mathematical reasoning, text prediction — and practitioners want to know how much each individual training task contributes to performance on a particular target task, a problem the authors call task attribution. The direct answer, leave-one-out (LOO) retraining, requires K+1 full training runs for K tasks, which is computationally infeasible at scale. The paper's goal is to build a surrogate model that predicts target-task performance for any subset of training tasks, capturing nonlinear interactions such as synergy, antagonism, or XOR-type effects that earlier linear surrogates miss.

Key Contributions

1. A unified task-weighting framework and a theoretical link between linear surrogates and influence functions. Using a second-order Taylor expansion in the task-weight space, the authors show (Proposition 3.1) that linear surrogate coefficients estimated by minimizing the surrogate objective approximately recover the gradient of the performance landscape, up to two bias terms induced by the Bernoulli sampling distribution. Under near-zero second-order interactions, these coefficients approximately equal influence functions (Corollary 3.2).
2. Kernel surrogate models (KernelSM) that replace the linear mapping with a nonlinear function learned via kernel ridge regression in an RKHS, using an RBF kernel over the {0,1}^K subset space to represent second-order task interactions.
3. An efficient gradient-based estimation procedure that exploits a first-order Taylor approximation of pretrained model outputs around initialization W₀, treating (projected) gradients as features. This reduces the surrogate fit to a regularized multinomial logistic regression problem and eliminates repeated retraining; Gaussian random convolutions project gradient vectors into lower dimensions so the regression runs in seconds on CPUs.
4. Broad empirical validation across modular arithmetic reasoning in transformers, in-context learning with Qwen3-8B, and multi-objective reinforcement learning on Meta-World MT10, including downstream task-selection experiments.

Main Findings

• Linear surrogates and influence functions agree when the landscape is near-linear. In a binary classification experiment, the Linear Task Model and Influence Functions achieved Pearson correlations of 0.98 and 0.97, respectively, with leave-one-out scores, and correlated 0.96 with each other.
• Nonlinear interactions are real and matter. A numerical example with a two-layer MLP shows that the influence of a training subset exhibits strong nonlinearity that cannot be decomposed into individual-sample contributions; the RBF kernel surrogate captures these dependencies while the linear surrogate does not.
• RBF kernels have much lower residual error than linear surrogates. With ResNet-9 classifiers on CIFAR-10 and other settings, residual error was 4.4 ± 0.9 (linear) versus 1.0 ± 0.0 (RBF) on CIFAR-10; 4.6 ± 1.3 versus 1.5 ± 0.4 on modular arithmetic; 0.8 ± 0.2 versus 0.4 ± 0.1 on in-context learning; and 0.2 ± 0.1 versus 0.1 ± 0.1 on multi-objective RL.
• The first-order gradient approximation is accurate. Relative approximation error of ϵ_W(x) was 1.02 ± 0.69% on CIFAR-10, 2.40 ± 2.17% on modular arithmetic, 0.51 ± 0.04% on in-context learning, and 0.43 ± 0.73% on multi-objective RL. The abstract reports surrogate estimates with less than 2% relative error without repeated retraining, and the paper states the approximation is under 1% relative error across diverse datasets, on models ranging from small MLP classifiers to language models with up to 34B parameters.
• Better attribution scores (Table 1, LDS, mean ± std over five runs). KernelSM reached 0.30 ± 0.12 (Addition) and 0.52 ± 0.08 (Quadratic) in modular arithmetic; 0.37 ± 0.02 (SST-2) and 0.54 ± 0.01 (Coin flip) in in-context learning; and 0.80 ± 0.04 on Metaworld. The strongest baselines were linear surrogate models at 0.18 ± 0.02, 0.44 ± 0.09, 0.33 ± 0.05, 0.43 ± 0.05, and 0.76 ± 0.04 respectively. Influence functions scored 0.03 ± 0.01, 0.01 ± 0.01, 0.16 ± 0.05, 0.05 ± 0.10, and 0.71 ± 0.11.
• Sizeable gains over prior methods. KernelSM improves LDS by over 42% on average versus linear surrogate models in modular arithmetic reasoning, achieves a 25% relative improvement in LDS on the Coin Flip ICL task, reports an 18% improvement in attribution accuracy over prior methods for Qwen3-8B on sentiment classification and mathematical reasoning, and a 5% improvement in Meta-World MT10 environments with dynamically shifting data distributions. Overall, the abstract and introduction report 25% higher correlation with leave-one-out ground truth relative to linear surrogates and influence-function baselines.
• Scaling does not degrade performance. Using 500 prompt samples to simulate a task-scalable setting, LDS values remained comparable to results on 50 tasks: 0.31 on SST-2 and 0.53 on Coin-Flip.
• Downstream task selection improves. Selecting the top-4 examples by attribution score for ICL prompt selection and co-training each RL target task with the top-3 attributed tasks, KernelSM achieved a 40% lower loss in the prompt selection task in ICL and improved target rewards by 15%. The abstract and introduction describe a 40% improvement in demonstration selection, while Section 4 states KernelSM improves the performance of baselines by 41%. The paper also reports that runtime is comparable to fitting linear surrogate models; the truncated Table 2 reports timings in minutes (for example 17 ± 1 for influence functions, 3 ± 1 for TracIn, and 2 ± 1 for Trak), though the row labeling in the available content is ambiguous.

Methodology in Plain English

The authors treat the training pipeline as a black box that maps a binary vector s — indicating which of K tasks are included — to a target-task performance score F(s). Instead of retraining the model once per subset, they sample m subsets and learn a surrogate function that predicts F(s) for any subset.

Their first step is analytical: they expand F(s) around the all-tasks point using a second-order Taylor expansion and plug it into the linear surrogate objective, showing that the fitted linear coefficients resemble influence functions when second-order interactions are negligible. Since real multi-task training has such interactions, they then replace the linear surrogate with a kernel ridge regression model in a reproducing kernel Hilbert space, using an RBF kernel k(s^(a), s^(b)) = exp(−γ‖s^(a) − s^(b)‖²) whose geometry matches the binary subset space. Hyperparameters λ and γ are chosen by cross-validation.

To avoid training m separate models to obtain each F(s^(i)), they approximate model outputs with a first-order Taylor expansion around a pretrained initialization W₀, f_W(x) ≈ f_{W₀}(x) + ⟨∇f_{W₀}(x), W − W₀⟩. With this approximation, estimating the best parameter perturbation on each subset becomes a regularized multinomial logistic regression problem with projected gradients as features. Because logits and gradients at W₀ need to be computed only once for all training samples, the entire procedure can run on CPUs and scale to large m. Gradient vectors are compressed with Gaussian random convolutions so the regression takes seconds, and the Johnson–Lindenstrauss lemma is used to bound the accuracy of this estimation relative to true test performance (Proposition A.4). The full pipeline is given as Algorithm 1 (KernelSM) with the estimation subroutine as Algorithm 2 (GradEx).

Evaluation uses the linear datamodeling score (LDS), computed as the Spearman correlation between real model outputs on sampled subsets and the attribution method's estimates, and compares against influence functions (with LISSA), TracIn, TRAK, SOURCE, Bayesian influence functions (BIF), and linear surrogate models.

Why This Matters

Impact on research. The paper supplies a formal bridge between two previously separate families of attribution methods — linear surrogate modeling (datamodels-style approaches) and influence functions — and shows exactly when they coincide. It then argues, both theoretically and empirically, that second-order interactions require a kernel-based surrogate, giving the field a principled reason to move beyond linear attribution and a concrete estimator to do so without repeated retraining.

Real-world applications.

• Multi-task learning: guiding architecture design and loss-reweighting strategies based on how tasks interact.
• Multi-group learning: revealing how training on different demographic groups shapes model behavior.
• In-context learning: determining how adding or removing a single demonstration affects predictions, which directly supports demonstration or prompt selection.
• Multi-objective reinforcement learning: understanding how competing reward signals influence the learned policy, and deciding which tasks to co-train for positive transfer.

Industry relevance. Practitioners training large models on heterogeneous data mixes need to know which data sources help or hurt a target capability, but cannot afford K+1 retraining runs. A surrogate that can be fit with gradients computed once, largely on CPUs, and yields attributions usable for data or demonstration selection, is directly applicable to data-mix curation and multi-task training decisions.

Future Directions

• Extending the second-order analysis to higher-order task interactions, since Corollary 3.2 only covers the regime where second-order interactions are near zero and the paper motivates "more expressive attribution methods" beyond it.
• Strengthening the performance guarantees for the gradient-based estimator beyond the stated Proposition A.2 and the Johnson–Lindenstrauss-based bound in Proposition A.4, which rely on the assumption that the first-order approximation error is small.
• Testing the approach at larger task counts and on a wider range of large language models; the paper's scaling check uses 500 prompt samples and reports approximation evaluation on models up to 34B parameters, leaving broader scaling behavior open.
• Improving estimation of individual item attributions from the kernel surrogate — the current approach averages predictions over sampled subsets that include an item, which leaves room for more direct or more sample-efficient per-item attribution.

Target Audience

Researchers and advanced practitioners in machine learning interpretability, data attribution, and multi-task learning who need efficient tools for understanding task composition effects. It is also relevant to engineers curating data mixtures or selecting in-context demonstrations for large language models, and to reinforcement learning practitioners working with multi-objective or multi-task benchmarks. Readers without a background in influence functions, kernel methods, or Taylor-expansion-based analysis will find the theoretical sections demanding.

Authors’ 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 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.