Efficient Estimation of Kernel Surrogate Models for Task Attribution
Authors: Zhenshuo Zhang
Minxuan Duan
Hongyang R. Zhang
Venue: ICLR 2026 Poster (2026)
Keywords: Model interpretability
Data attribution
Kernels methods
Primary area: interpretability and explainable AI

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 surrogate estimates with less than 2% relative error without repeated retraining. Experiments across multiple domains---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 task selection, kernel surrogate models further yield a 40% improvement in demonstration selection for in-context learning and multi-objective reinforcement learning benchmarks.

Published as a conference paper at ICLR 2026
EFFICIENT ESTIMATION OF KERNEL SURROGATE MODELS FOR TASK ATTRIBUTION
Zhenshuo Zhang" Northeastern University
Minxuan Duan" Northeastern University
Hongyang R. Zhang" Northeastern University
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
a target task, a problem we refer to as task aftribution. 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 's performance for any subset training tasks has emerged in recent literature. Prior work focuses linear sur- rogate models, which capture first-order relationships, but miss nonlinear interac- tions such synergy, antagonism, XOR-type effects. In this paper, first consider a unified task weighting framework for analyzing task attribution meth- ods, and show 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 effi- ciently learn the kemeel surrogate, we develop a gradient- based estimation proce- dure that leverages a first-order approximation of pretrained models; empirically, this yields accurate surrogate estimates with less than 2% relative error without re- mathematical
peated retraining. Experiments across multiple reasoning in transformers, in-context learning, and multi-objective reinforcement the effectiveness kernel surrogate models. They achieve 25% higher correlation with the leave-one-out ground truth than linear surrogates and refluenceffunction baselines, enabling more accurate and scalable task attri- bution. When used for downstream task selection, kernel surrogate models further yield a 40% improvement in demonstration selection for in-context learning and multi-objective reinforcement learning benchmarks.
INTRODUCTION
Modern AI systems are trained to perform well across a diverse set of tasks, ranging from ma- chine translation and code generation to object recognition and mathematical reasoning, This broad, multi-task capability raises a central question in interpretability: how do individual training tasks contribute model performance?
In this work, we study task attribution, the problem of quantifying the influence of each training task on downstream behavior. Beyond being conceptually fundamental, task attribution has direct impli- cations across several application domains. In multi-task learning, understanding task relationships can guide the design neural architectures and reseighwegg strategies (Wu al., 2020: Yang et al., 2025). In multi-group learning (Deng & Hsu, 2024). it can reveal how training on different demographic groups shapes model behavior. In in-context learning (Garg et al., 2022: Zhang al., 2025). it parallels the question of how adding or removing single demonstration affects predic- tions Min et al. (2022). Related questions also arise et al., 2020; Zhang et al., 2026), where one wants to understand how competing reward signals in- fluence learned policy. Accurate task attribution can also help understand when the representations generated by pre- trained large-scal models can be used in downstream tasks (Wu et al., 2023).
objeetvb reinforcement learning (Yu
"Email Address: (zhang.zheml duan.mi, mi, ho.zhang) @northeastem.edu.

---

Published as a conference paper at ICLR 2026
A natural estimator of task influence is leave-one-out (LOO) retraining, which measures the change in performance when one task is withheld from training. However, if the training set contains K tasks, computiing LOO scores requires K + 1 full training runs, which is prohibitive. More efficient proxies for the LOO scores use influence functions (Koh & Liang, 2017; Grosse et al., 2023), which estimate how model predictions would change if training sample were added or removed. These methods require evaluating Hessian-vector products, and are usually applied at the terminal state once the model has memorized the training data (Brown et al., 2021). However, computing influ- ence scores for all K' tasks still requires repeatedly Hessian approximations (Basu et al., 2021; Kwon et al., 2024). A complementary line of work datamodels (Ilyas al., 2022) and data attribution (Park et al., 2023; Bae et al., 2024)-buille surrogate functions that approximate black-box model behaviors. These methods treat the training pipeline as a function mapping data (or task) subsets to predictions, and learn a separate surrogate to emulate the mapping. Once trained, the surrogate provides a faster mechanism for evaluating how subset combinations contribute to a given test dis- tribution. Recent work has primarily focused on linear surrogate models, which enjoy appealing theoretical properties, as they can capture first-order effects and admit linear sampling complexities in multi-task settings (Li al., 2023; 2024a; Zhang et al., 2025). However, linear surrogates fail to capture second-order interactions where task effects are non-additive (See illustration in Figure 2).
In this paper, we revisit the surrogate modeling objective through a unified task weishheigt frame- work. We begin by analyzing linear surrogate models using a second-order Taylor expansion in the weight space. We find that linear surrogate models estimated via minimizing the surrogate model- ing objective are approximately equal to the influence functions, up to third-order expansion errors (Proposition 3.1). In particular, when the second-term expansion error is small, then linear sur- rogates and influence functions are approximately the same (Corollary 3.2). A key technical tool involves applying the delta method to certain covariance statistics in the surrogate regression.
Motivated by these observations, we propose kernel surrogate models (KERNELSM), which extend surrogate modeling beyond the additive structure inherent in linear methods. Kernel surrogates- such as those based on radial basis function (RBF) sernels-naturally capture geometric relation- ships in the 0-1 subset space. A direct kernel estimation, however, would require computing model outcomes on many task subsets. Instead, we design an efficient estimation that uses gradients as fea- tures, leveraging a first-order approximation of model outputs in its parameters. Empirically, we find that the accuracy of this first-order approximation is under 1 5 relative error across diverse datasets, including CIFAR-10, modular arithmetic, in in-context learning, and multi-objective RL benchmarks, on models ranging from small MLP classifiers to language models with up to 34B parameters. The- oretically, we bound the error of this estimation assuming the errors are small (Proposition A.2).
We evaluate the kernel surrogate modeling framework across broad range of datasets and archi- tectures. On modular arithmetic reasoning where the input-output mapping can be XOR, division, quadratic-kernel surrogates improve attribution accuracy by up 42% relative to exist- ing approaches (Li et al., 2023; llyas et al., 2022; Park et al., 2023). We also observe similar gains on in-context learning and sequential decision-male making tasks. Applied to the Qwen3-8B model on sentiment classification and mathematical reasoning, kernel surrogates improve attribution accuracy by 18% over prior methods.
reinforcement learning with soft actor-critic on the Meta- World MT10 benchmark (Yu et al., 2020), our method provides 5% improvement on environments with dynamically shifting data distributions. Overall, kernel surrogates consistently outperform prior methods across all settings, increasing correlation with leave-one-out estimates by 25% relative to existing attribution methods. Finally, for downstream task selection, our approach yields 40% lower loss in both LLM inference (demonstration selection for in-context learning) and multi-objective optimization on Meta-World, World, while using runtime comparable to fitting linear surrogate models.
Taken together, this paper provides an efficient estimation algorithm for the lask-attribution prob- lem via kernel surrogate models. First, we analyze linear surrogate models, showing that influence functions and the estimated regression coefficients are approximately the same under certain as- sumptions. Second, we introduce KERNELSM, along with an efficient gradient-based estimation via kernel ridge regression. Thus, our approach is both scalable and flexible for modeling nonlinear task relationships in black-box systems. Third, we evaluate KERNEL in a variety of empiri- cal settings including in-context learning, mathematical reasoning, and reinforcement learning. We provide the code to replicate our empirical results at http:///juubbcoomiiimoooeeeees murrogate-moeels.
2

---

Published as a conference paper at ICLR 2026
2 PRELIMINARIES
We begin by introducing the notation and a task weighting framework. Within this framework, provide a formal definition of task attribution. We then present influence functions and task modeling techniques from the data attribution literature within this unified framework, which also serve as state-of-the- baselines for evaluating task attribution.
Problem setup. Suppose the training dataset consists of K tasks, denoted by T1. T2..... . Each task contains nı samples drawn from its task specific distribution, i.e., Tk = {[Tk.j) yk.j)) }=
for k 1, : K. The total number of training samples is equal to n
=k==1 ===== 14--
We consider model fw such multitask network parameterized by W IRd, which is shared across all tasks. For each task k, we denote the task-level loss by & Ts). To specify which tasks are included in training, where 8k 1 indicates that task & is selected in and $ę 0 otherwise. The corresponding weighted empirical loss is then defined as
introduce K-dimensional binary vector [s1,-.,
K
K Σ
i(fw,s)
Σ
$j
sulu(fw, TT).
For example, when 88 = for all k; the weighted loss L(Jw, 8) reduces to the average loss over all tasks. We denote by W(s) arg minw L(fw,s) the minimizer of the weighted loss, and by (W))) the corresponding model.
Influence estimation. We aim to evaluate the influence of training tasks on a target test task Tisst To this end, we define performance metric F(s) = E (w Tast) which measures the test loss of the model trained with subset of training tasks indicated by
natural way to quantify the contribution of a task through leave-one-out (LOO) retraining, which measures the change in performance when task k is excluded from the training set
It:00 F( F(11]x - -e),,
where ę E RK is the all-ones vector corresponding to training on all tasks, and E {0, is the basis vector corresponding to the k-th coordinate.
Definition 2.1. The ground-truce truth task attribution score of task is given by I for k 1...., K. K.
The other approach for capturing task influence is by computing influence functions the weighted loss. Given s, the influence of task on F(s) is formally given by (Koh & Liang, 2017)
86 W (Ww,Ta
(s) = V W [(S) Viw
(2)
j==
In words, (s) quantifies how changes when task k is infinitesimally up-weighted. The gradient captures task k's direct contribution, and the Hessian inverse accounts for curvature of the loss. The outer product with Ww Ps)) maps the perturbation its effect L. We provide the derivation of equation (2) in Appendix A.
Linear surrogate models. An alternative approach to quantify task influence is through task mod- eling (Li et al., 2023). The key idea is to sample binary vectors s(1) s(2),...,, , from {0, 1 each indicating a random subset of tasks. For each s()), we train the model the corresponding weighted loss to obtain (s) and record its performance F(s()).. We then fit linear surrogate model, parameterized by intercept a E IR and coefficients
[B....., .x.) RRK by minimizing a- - 33)
R(o, EgwD (F(s) a
where D is distribution of such as the Binomial distribution. The optimal linear surrogate model is obtained by minimizing the empirical loss R( () on m sampled pairs ((s(), F(s())))
3

---

Published as a conference paper at ICLR 2026
3 METHODS
In this section, we present kernel-based surrogate modeling approach for influence estimation. First, we analyze surrogate modeling together with influence functions using a second-order Tay- lor expansion in the task weighting space. We develop second-order approximation of the linear surrogate model coefficients and find that it approximately equal the influence function. Specif- ically, linear surrogate models provide an estimation of the influence functions if second-ord order task interactions are nearly zero in the composition. Second, we design kemel-based surrogate mod- eling procedure to learn the nonlinear interactions between tasks. Finally, we introduce efficient gradient estimation algorithm efficiently learn the kernel surrogate.
3.1 ANALYZING SURROGATE MODELING IN THE ATTRIBUTION SPACE
k- Second-order analysis of linear surrogate modeling. We center analysis around s* which corresponds to the uniform weight applied to all tasks. Then, we characterize task contribu- tions with respect this global model. Expanding around the anchor point s*, we approximate the performance function F(s) using a second-order Taylor expansion as follows:
F(s) = F(s*)+ + [FF(s')] sss
s*) (s
(4)
Finst-order linear effects
Second-order interaction eflictts
where the Hessian matrix Hg (2*)(*) captures the nonlinear interactions between tasks.
We next connect this expansion to linear task modeling, which approximates F(s) with linear sur- rogates of the form s.. To analyze the behavior of linear surrogates, we apply a second-order expansion of F(s) in equation (4) into the surrogate modeling objective, R(a, 3). The following
proposition formalizes the resulting characterization of the linear regression coefficients from mini- mizing R(a, 33 to get the regression coefficients â; 3.
Proposition 3.1. Let F 1) + R and s* []]K- Assume the third (partial) derivatives of F()) are bounded by C3 for some small enough C3 > Suppose each coordinate in S is drawn independently from a Bernoulli distribution with probability for some fixed E (0, 1). Let & be the minimizer of the linear surrogate objective R(a, 3) on m random samples ((2), , s(m)) Then, with probability 8 over the randomness of m sampled subsets, we have
| B Ve-Va(ss)-aa -0 2 [Ha] cakkp- +
K + log(d-1)
(5)
This result shows that the surrogate coefficients 3 recover the gradient of F(s*) with respect to the weights up to two terms induced by the Binomial distribution. The first corresponds to a sampling shift due to the mean of the Bernoulli draw, while the second reflects the variance of task inclusion the random subsets. The key idea is based on the second-order delta method applied to certain certain covariance stutistics. A detailed proof of this proposition is in Appendix A.
Connection between linear surrogate models and influence functions. This result also shows the connection between the linear task modeling and influence functions. Specifically, under a first-order approximation (||H// S C2), both methods estimate the same underlying quantity (sF(s*). Un- der this assumption, the second-order bias terms in equation (5) vanish, showing that the linear task
modeling coefficients approximate the gradient of the performance landscape: 3 F(F(st). Mean- while, one can show that the influence function computes 7 (25"), (xxx"")).,III((( VzF(s*). We summarize this discussion precisely follows.
Ix(s*)]
Corollary 3.2. In the setting of Proposition 3.1, assume further that the second-order task interac- tions are near zero, i.e., the spectral norm of |H ll I C2 for some small 02 > O, the linear surrogate model coefficients satisfy that with probability at least 1 over the randomness of s(1) ,..., (m)
for any 6 > 0:
+ +1/2) | - S +0 (1cx33/20))))) +0 +0
( +lg(000
log(
(6)
4

---

Published as a conference paper at ICLR 2026
Correlation = 0.98
Correlation 0.97
Correlation 0.96
0.04
0.04
0.04
LOO
0.02
LO0
0.02
0.02
0.00
0.00
0.00
-0.02
0.02
0.02
4
2
0
2
4
-
0
2
B
x10
X10
X
(a) LOO / Linear Surrogate
(b) LOO / IF
(c) IF Linear Surrogate
We compare influence functions (IF), leave-one-our (LOO) retraining, and linear surrogate models. Each point corresponds to the individual effect of removing one training example.
Figure 1:
Thus, provided that 02, C3 are small enough and m 2 O(K + log(6-1)), then influence functions approximate linear surrogate models, while linear surrogate models capture additive relations. In particular, more expressive attribution methods are needed to express second-order task interactions. We provide the proof of this result in Appendix A.2.
Next, we experiment with a binary classification to empirically validate the connection, where the first-order approximation holds. We compare the estimates from both the linear task model and influence functions against the ground-truth leave-one-out (LOO) retraining scores. In Figure 1, we find that both methods produced estimates that aligned closely with the ground truth: the Linear Task Model and Influence Functions achieved high Pearson correlations of 0.98 and 0.97, respectively, with the LOO scores. The estimates from the two methods are also highly correlated with each other, with a Pearson correlation of 0.96. This demonstrates that when the performance landscape is nearly linear, both the empirical and analytical approaches indeed converge to the same underlying attribution scores. A detailed description of the experimental setup is available in Appendix B.1.
3.2 LEARNING KERNEL. SURROGATE MODELS
To build intuition for why nonlinear data model is necessary, we present numerical example in Figure 2. We use binary classification task with a two-layer MLP as the base classifier. The final goal for the surrogate model is to predict the MLP output training data it was trained on. We specifically analyze the effect of different subsets of training data sampled from near the MLP decision boundary. The detailed setting is in Appendix B.I.
fixed test sample, given the subset of
The influence of a training subset exhibits strong non- linearity that cannot be decomposed into individual sample contributions. When specific samples combine in training subset, they produce changes in the MIP learned function, causing transitions in test point predictions. The linear sur- rogate model does not capture these interactions. By con- trast, the RBF kernel surrogate model captures these depen- dencies by modeling the joint influence sample combi- nations. This demonstrates that non-linear surrogate mod- els are necessary for capturing nonlinear interactions.
1.0
Prediction
Goal
0.5
Lincar
Kernel
0.0
Subset
Figure 2: Illustrate linear vs. kernel
models on a decision boundary.
Estimating surrogate models using kernel ridge regression. To address the limitation of linear surrogates, we propose kernel surrogate models, replacing linear mapping with a nonlinear function learned via kernel ridge regression. This enables modeling complex, non-additive task interactions.
We learn kernel surrogate function 1) - within a Reproducing Kernel Hilbert Space (RKHS) Hj by minimizing the following objective:
min H
Σ (F(s(0) - (0() + Allgol1
(7)
5

---

Published as a conference paper at ICLR 2026
Table 1: We investigate the surrogate model per- formance with different kernels.
Table Relative approximation error of ew (2), tested on four different tasks.
Residual error CIFAR-10
Linear
RBF
Relative error
4.4
1.0
CIFAR-I0
1.02 1...9.699... %
Modular arithmetic
4.641.3
1.540.4
Modular arithmetic
2.40 %
In-context learning
0.8+0.2
0.440.1
In-context learning
0.51 %
Multi-objective RL
0.240.1
0.140.1
Multi-objective RL
0.43+o.73%
where À > 0 is regularization parameter and K. is an m x m kernel matrix with entries K,j k(s(t), ,(()). Since input vectors are binary, we use the Radial Basis Function (RBF) kernel:
k(s(() (0) s ())) By the representer theorem (Schölkopf al., 2001), the minimizer takes the form:
exp ( y ||ss - 11))
(8)
m
go (s)
Σ ())(s()ss))
The coefficient vector 0 [01,.... mm]] is computed 0 (K AId) where F is the vector of observed outcomes and Id denotes the identity matrix.
This kernel provides flexibility for capturing complex relationships while maintaining computational tractability. Hyperparameters À and y are selected via cross- validation. Unlike linear models that assign fixed coefficients to individual tasks, the kernel task model learns global non-linear function that predicts performance for any task combination, thereby capturing intricate inter-dependencies.
We empirically validated this choice by comparing the RBF kernel with linear and polynomial ker- nels with degrees {1, 2} in Table 1. We use ResNet-9 network classifiers trained on the CIFAR-10 dataset (Krizhevsky & Hinton, 2009) as the base model. The detailed description is in Appendix B.2. We find that the RBF kernel achieves much lower residual error compared to the linear surrogate model. We discuss more kernels in Appendix B.2.
Remark 3.3. A major property of RBF kernels (Poggio & Girosi, 2002) is its universal expressivity over any distance functions the space {0,1) 1)K Moreover, the RBF kernel encodes an inductive bias that matches the geometry of the subset space {0, 1) (See also, the use of heat kernels for community detection (Kloster & Gleich, 2014)). We formalize this in Proposition A.3, Appendix A.5.
3.3 EFFICIENT ESTIMATION VIA PROJECTED GRADIENTS
One downside of kernel surrogate models is that they involve training m separate models to obtain F(s()), one on each g(f) Rather than pursuing this repeated training, we introduce an efficient algorithm to approximate the model outcomes instead. Consider first-order Taylor expansion of model outputs around the initialization Wa as follows:
fw(x) = fw.(x)) + fW (x), - WW) + ew(x),
where (x)) is the first-order Taylor expansion error of fw, on input The intuition around this estimation algorithm can be related to how neural networks learn features through the gradients (Malladi et al., 2023; Radhakrishnan al., 2024). Empirically, fw( in Table 2, which is with 2% across various test environments, and report further evaluation of this approximation on LLMs with up 34B parameters in Table 5, Appendix A.4..
also evaluate ew (ar) relative to
Next, we substitute approximation (9) of fw (x) into the empirical loss to estimate the optimal fine-tuning parameter perturbation on top of the initialization Wo for every s(t) Specifically, we formulate an approximate objective on each subset s(t) as:
Z arg min
Σ (T.
la (f),()) Ww (x),
K
(4)
ZERa
assuming the expansion error é (z) is negligibly small. This corresponds to (regularized) multi- nomial logistic regression problem, with (projected) gradients serving features.

---

Published as a conference paper at ICLR 2026
Algorithm Estimating Kernel Surrogate Models (KERNELSM)
Input: K training tasks "i,..., Tk. pretrained model fwo, test dataset and metric F())
Subset sampling distribution D, number of subsets m, Ez-regularization parameter X, kernel function k(:,))
Requires:
Output: An estimated kernel surrogate model (KSM)
/*
Generate random subsets and estimate test performance
1: D. + 0: Initialize surrogate dataset
2:
(fw, (2), fww t- Compute logits and gradients for all relevant samples
3:
for
I,..., do s(1) D
4:
5:
F(s(0) - GradEx(s()) (Swo(z) (swa(=)), Apply Algorithm 2
6: Do DDD ((so), F(s())))
7: end for
Learn kernel surrogate model using kernel ridge regression 8: K, + -k(s(*) ,())) Construct the m x m kernel matrix K.
1
9: y+ [F(((1), ..... [T Compute the outcome vector
10: 0 (K + A x-dddm)) y: KRR coefficients
11: return go(s) = Ei= (g1(((0)
Based on this estimation algorithm, if we want to minimize the approximate loss all the subsets ((1), s(2) ...., s(m)), we only need to compute the functional values and the gradients at the initializa- tion Wo (such as a pretrained LLM) once, including fw, (x) and fw. for all x in the training dataset. This eliminates the need for repeated training and dramatically speeds up the estimation procedure, and is practical even for very large m since the estimation procedure can be executed entirely CPUs.
Finally, we estimate the model output on a test sample x, y using the same linear approximation:
f(x) = fw. + fW. (x). Zio)).
The test performance of each ( is then evaluated as the loss ((f(z), y) on the test dataset. Another challenge involved in this procedure is that the gradient vectors can have large number of features. We address this problem by applying Gaussian random convolutions to project the gradient vectors into much lower dimensions. Then, solving the regression problem in the reduced space only takes several seconds, which can be run quickly on many subsets m. Due to space limit, we summarize the detailed procedure in Algorithm 2, Appendix A.3. Taken together, summarize the overall procedure in Algorithm 1.
Performance guarantees. Assuming that the first-order approximation is small, bound the accuracy of the above estimation procedure relative to the true test performance using Johnson- Lindenstrauss lemma. For details, see the result statement and its proof following Proposition A.4.
4 EXPERIMENTS
Our experiments investigate three key questions: (1) How does KERNELSM compare to existing attribution techniques on complex tasks such as modular arithmetic reasoning, in-context learning, and multi-objective reinforcement learning? (2) Can the attributions generated by KERNFLSM be effectively leveraged for downstream task selection?
(3) How does complex loss landscape impact the performance of KERNELSM compared to methods relying on linear surrogate models? We evaluate our approach across three settings, including arithmetic reasoning, in-context learning, and multi-objective reinforcement learning. We find that KERNELSM can improve over existing influence estimation methods by 25%. When applied to downstream task selection, KERNELSM can improve the performance of baselines by 41%. Various robustness checks and ablation studies validate the consistency of KERNELSM under different optimizers and sample sizes.
7

---

Published as a conference paper at ICLR 2026
4.1 EXPERIMENT SETUP
Datasets and models. Our primary controlled environment is arithmetic reasoning task, where we train a small decoder-only transformer. Consider learning a transformer model to perform mod- ular arithmetic operations over two numbers, in the form of o b (mod p) = c, where o is an arithmetic operation and p is prime. Inputs are composed of four tokens: a o, b, and =, with a and b generated between and p 1. The label is the result of the arithmetic operation c. Specifically, we use the addition task a + b (mod p) = and the quadratic task a² = b² + ab (mod p) = c, where we set p = 97.
For more complex scenarios, we evaluate a Qwen3-8B model on in-context classification tasks. We construct prompts with in-context examples followed by a query. For attribution computation, we treat each example's input embedding as the sample input, and the query's s cross-entrops loss as the model score Fss). We evaluate on two datasets, including SST-2 and coin flip.
We also evaluate attribution in multi-objective reinforcement learning using the Weta-World MT10 benchmark. It contains ten different manipulation tasks. We adopt the Soft Actor-Critic (SAC) algorithm as our training protocol. Each task is treated as sample for attribution analysis, and we evaluate the model score by the average rewards of tasks.
Baselines. We compare KERNELSM with existing data attribution methods. Influence func- tions (Koh & Liang, 2017) employ implicit differentiation, computing the inverse Hessian via LISSA. TracIn (Pruthi et al., 2020) traces prediction changes throughout training by computing gradient products between training and test examples. Linear surrogate models (Ilyas et al., 2022) train linear models that predict test behavior from binary indicators of training sample inclusion across multiple model retrainings. Trak (Park et al., 2023) applies random projection and the one- step Newton method to approximate the datamodel results without retraining. SOURCE (Bae al., 2024) approximates unrolled differentiation by segmenting training into a few stationary phases and using gradient/Hessian statistics from a handful of checkpoints to estimate counterfactual parameter changes. Bayesian influence functions (BIF) (Kreer et al., 2025) extend classical influence functions by expressing influence as covariance under a localized Bayesian posterior and estimating via batched SGLD sampling.
Evaluation. We evaluate the data attribution methods by comparing them to the linear datamodeling score (LDS) (Ilyas et al., 2022). We compute LDS in the following steps: (1) Sample m fixed subsets of the training set, which are represented by s (mm) E {0, 1}K We sample s from a Bernoulli distribution with probability (2) For each s(O), we train a model (()())), and use a task attribution method to obtain an estimation (s)())). (3) We compute the LDS from the Spearman correlation p between the real model output and the task attribution estimation.
4.2 EXPERIMENTAL RESULTS
Attribution results. Our results in Table demonstrate that by moving beyond the linear assump- tions of prior work, KERNEL.SM achieves more accurate estimation of task attribution. The ad- vantage is most pronounced in tasks with complex, non-linear structures. For example, in modular arithmetic reasoning. KERNELSM improves the LDS by over 42% on average compared to linear surrogate models. Similarly, on the reasoning Coin Flip ICL task, where the language model ex- hibits intricate prompt-followigg behavior that falters linear approximations, our method's ability to capture higher-order interactions yields a 25% relative improvement in LDS. In multi-object rein- forcement learning, where the task's 's more transparent structure allows the linear baseline to perform well already (LDS of 0.76), KERNEL still shows an improvement to 0.80.
We further evaluate whether KERNEL SM scales to scenarios involving a large number of tasks. In our in-context lcarning experiments, we use 500 prompt samples to simulate a task-scalable setting. The resulting LDS values remain comparable to the initial results tested on 50 tasks: 0.31 on SST-2 and 0.53 on Coin-Flip These results indicate that the performance of KERNELSM does not degrade as K increases, demonstrating its robustness in scaling to settings with large K.
Task selection results. Next, we show that the accurate attributions from KERNELSM are applica- ble to downstream optimization tasks. We apply task attribution methods to find optimal groupings for downstream performance. For prompt selection in ICI., we select the top-4 examples with the
8

---

Published as a conference paper at ICLR 2026
Table 3: Summary of comparison results on task attribution. deviation from five independent runs.
We report the mean and standard
Modular arithmetic reasoning
In-context learning
Multitask RL
Methods
Addition
Quadratic
SST-2
Coin flip
Metaworld
Influence functions
0.03
0.01=0.01
0.16+0.05
0,05++.10
0.71
TracIn
1.101
0.32
0.21 10.05
0.32
TRAK
0.14
0.28=0.06
00.08
0.23
0.42
SOURCE
0. 0.22++0.5
0.28
0.27
BIF
0.18 16.16
0.42.0.07
1..020.00
10.08
0.22 1..24
Linear surrogate models
0.18+0..02
0.44
0.33+0.05
0.43
0.76
KERNELSM
1..11
0.52
0.37 10.02
0.54 10.01
0.80 0.80.0.04
Table 4: We evaluate KERNELSM on the downstream task selection. We measure performance by the loss ICL tasks and the target rewards in multi-objective RL tasks. We measure the running time in minutes.
Loss (1)
SST-2
Coin flip
Time (L)
Rewards (t)
MT 10
Time (L)
Influence func.
0.23
1.6241.33
17+1
Influence func.
16.3+1.4
2+1
TracIn
0.25+40.00
0.5440.08
3+1
TracIn
16.042.3
3+1
Trak
0.32+0.48
0. 0.7240.87
2+1
Trak
15.844.5
1+1
Lin. surrogate
0.20
0.05
2+1
Lin. surrogate
13.1+4.3
1+1
KERNELSM
0.16400.8
0.02
2+1
KERNELSM
18.845.6
1+1
highest attribution score. For synergistic task selection in RL, we co-train each target task with the top-3 tasks identified by our method to encourage positive transfer.
To derive individual attribution for each item from the kernel surrogate model, we adopt a random ensemble method. We sample multiple subsets and use our kernel surrogate to obtain performance prediction for each. The final attribution for any given item is then calculated as the average predic- tion score of all sampled subsets in which that item was included.
As shown Table 4, KERNEL achieves a 40% lower loss in the prompt selection task in ICL, and improves the target rewards by 15%. These results demonstrate the benefit of using KERNELSM over linear surrogate models.
4.3 REGULARIZATION OF HESSIAN
Recall from equation (5) that the error of surrogate models is influenced by the second-order terms. Next, we compare the accuracy of linear task modeling by measuring the curvature of the loss Hessian LWW), by controlling it with different Hessian regularization. We report results on modular quadratic task trained using a two-layer transformer. We compared three methods, cluding the standard SGD without explicit Hessian regularization, and two Hessian regularization methods: sharpness-aware minimization (SAM) (Foret al., 2021), ), and noise stability optimization (NSO) (Zhang et al., 2024). We measure the Hessian trace on the global model, and the LDS and residual error of the linear surrogate model are measured on subsets with a 0.9.
In SGD, the Hessian trace increases continuously throughout the training process. This worsens the linear surrogate model's performance, characterized by a decrease LDS. Eventually, the LDS dropped to 0.01. By contrast, both SAM and NSO constrain the Hessian trace. This regularization results in improved model fit, whose LDS remained relatively consistent throughout training.
5 RELATED WORK
Our work is related to the literature on training data attribution methods. Data attribution methods have been extensively studied to quantify the influence of individual training samples model behavior (Koh & Liang, 2017; Brown et al., 2021; Ilyas et al., 2022). In this work, we extend these concepts by measuring the contribution of groups of samples that form different training tasks.

---

Published as a conference paper at ICLR 2026
10
0.4
0.4
Hessian
LD5
LDS
10%
0.2
0.2
SGD
SAM
10
N30
0.0
0.0
0.5
1.0
1.5
2.0
0.5
1.0
2.0 Epochs x10*
1.5
0.5
1.0
1.5
2.0
Epochs x10
Epochs X x10*
(a) Hessian trace
(b) Linear model
(c) RBF Kernel
Figure 3: We investigate both linear and kernel surrogate models' fit under different Hessian reg- ularization during model training. The linear surrogate model does not 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.
The first line of approach is based on the influence function (Koh & Liang. 2017). These methods compute the influence function. The main challenge is that the influence function involves Hessian computation. Instead of computing the exact Hessian inverse, recent approaches use the Gauss- Newton Hessian approximation (Choe al., 2024) sparse gradient compression (Hu al., 2025). Influence functions can also be applied to non-decomposable loss functions, such as contrastive loss and preference loss, allowing for data attribution across a broader class of models (Deng al., 2025). The second line of approach studies the trajectory of model training, rather than just examining the final state of the trained model. These methods aim trace the training dynamics and predict the counterfactual trajectory that would result from different training set (Bae et al., 2024). Recent work calculates the exact influence of training data on single, deterministic model instance by leveraging large-scale meta-g aata-grrdient computation (Engstrom et al., 2025; Ilyas & Engstrom, 2025). The third line research involves using a surrogate model to approximate the behavior of the original complex model and then computing the surrogate model. Recent work shows that the model output can be linearized in a certain local area (Li et al., 2024b). Under this assumption, one first-order Taylor expansion
rephrase the model output provides an efficient estimation for the leave-one-out score (Li et al., 2023) and for datamodels (Park al., 2023). Building on surrogate models, our work non-linear interactions between tasks. The use of kernel ridge regression, also known as Tikhonov regularization, can be related earlier literature on kernel methods (Dicker et al., 2017), and it may be worth further examine the role of regularization in task attribution. The RBF kernel that we have used is also related to the heat kernel, which is highly effective for seed set expansion in community detection (Kloster & Gleich, 2014).
the parameters. This approach
kernel-based model better capture the
A limitation of this work is that our evaluation primarily focuses on prediction tasks. Extending our attribution framework to examine other aspects of language models such planning abilities remains an open question, which a fundamental aspect of human intelligence (Wang al., 2024). In particular, this approach may explain the generalization gaps between SFT and RL (Wang al., 2025); we leave a detailed exploration of this connection to future work.
6 CONCLUSION
In this paper, we study the problem of task attribution, which aims to quantify the influence of dif- ferent training tasks on model's final performance. We first unify the theories of existing influence functions and linear surrogate models, highlighting their limitations in handling non-linear interac- tions between tasks. To overcome this challenge, we propose new kernel-based attribution method. This approach effectively captures complex task interactions and does not rely specific assump- tions about the training process. This makes it broadly applicable to emerging paradigms, including in-context learning and reinforcement learning. Furthermore, we introduce approximation technique that avoids expensive retraining, ensuring our method is scalable. We val- idate our method's effectiveness with experiments across several distinct domains. The experiment results demonstrate that our approach provides more accurate attribution than existing methods in a diverse range of settings that are relevant to modern AI systems.
efficient gradient
10

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGMENT
Thanks to Chris Ré for insightful discussions about this work. We appreciate the anonymous referees for their constructive comments and suggestions. This work supported in part by NSF award IIS- 2412008. The views and opinions expressed in this paper are solely those of the authors and do not necessarily reflect those of the National Science Foundation.
REFERENCES
Juhan Bae, Wu Lin, Jonathan Lorraine, and Roger Grosse. Training data attribution via approximate unrolling. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024.
Samyadeep Basu, Phil Pope. and Soheil Feizi Influence functions in deep learning are fragile. In International Conference on Learning Representations, 2021.
Gavin Brown, Mark Bun, Vitaly Feldman, Adam Smith, and Kunal Talwar. When memorization of irrelevant training data necessary for high-accuracy learning? In Proceedings of the 53rd annual ACM SIGACT symposium on theory of computing, pp. 123-132, 2021.
Sang Keun Choe, Hwijeen Ahn, Juhan Bae, Kewen Zhao, Minsoo Kang, Youngseog Chung, Adithya Pratapa, Willie Neiswanger, Emma Strubell, Teruko Mitamura, Jeff Schneider, Eduard Hovy, Roger Grosse, and Eric Xing. What is your data worth to gpt? l1m-scale data valuation with influence functions. arXiv preprint arXiv:2405. 13954, 2024.
Junwei Deng, Weijing Tang, and Jiaqi W. Ma. A versatile influence function for data attribution with non-decomposable loss. In Forty-second International Conference on Machine Learning, 2025.
Samuel Deng and Daniel Hsu. Multi-group learning for hierarchical groups In Proceedings of the s/ International Conference on Machine Learning, pp. 10440-10487, 2024.
Lee H Dicker, Dean P Foster, and Daniel Hsu. Kernll ridge vs. principal component regression: Minimax bounds and the qualification of regularization operators. Electronic Journal of Statistics, 2017.
Logan Engstrom, Andrew Ilyas, Benjamin Chen, Axel Feldmann, William Moses, and Aleksander Madry. Optimizing ml training with metagradient descent. arXiv preprint arXiv:2503.137511 2025.
Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware minimiza- tion for efficiently improving generalization. International Conference on Learning Represen- tations, 2021.
Shivam Garg, Dimitris Tsipras, Percy Liang, and Gregory Valiant. What can transformers learn in-context? case study of simple function classes. Advances in neural information processing systems, 25:30583-30599 2022.
Roger Grosse, Juhan Bae, Cem Anil, Nelson Elhage, Alex Tamkin, Amirhossein Tajdini, Benoit Steiner, Dustin Li, Esin Durmus, Ethan Perez, Evan Hubinger, Kamilè Lukošiūtē, Karina Nguyen, Nicholas Joseph, Sam McCandlish, Jared Kaplan, and Samuel R. Bowman. Studying large lan- guage model generalization with influence functions. arXiv preprint arXiv:2308. 03296. 2023.
Pingbang Hu, Joseph Melkonian, Weijing Tang, Han Zhao, and Jiaqi W Ma. Grass: Scalable influ- ence function with sparse gradient compression. arXiv preprint arXiv:2505. 18976, 2025.
Andrew Ilyas and Logan Engstrom. Magic: Near-optimal data attribution for deep learning, arXiv preprint 2504 16430, 2025.
Andrew Ilyas, Sung Min Park, Logan Engstrom, Guillaume Leclere, and Aleksander Madry. Data- models: Predicting predictions from training data. In International conference on machine learn- ing, 2022.

---

Published as a conference paper at ICLR 2026
William B Johnson. Extensions of lipshitz mapping into hilbert space. In Conference modern analysis and probability, 1984, pp. 189-206, 1984.
Kyle Kloster and David F Gleich. Heat kernel based community detection. In Proceedings of the 20th ACM SIGKDD international conference on Knowledge discovery and data mining, pp- 1386-1395, 2014.
Pang Wei Koh and Percy Liang. Understanding black-box predictions via influence functions. In International conference on machine learning, pp. 1894, PMLR, 2017.
Philipp Alexander Kreer, Wilson Wu, Maxwell Adam, Zach Furman, and Jesse Hoogland. Bayesian influence functions for scalable data attribution. In High-dimensions Learning Dynamics, 2025.
Alex Krizhevsky and Geoffrey Hinton. Learning multiple layers of features from tiny images. 2009.
Yongchan Kwon, Eric Wu, Kevin Wu, and James Zou. Datainf: Efficiently estimating data influence in lora-tuned lims and diffusion models. In The Twelfth International Conference on Learning Representations, 2024.
Dongyue Li, Huy Nguyen, and Hongyang R. Zhang. Identification of negative transfers in multitask learning using surrogate models. Transactions Machine Learning Research, 2023.
Dongyue Li, Aneesh Sharma, and Hongyang Zhang. Scalable multitask learning using gradient- based estimation of task affinity. In Proceedings of the 30th ACM SIGKDD Conference Knowl- edge Discovery and Data Mining, pp. 1542-1553, 2024a.
Dongyue Li, Ziniu Zhang, Lu Wang, and Hongyang R. Zhang. Scalable fine-tuning from multiple data sources: first-order approximation approach. In Findings of the Association for Computa- tional Linguistics: EMNLP 2024, pp. 0608-6623, 2024b.
Sadhika Malladi, Alexander Wettig, Dingli Yu, Danqi Chen, and Sanjeev Arora. A kernel-based view of language model fine-tuning. In International Conference on Machine Learning, pp- 23610-23641. PMLR, 2023.
Charles Micchelli. Interpolation of scattered data: distance matrices and conditionally positive definite functions. In Approximation theory and spline functions, pp. 143-145. Springer, 1984.
Sewon Min, Xinxi Lyu, Ari Holtzman, Mikel Artetxe, Mike Lewis, Hannaneh Hajishirzi, and Luke Zettlemoyer. Rethinking the role of demonstrations: What makes in-context learning work? In Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pp- 11048-11064, 2022.
Sung Min Park, Kristian Georgiev, Andrew Ilyas, Guillaume Leelere, and Aleksander Madry. Trak: Attributing model behavior at scale. In International Conference on Machine Learning, pp- 27074-2711 PMLR, 2023.
Tomaso Poggio and Federico Girosi. Networks for approximation and learning. Proceedings of the IEEE,
2002.
Garima Pruthi, Frederick Liu, Satyen Kale, and Mukund Sundararajan. Estimating training data influence by tracing gradient descent. Advances in Neural Information Processing Systems. 33: 19920-19930, 2020.
Adityanarayanan Radhakrishnan, Daniel Beaglchole, Parthe Pandit, and Mikhail Belkin. Mecha- nism for feature learning in neural networks and backpropagation-ffer machine learning models. Science, 14666666:)111114446 2024.
Aravind Rajeswaran, Chelsea Finn, Sham M Kakade, and Sergey Levine. Meta-learning with im- plicit gradients. Advances in neural information processing systems, 32, 2019.
Marco Tulio Ribeiro. Sameer Singh, and Carlos Guestrin. "why should trust you?" explaining the predictions of any classifier. In Proceedings of the 22nd ACM SIGKDD international conference knowledge discovery and data mining, pp. 1135-1144, 2016.
12

---

Published as a conference paper at ICLR 2026
Bernhard Schölkopf and Alexander J Smola. Learning with kemels: support vector machines, regularization, optimization, and beyond. MIT press, 2002.
Bernhard Schölkopf, Ralf Herbrich, and Alex J Smola. A generalized representer theorem. In International conference on computational learning theory, pp. 416-426. Springer, 2001.
Siwei Wang, Yifei Shen, Shi Feng. Haoran Sun, Shang-Hua Teng, and Wei Chen. Alpine: Unveil- ing the planning capability of autoregressive learning in language models. Advances in neural information processing systems, 37:1 19662-119688, 2024.
Siwci Wang. Yifci Shen, Haoran Sun, Shi Feng, Shang-Huu Teng, Li Dong, Yaru Hao, and Wei Chen. Benefits and pitfalls of reinforcement learning for language model planning: theoretical perspective. arXiv preprint arXiv:2559.226133 2025.
Chenwei Wu, Holden Lec, and Rong Ge. Connecting pre-trained language model and downstream task via properties of representation. Advances in Neural Information Processing Systems, 36: 47216-47238, 2023.
Sen Wu, Hongyang R. Zhang, and Christopher Ré. Understanding and improving information trans- fer in multi-task learning. In International Conference on Learning Representations, 2020.
Fan Yang, Hongyang R. Zhang, Sen Wu, Christopher Ré, and Weijie Su. Precise nigh-dimensional asymptotics for quantifying heterogeneous transfers. Journal of Machine Learning Research, 26 (113):1-88, 2025.
Tianhe Yu, Deirdre Quillen, Zhampeng He, Ryan Julian, Karol Hausman, Chelsea Finn, and Sergey Levine. Meta-world: A benchmark and evaluation for multi- task and meta reinforcement learning. In Conference on robot learning. pp. 1090- PMLR, 2020.
Hongyang R. Zhang, Dongyue Li, and Haotian Ju. Noise stability optimization for finding flat minima: hessian-based regularization approach. Transactions on Machine Learning Research, 2024. ISSN 2835-8856.
Zhenshuo Zhang. Minxuan Duan, Youran Ye, and Hongyang R. Zhang. Scalable multi-objeccivv and meta reinforcement learning via gradient estimation. In AAAI, 2026.
Ziniu Zhang Zhenshuo Zhang, Dongyue Li, Lu Wang, Jennifer Dy, and Hongyang R. Zhang. Linear-time demonstration selection for in-context learning via gradient estimation. Empirical Methods in Natural Language Processing. 2025.
13

---

Published as a conference paper at ICLR 2026
A COMPLETE PROOFS
Derivation of influence functions. Let S = {(11, 31), (22, y2)..... (In, vn.} be a dataset in- cluding n samples drawn independently from an unknown data distribution. Let fw denote model with parameters W E Rd. Let L(Jw) denote the empirical loss of the model fw S. averaged over the training data samples. The influence function (Koh & Liang, 2017), which is defined by how much trained model changes after adding or removing one sample (zz y). is given by
To derive this result, let the perturbed parameter after adding input with step size be given by:
We.= arg min (i(fw) t (ff(())
(10)
W
Define the parameter change A, W.,: W and note that, since W does not depend on thus write:
dW dA.
Based on first-order optimality conditions, from equation (10), we get:
vi(fw.. (lJW..)+ VtvVV((....))) = O.
Since W W O, we perform Taylor expansion centered at W up to first-order
vi(fm) + Ve(fw (z).u) (v v((lf(()),)) A. = 0,
which is approximately zero after dropping the second-order term.
After setting approaching zero, and solving for A., get:
A. - (v((f) + VVVl (z), =-
since vi(ww) = 0. Through the chain rule, the influence on any differentiable funetion F(s) becomes:
I.(s) = F())
By aggregating over = 1,2, K, we conclude the proof of equation (2).
Notations. Given square matrix X Rddd diag [X] Rd to denote the diagonal entries of organized as a vector. Let |X|| denote the largest singular value of X. Let ||>||| denote the Euclidean norm of a vector : E Rdd.
We follow the convention of big-O notations. We (nn Ogg(n) to indicate that there exists a constant e > 0 so that for large enough n, f(n) S c. g(n). We also write f(n) S g(n) to indicate that (n) (g(n)).
A.1 PROOF OF PROPOSITION 3.1
Proof. We denote by f E[s] = 1p]x- First, we derive the population ordinary least squares (OLS) coefficients for the linear predictor
g(s) = a Σ B,sh-
j=1
Let a² and 3* = (Bi. 3kk) denote the minimizers of the population squared loss
K
a" 8* + arg min E
(F(s)
Σ
8ys;)
)²
1=1
14

---

Published as a conference paper at ICLR 2026
Taking the derivative with respect to and setting it gives
E F(s) -
Σ 3J5 0,
(11)
We differentiate with respect each i and set the derivative to zero:
E (F((s)) a' -E =0 i= 1..... K.
Expanding the expectation and using linearity yields
E[s] - a''[s,]]
Σ 83 E(s,s,] 0.
(12)
From equation (11), we have
a² [F(s)]
Σ
33 [sy]
(13)
SND
SND
j=1
By multiplying both sides of equation (13) with s; and taking expectation on again, we get
E E Σ E
(14)
Subtracting and adding both sides from equation (14) into equation (12) leads the following covariance form
Covls;, F(s))
Σ 8J Covls,, syl,
K,
which are the population normal equations. Since subsets independent, have Sov = 0 for i * j.
so the covariance matrix of diagonal. Thus each normal equation decouples: Cov(se, F(s) 3 Var[sa].
Solving for ft gives the uni-variate OLS coefficient
3 Covist, F(s)]
(15)
Var[sk]
Next, we decompose the covariance using second-order Taylor expansion as:
Cov|sk, F(s)] = Cov(Sk, F(s*)) + Cov [Sk; G (s s*)] + Cov [k, (s - s*) Hg T s
A
Au
As
RR((s);
where we denote the gradient and Hessian as G.,, Hg. respectively, and the expansion error term as Ra (Sa). Notice that A = 0. Thus, A2 is the first-order term, and A3 is the second-order term.
Next, by the independence between different coordinates of s, we have that
A 9& Var[ss],
where 9k (VFF(s))).
For Ag. write s² = (s 11) + (μr -s) and expand the Hessian product as:
Q = - n) Ha(s μ) + 2(μ - 5*) H.(s p) + (4 s*) Ha -s)).
15

---

Published as a conference paper at ICLR 2026
For the second term above,
K
Cov(sk: (μ 5*) Hs(s [+]] = Var(sH [sk]
Σ
H
(k, - -s))
j=1
using symmetry of Hg, where Hg (k; denotes the jk, entry of Hg. For the first term,
K
(s - 11) H.(s (1) = _ Ha(i, 12)² +2
Σ
(s(i,))(s: -ms((s; - M4)).
ISKjSK
All cross terms with i # vanish by independence and zero mean of 11j. leaving
Cov Sk: (s 16) (g 14] - Ha, Cov (Sks (Sk )
2
With X2 = phx Cov (sk x _ and for Bernouli(p),
E = p(1 p) 2p).
Hence
Cov(sk; Q1
Var[sk]
Σ
H (k,k,))) -s) + HH, p)(1 -2p),
(16)
1=1
which yields
K
3 gk Σ
(k. -8)) Hk, - 2p) I (Ra(ss)/Vartte
(17)
By applying equation (17) over
.,,, K, we have shown that
8 VvF((s)) - (
1 diag H.
S
Rs(sk)((p11 pp)) p)), 1=1
(18)
2
where Var[ss] = p(1 pp) since 8& is a Bernoulli random variable with probability p. Assuming the third-order derivatives of F(s) bounded by some small enough cs, then, have that
Σ
Ra (sx) O (cs [ |s
s sT)² So 80g 88g (1-p)) p))
(19)
By applying equation (19) to equation (18), we obtain an error term as O (cggk3/2p-1)
The last step is to bound the minimizer between empirical loss ft( 3) and the population loss R(a, 8). Recall that 3 is the empirical risk minimizer on m samples. Since we are working noiseless setting, we have that
K + log(0-1) | - $ C(I[5(]
(20)
m
for some fixed constant C O. This can be shown using standard OLS analysis and concentration bounds on the sample covariance matrix. In the case of Bernoulli sampling, we have that the popu- lation covariance matrix p(1 - p) Id, thus <(E) = Thus, by combining equation (20) with the Taylor expansion errors above, have concluded the proof of equation (5).
Corollary A.1. In the setting of Proposition 3.1, the intercept and the residual error of linear sur- rogate models are given by
= E[)(s)) BT
K CovisE Q1
min E (F(s) - a - -T 3Tss)2 Var[Q] a
Var
where Q defined in equation (16).
16

---

Published as a conference paper at ICLR 2026
Proof. First, the intercept term follows from equation (13).
Next, consider the residual mean-squared error term. Let X = s - μ (zero mean, independent coordinates). The optimal linear predictor removes the constant and all components in the span of {-+}. Therefore
Cov(X Q] [
E Var (
B[Q]
bx
Var[Q]
Because Cov(X, Q] = Cov(sH: Q) and Var(X,] = Var[sk] = =p(1 p), this gives
K Cov|SR, Q1?
3 = min E [F(s) - a - 3J] ~ Var(Q]
(21)
a
Var [Sk]
remains to compute Var[Q]. Since Q = 1 H
HgX,
K
X H Σ (1,0) Σ H,,(,))x,,,
ISKjSK
+ VarlX H,X] = Σ
+4 Σ
Varl))))
ISKjSK
by independence and mean zero, eliminating mixed covariances. Thus
Var[Q) Var HH,X]
H.X].
For Bernoulli (p). X, E - p. -p) with
E[X?]| p(1 - p),
E[x = p(1 - p)4 + (1 - p)p²
Thus,
Var Var[X?] |[X;| E[ p(1 - p)(1 2p)
Var Var[X;X;] = || = (p(1 - ( (1)
Substituting back to Q gives
Var[Q] 1
Σ
p) 2p) +4
Σ (p(1)))(((1ppp))
(22)
ISKGSK
Plugging equations (16) and (22) into equation (21) gives the detailed residual mean-squared error.
A.2 PROOF OF COROLLARY 3.2
Next, we build on the above proof to derive Corollary 3.2. We show that the influence function estimates the gradients of the performance function. Our goal is to show that the vector of influence functions I = (ZI,...,III is equal the gradient vector VsF(s).
The optimal parameters (ss) arg miny L(fw. s). The empirical risk minimizer corresponds to the uniform weight vector s* = [1]x- Let F(s) be any differentiable performance metric that depends the optimal parameters, e.g., the loss on test sample (x. F(s) = (x),)) The gradient of this function is given by (s*).
Proof. Recall that the influence function uses an additive perturbation c for single task The perturbed objective is:
1 K
i(sw,
((w(())))))))
+ c
(((((x))))))
K
17

---

Published as a conference paper at ICLR 2026
We can rewrite this expression by collecting terms for each sample:
c(fw,c)= (L+)))((((Z)).)) + Σ [((f((zz).v))
K
Thus, the additive perturbation a specific instance that corresponds a path in the K K-dimensional weight space s:
s(e) = (1,1,.... Lts ....... 1)T = e.
i position
where e; refers to the i-th standard basis vector.
The influence function for task i, (s), is defined as the derivative of with respect to e, evaluated at 0. We can compute this derivative using the multivariate chain rule F(s) as:
dF(((c)) de
I.(s)
c=0
The chain rule states:
dF(((c))
ds(c) de
[V
(23)
de
By definition, s(E) =s ej, thus
ds(c) Ci-
de
Substituting this back to equation (23), we get:
dF(((()) ((((e))
aF(s(e))
de
Os;
Thus, the derivative of F(s(c)) with respect to is exactly the partial derivative with respect to the weight s,. Finally, consider O. We have s(0) s* Therefore:
(s) OF (s(e))
(F(s)
Oe
Os:
c=0
1==
Since this holds for all i 1,..., K, the vector of influence functions is identical to the gradient vector the performance function:
i = [I'(s), [s)) = v.F(s*).
Next, recall that ||H|| S 02 by assumption. As a result, the sampling bias and variance terms in equation (5) are at most:
C2 c2VK K + c2/2.
In summary, we have shown that | - S c2v K +c2/2. Together with the concentration bound between 3 and 3* from equation (20), we conclude that
K log(6
B 7 s (2(V + 1/2) + 0(2)k3/2--1 + +06
m
with probability at least 1 - 8 over the randomness m random subsets drawn from D, which completes the proof of equation (6).
18

---

Published as a conference paper at ICLR 2026
Algorithm 2 Gradient Estimation (GradEx) in the case of multi-class classification
Input: Subset vector s ( precomputed logits {fw, and gradients {Vw fw, test dataset Tiest and
test loss
Requires: fz-regularization À > O; random projection matrix P E R
Output: Estimated performance on the test task
1: Construct the sample subset s(i) = 2 = (z, 92 } from tasks selected in s
Step 1: Project gradients and solve low-dimensional regression problem */ 2: Compute projected gradients for all samples in the subset: 9j = PV fwo
|*
3: Define the low-dimensional objective over E RR:
1s1
X
T(Z))
Σ ECE (fw) (oj- 2),
4: Compute the optimal low-dimensional regression by minimizing this convex objective:
=aa min L(Z)
ZER*
1*
Step 2: Estimate test performance
5: y
6: for each (2r, v) E Tiea do
7:
Project the test gradient: ğ == PV fw, (z)
8:
Approximate the test logits using the low-dimensional projections:
f(x) = Swa (xr) + 2*)
9:
Accumulate the loss: y(t) + (0 + ((f(x),))
end for
11: (()) <- /
Normalize the loss over the test dataset
12:
return y)()
A.3 GRADIENT ESTIMATION ALGORITHMM
Gaussian random projection. This method directly addresses the high dimensionality of the gra- dient by projecting it into a low-dimensional subspace. We employ random matrix P E Rexd (where l: d) to map the gradient V /w((x) E Rd compressed representation
fw(()= = PV E R
Motivated by the proximately preserved, we then solve the perturbation objective using these low-dimensional fea- tures. The optimization is performed over k-dimensional vector, making the problem tractable and independent of the original parameter count d.
Lemma, which guarantees that pairwise distances are ap-
Gradient estimation for in-context learning. For our experiments on the ICL task, we define the loss function and compute gradients in the embedding space. This approach enables us to avoid perturbation-base objective optimization methods, as we already know difference in em- beddings. Thus, we can directly perform first-order estimation of the logit outputs for different prompts based on the gap between their embeddings. Consider the model output of a prompt subset S on input z, denoted as fw (0(S,))). Given anchor prompt So. the first-order approximation of fw around the embedding vector Φ(S,, x)) is given by:
(w(( Jw (w .)) (W x)), ((S, x)) (sSo. x)) tes...
(24)
A.4 FIRST-ORDER APPROXIMATION ERROR
First, kernel analyses of fine-tuning large transformers show that, in the kemel behavior regime, the model evolution is well aplraaaproxitttee by its first-order Taylor expansion around the pra-trdined parameters W, such that
fw(x: Sw, (2: + (V.wo (w), W Wo) +ew(x),
19

---

Published as a conference paper at ICLR 2026
Table 5: We report the approximation error on different sizes of models on SST-2 and Coin flip.
SST-2
Coin flip
DeepSeck-LLM-77B
7.644.5 10
6.5_2.0 X 10
Llama-3.1-8
1.840.1 x 10-2
7.5 x 10-3
Llama-2-133
7.7
10=4
1.520.7 10-3
CodeLlama-34B3
7.4
10-4
8.521.1 10 10-4
with Taylor expansion residual error term (w(x) on æ. This justifies working in the linearized regime around Wo even for highly nonlinear LLMs.
We show the approximation error various models with parameters from 7B to 34B. We test SST-2 and Coin flip datasets. Our findings in Table show that the errors remain negligible across different architectures and model sizes.
Second, assuming the first-order approximation error is small, we show that the gradient estimation algorithm gives accurate estimates. Let ((fw) be the empirical loss, assume that:
The average Taylor expansion residual bounded by 8:
E (x, fwo (z, V fw, (z, v) T (W -Wo]] s 6,
(x.y)
Gradients Wo are bounded by G, and
The optimization is restricted to domain of radius most D, with random projection that distorts inner products by at most e.
Denote by (SS) the parameter produced by the gradient-based -based estimator (obtained from regression in the projected gradient feature space, as in our gradient estimation step). One can show that
i((((s)))
min i.[w) + 25 4GDr.
WED
Thus, provided that the linearization analysis ensures that the first-order approximation error ew (ar) (and hence 8) small, the gradient-basee based estimator W(S) = W, Z achieves training loss within 28 + 4GDe of the minimum loss in the search domain. We summarize this discussion into the following proposition.
Proposition A.2. Let D C Rdd be a search space whose radius is most D. Suppose the gradient of fwo at the initialization Wo in the training set is at most G in Euclidean norm. For each task i 1,2...... n, let T, denote the training data. Suppose that for every
Σ (R.))TT
(fw(x,)) y) -f., (z. v) Vw fw. (x, (T (W WWW] 28.
Ti
Provided that the random projection dimension k satisfies k = 0 log the training loss of W(S)
is bounded away from the minimum training loss for any 5 C (1, n as (.J S min i.(w) + 25 + 4GD.
(25)
WED
The proof is based on the Johnson-LLindenstraus lemma (Johnson, 1984), which asserts that when
k
log N
for any 9: with ||9=|| G and any W, Wo in D, have
Kg, W Wo) - (Pt,, P(W - Wo) - <) W Woll < 2GDe
The rest of the argument can be completed via simple inequalities.
A.5 EXPRESSIVITY OF RBF KERNELS
Proposition A.3 (Expressivity of RBF kernels), Let X = {0, and k the Gaussian RBF kernel. Then:
20

---

Published as a conference paper at ICLR 2026
Exact interpolation: There exists a function Hy such that g(s) F(s) for all X. i) Universal approximation: Consequently, for any 0, there exists f E Hg such that max F(s) f(s)] S.
Proof. We first prove Part i). The set X consists of distinct points in I is well cell-established fact that the Gaussian RBF kemee strictly positive definite on any finite set of distinct points (Micchelli, 1984: Schölkopf & Smola, 2002).
Let m = | be the number of tasks. We enumerate the elements X as s!",...., s(() and let
Vj F(s()) be the target values. We define the Gram matrix Kg E mR with entries
Ko., k(s()) s01).
Since the kernel is strictly positive definite, the matrix K is strictly positive definite and therefore invertible (full rank). This implies that the linear system Kg0 has a unique solution I We construct the function
go (s) =
0s (s).
1=1
By construction, ge Hy and satisfies ())) y, F(s(i)) for all j. This proves Part
Part ii) follows immediately from Part i). Since g matches F' exactly on X, the approximation error zero, which is always less than or equal to any
The intuition from the above result is that {0, 1)*, two subsets that differ only few coordi- nates are close in Hamming distance, and intuitively their performances tend to be similar. The RBF kernel directly reflects this structure:
knar(s,s') exp (-203
1 20
s'112
decays smoothly as the Hamming distance between and s' increases. Thus, kernel regression assigns stronger influence to nearby subsets and weaker influence to distant ones, capturing the local smoothness we observe in practice.
By contrast, polynomial kernels do not encode distance-based locality. A degree-d polynomial kernel expands inputs into global monomials (e.g., 8;8j, 848,8k, ...), subsets with small Hamming distance are not necessarily mapped to nearby points in the feature space.
B OMITTED EXPERIMENTS
We begin by describing additional details that are omitted from the experimental setup in Section 4.1.
Modular arithmetic tasks. We consider the following functions with 97: + (mod p) = c, ab-bb
and + b2 (mod p) = c. For each task, we generate complete dataset by iterating through all possible pairs of (a. b), where a, E {0, 1.... 1}. This results in dataset of 97 x 97 9, 409 unique equations for each function. Each equation is formatted as a sequence of tokens: a. op, b, where op is for addition and placeholder token for the quadratic function. We randomly split the whole dataset into a training set comprising 90% of the data and test set comprising the remaining 10% of the dataset.
We use standard two-layer decoder-only transformer with embedding dimension 128 as the clas- sifier for all modular experiments. The model is trained to predict the correct token for the result c, given the preceding sequence (a, op, b. =). The cross-entropy loss and gradients are computed based on the model's 's output logits the final token position.
To support our task attribution analysis, we use a grouped subset sampling strategy that partitions the training data based the values of both operands, and b. We first partition the range of
21

---

Published as a conference paper at ICLR 2026
possible values for each operand, {0, 1.......),.. into 5 disjoint intervals. Specifically, training example (a, b, c) is assigned to a group Gt, where = |a/20] and bb/20 This creates a x grid of 25 groups, where each group coressonds to specific rectangular region in the input space. Subsets of the training data are then constructed by selecting specific combinations of these groups. This method enables fine-grained analysis of how the model learns from different regions of the input space, allowing for the systematic construction of task vectors, We sample 50 subsets with sampling ratio 0.9, and use 40 subsets the surrogate training set.
In-context learning. We use the Qwen3-8B as the base model and evaluate it on a sentiment classi- fication dataset, SST-2, and reasoning task, the coin flip. The SST-2 dataset is binary sentiment classification dataset composed of movie reviews labeled as either positive or negative from the GLUE benchmark. The number of queries is 450. The Coin-Flip dataset is an arithmetic reason- ing task where the model reads a natural language description of a sequence of fair coin flips and must predict the final outcome (heads or tails). The number of queries is 869. For each dataset, we use the first 50 candidate demonstrations as the data to be attributed. We sample 200 subsets, each containing prompts, and use 80 of these subsets as the surrogate training set.
Multi-objective reinforcement learning. We use MTI0 from the Meta- World benchmark (Yu et al., 2020), which consists of 10 diverse robotic manipulation tasks. The agent's s observation includes the environment state and a one-hot vector that specifies the current task. A sparse reward for moving the objective its goal position. The 10 tasks in MT10 are: reach, push, pick-place. door-open, drawer-open, drawer-close, button- citton-press-tee pross-topdown, peg-insert-side, window-open, and box-open.
For the task attribution evaluation, we use the Soft Actor-Critic (SAC) as the training algorithm, and take the reward of each task as the F(s). We sample 50 subsets, each containing 7 tasks, and use 40 subsets as the surrogate training set.
Hass training regularizers. Sharpness- method that aims find parameters lying in flat neighborhoods of the loss landscape, rather than just minimizing the loss at a single point. Concretely, instead of minimizing the empirical loss Ls( (ww). SAM minimizes the worst-case loss in an bp-ball of radius p around wt
Minimization (SAM) is an optimization
I.S.M min SAM (w)) min max Ls(w + E).
||||,SP
Using first-order Taylor approximation, the inner maximization has approximate solution
Ls
E(w) = PIV
for the common case p = 2. Each SAM step then computes the gradient of the loss at the perturbed weights, L.w + Eww),, and updates w in that direction, which biases training toward flatter minima that empirically generalize better than the sharp minima found by standard SGD.
NSO (Zhang et al., 2024) is an optimizer that explicitly encourages flat minima by minimizing noise-perturbed loss and thereby regularizing the trace of the Hessian. Given an empirical loss f(W) and zero-mean noise distribution P. NSO considers the smoothed objective
F(W) == Eu Ef + U)},
and uses a two-point noise injection scheme; each step it samples U P and averages gradients at symmetrically perturbed weights,
G(W,U) = (o(W +0) + g(W -)),
which cancels the first-order term in the Taylor expansion and keeps the second-order term
U V'F(U...
For isotropic noise with covariance Σ =V 22 Id), this yields
F(W) ~ ((W) + 12 ((W))
Thus, rumning gradient descent on F(-) with Gaussian convolution is approximately the same as running gradient descent on f(-) plus a penalty term of (a²/2). Tr[V2 S(W).
22

---

Published as a conference paper at ICLR 2026
Table 6: Top positive and negative influential prompt samples identified on the coin-flip task.
Query
Top positive prompt samples
Top negative prompt samples
Q: A coin heads up.
Q: A coin is heads up. Erdener flips the coin. Ismari does not flip the coin. Is the coin still heads up? (no)
Q: A coin is heads up. Pachl flips the coin. Lissett flps the coin. Is the coin still heads up? (yes)
Kielmeyer flips the coin. Jevgenij does not flip the coin. the coin still heads up? (no) Q: A coin is heads up. Bonnitta does not flip the coin. Ellise does not flip the coin. Is the coin still heads up? (yes)
A coin is heads up. Britting- does not flip the coin. Alilet does not flip the coin. Is the coin still heads up? (yes)
Q: A coin is heads up. Kimyetta flips the coin. Raynel does not flip the coin. Is the coin still heads up? (no)
Q: A coin is heads up. Roeland does not flip the coin. Joeliz flips the coin. Is the coin still
Q: A coin is heads up. Kulju does not flip the coin. Afrodisio flips the coin. Is the coin still heads
Q: A coin is heads up. Pachl flips the coin. Lissett flips the coin. Is the coin still heads up? (yes)
heads up? (no)
up? (no)
Attribution baselines. Influence functions (Koh & Liang, 2017) are proposed to measure the in- fluence of individual samples, which can also be adapted estimate the influence of entire tasks by considering a task-welt weighted loss function. The influence of infinitesimally up-weighting task Ta on the test performance given by:
(s) (Vw F(s) (VW Vvw
la lu(fw, Ej=183
8k
Here, = Lfw is the Hessian of the total training loss. In our implementation, we approximate the inverse Hessian-vector product H (where is the gradient of the task task-specific loss term) by solving a Lasso regression problem, which avoids direct matrix inversion.
TracIn (Pruthi et al., 2020) traces the influence of training data through the optimization trajectory, avoiding Hessian computation. We adapt this method from the sample level to the task level by measuring the correlation between task gradients over the course of training. The influence of training task Ti on a test task Test is approximated by summing the dot products of their respective loss gradients across various training checkpoints:
T
TracIn(Tk, Tiea)
Σ nh [VSll (la(JW,
where W, are the model parameters and 7/s is the learning rate checkpoint 4.
TRAK (Park et al., 2023) offers an efficient algorithm for data attribution by linearizing the model and using random projections. We adapt it attribute influence at the task level. The core idea is to represent each task by an average of its constituent samples' projected gradients. Specifically, for each sample z in a task, a feature vector computed from the gradient of a model output function fw. To get feature vector for an entire task Tj. we average these features over all its samples. This tevkl feature is then projected into a low-dimensional space using a random matrix P. The attribution score for the test task Ties is then computed as: ((Ties)) T ((Tuss) & is the projected feature vector projected features for the K training tasks, and Q is K K diagonal weighting matrix. The final scores are averaged over an ensemble of models to ensure robustness.
TT T Tb) - Q. Here, the test task, is the K matrix of stacked
Linear surrogate models (Ilyas al., 2022; Li et al., 2023) learn linear mapping from data weights to model predictions. Given a set of models trained on different data subsets, this approach fits linear function f f0, 1] - R that predicts test performance from subset indicators. The coefficients of this linear model serve as attribution scores, capturing the marginal contribution of each training example across multiple training runs.
23

---

Published as a conference paper at ICLR 2026
10
10-
Linear
Error
Error
Error
Error
RBF
2-
0.2
0.4
0.6
0.8
0.2
0.4
0.6
0.8
0.2
0.4
0.6
0.8
0.2
0.4
0.6
0.8
Subset split
Subset split
Subset split
Subset split
(a) Modular addition
(b) Modular quadratic
(c) SST-2
(d) Coin flip
Figure 4: We investigated how the size of the surrogate training split affects the residual error of the linear and kernel models. We observed that across range of training splits from 20% 80%, the kernel model consistently yields a lower residual error than the linear model.
B.1 OMITTED EXPERIMENTAL RESULTS
Qualitative results. In Table 6, we present the most positively and negatively prompt samples obtained by KERNELSM on the Coin flip dataset. We use bold to show information that is the same as the query, and italics for information that different from the query.
Correlation between surrogate models and influence functions. We designed an experiment using an fz-regularized logistic regression model (with a 10-2 Eg penalty) on the Wisconsin Breast Cancer dataset. The data was standardized and split into a training set of 455 samples and a test set. Our analysis focused on quantifying the influence of each training point on the loss of single, randomly selected test point.
We computed three distinct valuation scores for every training point: the ground-truth leave-one-out (LOO) score, the first-order influence function (IF) approximation, and the coefficients from linear surrogate model. The LOO scores were obtained by exhaustively retraining the model from scratch (cold-start) after removing each training point individually and recording the change test loss. The IF scores were calculated as a first-order approximation using the Hessian of the full training loss, stabilized with a 10 damping term. To build the surrogate model, we generated 1000 subsets, each containing 430 training points, and retrained model each to measure the resulting test loss change. These loss changes were then used as targets in a final linear regression, whose coefficients, solved via ordinary least squares (OLS), provided the surrogate model scores.
Nonlinear numerical simulation. We consider binary classification task in feature space. X' I² with labels y E 13. The base learning algorithm, denoted by A, is a two-layer MLP. The network architecture consists of input layer, two hidden layers with 16 neurons each and ReLU activation functions, and final output layer with a softmax activation.
To analyze the influence of different training samples, construct series of related but distinct training subsets. We first define a small, fixed set of N anchor data points, Sanchor = { yı N }. Next, define candidate set of K additional data points, C = {ci, c2,..., Cx ) which are sam- pled along the path that traverses a critical region near this boundary. The experiment then focuses on a series of K training subsets, where each subset Sy is formed by combining the fixed anchor set with exactly one candidate point from C: S = Sanciue U [(g, [e)]: for j 1,..., K. We select fixed test point Tiea X whose prediction is sensitive to the location of the decision boundary. The goal is to attribute the model's prediction on Frest to the choice of the candidate point Cj.
B.2 ABLATION STUDIES
Kernel methods often require higher sample complexity to converge to an optimal solution com- pared to linear models. We conduct an ablation study to investigate this trade-off. We compared the performance of the kernel surrogate model with that of the linear surrogate model, varying the size of the data subset used for their construction, sampling from 20% to 80% of the total subsets. The results in Figure 4 show that the kernel surrogate model consistently outperforms the linear surrogate model across all tested subset sizes. Notably, even when the number of samples was small (e.g., at the 20% subset level), the kernel method still demonstrated better performance.
24

---

Published as a conference paper at ICLR 2026
Table 7: We compare the performance of KERNELSM under different and 7.
À =
A = 10-!
= 10
== 10-3
y= 10-1
10-2
y= 10m3
y= 10
7= 10-5
Error LDS
0.81
0.65-2.01
0.73
0.86-0.10
0.65_o.01
0. 0.66=0.01
0.87.0.01
1.01
1.02+.03
0.52=0.05
0.54g0.01
0.51
0.51=0.06
0.54
0.53
0.52
0.51
0.51+0.04
Hyperparameters analysis. We evaluate the sensitivity of our approach to different hyperparam- eter settings in the in-context learning task in Table 7. Specifically, we vary À from 10 to 1 and from 10-5 to 10-1, We find that our approach is relatively robust to changes in both À and 7, exhibiting stable performance across the entire range.
In our experiments, we set À = 10-1 and y = 1/n as the default configuration, where n denotes the number of tasks.
Comparison of kernels. We use the CIFAR-I0 dataset and the ResNet-9 network classifier for our experiments. The dataset consists of 60,000 32 x 32 color images in classes, with 6,000 images per class. We utilize the standard training set of 50,000 images. To align with the scope of our task attribution methodology, we partition the 50,000 training samples into 50 disjoint tasks. Each task is constructed by randomly sampling 1,000 data points from the training set without replacement. This results in 50 distinct tasks, which form the basis of our analysis. For the training of our surrogate model, we randomly select 30 of these 50 tasks.
We evaluate the performance of different kernels, Let s(a) and (b)) denote different subset indices.
The kernels evaluated our study are:
Polynomial kernel (Degree-1, 2, 3): The general form the polynomial kernel k(s!a), ( ((s(a))T&s)
(b)) + c)"
In our experiments, we set the constant c = and test for degrees d E (1,2.3).
Radial Basis Function (RBF) kernel: The RBF kernel defined by the following equation:
k(s(a), s (*)) = exp - >||s(a) - 12))
where 7 is parameter that controls the smoothness of the kernel function relative to the
distance between s(a) and g
In Table 8, we find that the linear surrogate model and 1egree-1 kernel have large residual error. The degree-2 kernel achieves the lowest residual error among the three polynomial kernels. The RBF kernel achieves the lowest residual error among all the kernels.
B.3 EXTENSIONS AND DISCUSSIONS
Model-agnostic meta- (MAML) seeks to learn model initialization that is optimized for rapid adaptation to new tasks, typically with only a few gradient steps. Our work on task attribution, which measures model's sensitivity to its training data, is connected to this goal. The error of first-order attribution methods is directly governed by the Hessian. For attribution, a large Hessian signifies a highly curved, non-linear performance landscape where linear approximations are unre- liable, leading to significant attribution error. In contrast, MAML leverages this same curvature as signal. The meta-gradient update in MAML involves differentiating through an inner-loop gradient step, a calculation that explicitly depends on the Hessian of the inner-loop loss with respect to the model parameters. There also exists mathematical parallel between the derivation of influence functions and the formulation of implicit MAMI. (iMAML) (Rajeswaran et al., 2019). Influence functions can be derived by applying the implicit function theorem to the first-order optimality con- dition of the perturbed loss, yielding expression for the parameter change that involves the inverse Hessian. Analogously, iMAML uses the implicit function theorem to derive an analytical expression for the meta-gradient that depends only the solution of the inner optimization, not the path taken to reach it.
Multi group learning aims to train single predictor that performs robustly across predefined set of subgroups, addressing the common failure mode where high average accuracy masks poor perfor- mance critical aab-populations (Deng & Hsu, 2024). The influence of task k on the performance
25

---

Published as a conference paper at ICLR 2026
Table We investigate the surrogate model performance with different kernels. We run each exper- iment with five random seeds to report the standard deviations.
Residual error
Linear model
Degree-1 1
Degree-2
Degree-3
RBF
CIFAR-10
4.4
3.8
1.610.1
2.7
1.0
Modular arithmetic
4.621.3
2.2+0.5
1.740.4
3.3
1.540.4
In-context learning
0.8
0.640.1
0.540.1
0.5
0.440.1
Multi-objective RL
0.2
0.240.1
0.1
0.1
0.1
of model evaluated on task j can be directly interpreted as the marginal contribution of group k's training data to the model's performance on group j. positive influence value for a loss-based metric provides clear signal of negative transfer. Standard influence functions, being first-order approximations. capture pairwise, additive effects and struggle to model more complex issues. For example, model may perform poorly due to higher-order interactions between the data subgroups that linear methods cannot detect. By using non-linear RBF kernel, model learns global function that captures combinatorial effects between groups This can also inform the design of robust learning algorithms. For example, the MGL-Tree algorithm for hierarchical groups decides whether use a general parent-level predictor or specialized child-level predictor by comparing their empirical risks. Our method could provide a more principled, causal signal guide this choice.
interscctional group (e.g.. women of specific race)
Discussions. We now discuss this work with several related literature. There is an earlier line of work called LIME (Ribeiro et al., 2016), which addresses the question of explaining a single prediction at the feature level by fitting a surrogate that maps feature perturbations of specific input to changes the model output.
Our work instead focuses on training-date attribution by modeling how changing the weights of training examples affects the model's behavior. Although both approaches rely on local surrogate modeling, they operate on different objects. As a result, our work is complementary to this important line of work on LIME.
In terms of what's new in this paper relative to this earlier line of work, our first observation is to connect linear surrogate models to influence functions; this is shown Section 3.1, where we find the Ordinary Least Squares (OLS) regressor is approximately equal to influence functions (Koh & Liang. 2017). This observation partially explains why linear surrogate models can perform well empirically (Ilyas et al., 2022).
At the same time, we identify a limitation of linear surrogate models in the context of task/data at- tribution, which motivates kerne- based surrogate modeling. Kernels offer universal function approximation framework, enabling the surrogate to capture nonlinear dependencies between train- ing data and model behavior. To the best of knowledge, no prior work in data attribution uses kernel-bee based surrogate.
In the context of task attribution, kernel surrogate models are challenging estimate due the need for repeated training. Our second contribution is an efficient gradient-based estimator for fitting the kernel surrogate model. The computational cost consists of three parts: pre-computing fw gradient estimation, and surrogate fitting. Except for the pre-computation step, which only involves computing gradients once on the training data, all the remaining computations can be executed quickly CPUs.
As another remark, our in-context learning focuses on text classification or linear regression cases where the ordering of the examples does not affect the outcome. This setting builds on the prior work of Min et al. (2022), who found that for many ICL classification tasks, model outputs are often not sensitive to permutations of the context examples. We verify this in our own experiments by eval- uating the model under multiple random permutations of the same demonstration set: the standard deviation across different orders is only about 11% of the mean loss, indicating that order effects are relatively minor in our setting. This justifies our choice to model only which demonstrations are included, rather than their specific ordering.
26

---

Published as a conference paper at ICLR 2026
Beyond this setup, agree that there important scenarios, such math reasoning with chain- oh-thought style prompts, where the order demonstrations plays a significantly larger role. Our framework be naturally extended to handle such cases. Concretely, ings of the same demonstration set as different inputs to the surrogate model. For each permutation of fixed set of examples,
treat different order-
would evaluate the corresponding performance F(s, T) and fit the kernel surrogate on (s, T) thereby assigning different scores to different orderings. In this way, our method can explicitly model and attribute order-sensitive effects when they significant. This would be an interesting question for future work.
27


Reviews:
[{"review_id": "nmz5S0vtC5", "reviewer": "Reviewer_3bLy", "summary": "This paper introduces KERNELSM, a kernel-based surrogate modeling approach for task attribution in multi-task learning. The method captures nonlinear relationships between training tasks and model performance, going beyond the linear assumptions of prior approaches. The paper provides (1) a theoretical framework connecting influence functions and linear surrogate models, clarifying their limitations; (2) an efficient gradient-based algorithm for estimating kernel surrogate models without retraining; and (3) comprehensive experiments across arithmetic reasoning, in-context learning, and multi-objective reinforcement learning, showing consistent improvements in attribution quality.", "questions": "1. While Appendix C.4 provides an empirical comparison among polynomial and RBF kernels, could the authors offer more theoretical justification or intuitive reasoning for why the RBF kernel performs best in this method?\n\n2. Is this method sensitive to the choice of hyperparameters (λ and γ)? Are there any guidelines or heuristics for selecting these hyperparameters in practice? I am asking this because some other data attribution methods (e.g., TRAK) have already had practical heuristics for hyperparameter selection. If KERNELSM is highly sensitive to λ and γ, the additional hyperparameter search could increase the computational cost of applying the method in practice.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "1. The paper is well-structured and clearly written. \n2. The theoretical discussion is solid; it connects influence functions and linear surrogate models, offering valuable insights into their relationship and the limitations of linear approximations.\n3. The method seems to perform well against all baselines, and the evaluation tasks are comprehensive (modular arithmetic reasoning, in-context learning, multitask RL).", "weaknesses": "In Table 4, the comparison seems to exclude the task-level adapted versions of TRAK and TracIn, and the paper also does not specify their hyperparameter settings, which may reduce the credibility of the results reported in Table 3.\n\nAlthough the evaluation covers a diverse and comprehensive set of tasks (including arithmetic reasoning, in-context learning, and multi-objective reinforcement learning), the actual number of training tasks in each setting remains relatively small (only 10–50).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "fzaqLRkSqU", "reviewer": "Reviewer_SaVr", "summary": "The paper presents a framework for estimating how each training task impacts performance on a target task. By extending beyond linear models, to capture some of the non-linear interactions between tasks, significant empirical performance gains are found to emerge. This kernel based approach is validated with a range of domains and found to correlate more favourably with the ground truth LOO approach.", "questions": "What motivated the choice of setting gamma to 10^-5? (section C4) \nIf a sweep of hyperparameters was run, how sensitive are the empirical results to these values?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The manuscript is well written and clearly presented.\n\nThe approach presented in the paper is novel, well motivated, and timely.\n\nA broad range of experiments were performed with a suitable set of baselines, and empirical performance of the proposed approach is consistently strong.", "weaknesses": "Motivations for the kernel hyperparameters could be more clearly presented.\n \nFigures 3 and 4 appear to be lacking uncertainty estimates", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "sJhPUEysw9", "reviewer": "Reviewer_Wqdp", "summary": "This paper studies \"task attribution,\" which is a task to quantify the influence of individual training tasks on a model's performance for a given target task. \n\nThe paper aims to become less computationally expensive compared to leave-one-out techniques while being more accurate than linear models which only capture first-order effects.\n\nThis work introduces the use of kernel surrogate models to capture nonlinear task interactions better. The authors develop a gradient-based estimation procedure for fitting these kernel models, leveraging first-order approximations to avoid repeated retraining. \n\nFinally, the proposed approach is validated on different tasks showing improved accuracy.", "questions": "There is a paper called \"Register Always Matters\" where they find how different genres of data in the pretraining affect the performance of the model on the test samples. I am wondering if one can use your method to validate their results? No need to answer if you don't know or etc.", "limitations": "", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "Computational efficiency (comparable to linear surrogate runtime) while being much more accurate (<1% relative error)\n\nClearly explain and support the fact about why linear surrogates work well only when the task does not have interaction (in section 3.1)\n\nSimple and smart method to get rid of multiple retraining of the model in different input subsets.", "weaknesses": "Minor point: you didn't clearly define KernelSM. Although it is obvious that it stands for Kernel Surrogate Model.\n\nThe discussion in the results is not well organized and is confusing. For instance, table 3 is not referenced and discussed anywhere\n\nSome components are not explained and well discussed (e.g. SAM and NSO) or they are in the appendix although it would be better to be in the main body", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "5colIccBBO", "reviewer": "Reviewer_yBwo", "summary": "The paper studies the problem of task attribution: quantifying how each training task affects performance on a target task. The authors propose KernelSM, a kernel surrogate model over subset-indicator vectors trained via kernel ridge regression to capture higher-order task interactions. The method aims to address two limitations of existing approaches: the computational cost of leave-one-out evaluation and the inability of linear surrogate models to capture nonlinear interactions. Experiments on modular arithmetic, in-context learning, and multi-objective RL show that KernelSM achieves higher correlation with leave-one-out estimates and improves downstream task selection.", "questions": "In tasks like in-context learning, the order of examples can substantially affect the final outcomes. Could the authors clarify how their method accounts for or mitigates this issue?", "limitations": "", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "* Task attribution is an important and broadly applicable problem, as demonstrated by the diverse experimental settings.\n* Modeling task interactions with kernel methods is a reasonable and well-motivated approach.", "weaknesses": "* In my opinion, the issue of capturing nonlinear interactions is relatively straightforward and has already been explored in feature or data-selection settings using linear surrogate methods such as LIME - KernelLIME. This limits the contribution, as the paper essentially replaces linear surrogates with kernel-based ones.\n* There is no theoretical support for the gradient-based estimator obtained by linearizing the model in Section 3.3. It is surprising that this approach works well for highly nonlinear models such as LLMs. The authors should provide more intuition for why this estimator is effective and whether it generalizes to larger model sizes.\n* Baselines: the paper doesn’t discuss and compare with several recent works beyond TRAK, such as [1, 2], which address related problems in data and task attribution.\n\n[1] Bae, Juhan, et al. \"Training data attribution via approximate unrolling.\" Advances in Neural Information Processing Systems 37 (2024): 66647-66686.\n[2] Kreer, Philipp Alexander, et al. \"Bayesian Influence Functions for Scalable Data Attribution.\" High-dimensional Learning Dynamics 2025.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)