Efficient Estimation of Kernel Surrogate Models for Task Attribution

Technical Summary: Efficient Estimation of Kernel Surrogate Models for Task Attribution

Problem Statement
Modern AI systems, particularly large language models (LLMs) and multi-task learners, are trained on diverse sets of tasks (e.g., translation, code generation, reasoning). A fundamental challenge in interpretability is task attribution: quantifying the influence of each individual training task on the model's performance on a specific target task. The ground-truth metric for this is Leave-One-Out (LOO) retraining, which measures the performance change when a specific task is excluded from training. However, LOO is computationally prohibitive at scale, requiring $K+1$ full training runs for $K$ tasks.

Existing efficient alternatives include influence functions (estimating the effect of infinitesimal data perturbations via Hessian-vector products) and linear surrogate models (learning a linear mapping from task inclusion indicators to performance). While linear surrogates are scalable, they are limited to capturing first-order, additive relationships. They fail to model nonlinear interactions between tasks, such as synergy, antagonism, or XOR-type effects, which are common in complex learning landscapes.

Methodology
The paper proposes KERNELSM (Kernel Surrogate Models), a framework that extends task attribution to capture second-order and higher-order nonlinear interactions between training tasks.

1. Theoretical Analysis of Linear Surrogates:
   The authors first establish a unified task-weighting framework. Through a second-order Taylor expansion of the performance function $F(s)$ around the uniform weight vector $s^* = [1]^K$, they analyze linear surrogate models.

   • They prove that linear surrogate coefficients ($\hat{\beta}$) approximate the gradient of the performance landscape ($\nabla_s F(s^*)$) up to terms induced by the sampling distribution and second-order interaction effects.
   • Crucially, they demonstrate that under the assumption of small second-order interactions (i.e., a nearly linear landscape), linear surrogate models are approximately equivalent to influence functions. This explains why linear methods work well in some settings but fail when task interactions are non-additive.

2. Kernel Surrogate Modeling:
   To address non-additivity, the authors replace the linear mapping with a nonlinear function learned via Kernel Ridge Regression (KRR) in a Reproducing Kernel Hilbert Space (RKHS).

   • Kernel Choice: They utilize the Radial Basis Function (RBF) kernel ($k(s^{(a)}, s^{(b)}) = \exp(-\gamma \|s^{(a)} - s^{(b)}\|^2)$). This choice is motivated by the geometry of the binary subset space $\{0, 1\}^K$, where the RBF kernel naturally encodes locality based on Hamming distance, allowing the model to capture complex dependencies between task combinations.
   • Objective: The surrogate function $g_\theta$ minimizes the empirical loss between observed performance $F(s^{(i)})$ and predictions, regularized by the RKHS norm.

3. Efficient Gradient-Based Estimation:
   A major bottleneck in surrogate modeling is the need to train $m$ separate models to obtain performance values $F(s^{(i)})$ for $m$ sampled subsets. The authors introduce an efficient estimation procedure to avoid repeated training:

   • First-Order Approximation: They leverage a first-order Taylor expansion of the model output $f_W(x)$ around a pretrained initialization $W_0$: $f_W(x) \approx f_{W_0}(x) + \langle \nabla f_{W_0}(x), W - W_0 \rangle$.
   • Gradient Features: Instead of retraining, they estimate the optimal parameter perturbation $Z^*_s$ for a subset $s$ by solving a regularized regression problem where the features are the gradients of the pretrained model ($\nabla f_{W_0}$).
   • Dimensionality Reduction: To handle high-dimensional gradients, they apply Gaussian random projections (Johnson-Lindenstrauss lemma) to compress the feature space, enabling fast computation on CPUs.
   • Accuracy: Empirical results show this first-order approximation incurs less than 2% relative error across diverse datasets and model sizes (up to 34B parameters).

Key Contributions

• Theoretical Unification: The paper establishes a formal connection between linear surrogate models and influence functions, showing they are equivalent under conditions of negligible second-order task interactions.
• Kernel Surrogate Models (KERNELSM): The introduction of a kernel-based approach (specifically RBF) that effectively captures nonlinear task interactions (synergy/antagonism) which linear models miss.
• Scalable Estimation Algorithm: A gradient-based estimation procedure that approximates the performance of any task subset using gradients from a single pretrained model, eliminating the need for repeated retraining while maintaining high accuracy.
• Empirical Validation: Comprehensive evaluation across modular arithmetic reasoning, in-context learning (ICL), and multi-objective reinforcement learning (RL).

Experimental Results
The authors evaluate KERNELSM against baselines including Influence Functions, TracIn, TRAK, SOURCE, and Linear Surrogate Models.

• Attribution Accuracy: KERNELSM consistently outperforms baselines. In modular arithmetic tasks (involving XOR and quadratic structures), it improves attribution accuracy by up to 42% relative to existing methods. Across all domains, it achieves a 25% higher correlation with the LOO ground truth compared to linear surrogates and influence functions.
• Downstream Task Selection: When used to select optimal subsets for downstream tasks:
  • In-Context Learning: KERNELSM yields a 40% lower loss in demonstration selection compared to linear surrogates.
  • Multi-Objective RL: It improves target rewards by 15% in the Meta-World MT10 benchmark.
• Scalability: The method scales to large numbers of tasks (e.g., 500 prompts) without performance degradation and runs efficiently on CPUs.
• Robustness: The kernel approach remains robust under different optimizers (SGD, SAM, NSO) and Hessian regularizations, whereas linear surrogates degrade significantly when the loss landscape is highly curved (high Hessian trace).

Significance and Claims
The paper claims that KERNELSM provides a scalable and flexible solution for task attribution in black-box systems. By moving beyond the additive assumptions of linear models, it enables a more accurate understanding of how complex, nonlinear interactions between training tasks shape model behavior. The authors emphasize that their approach does not rely on specific assumptions about the training process (beyond the first-order approximation of the pretrained model) and is broadly applicable to emerging paradigms like in-context learning and reinforcement learning. The work bridges the gap between theoretical influence functions and practical, scalable data attribution for modern large-scale AI systems.