Interpretable AI with Local Distillation

Erin Craig ^†^†thanks: Corresponding author. Email: ercr@umich.edu. Affiliation: and Affiliation: Yiling Huang Department of Biostatistics Affiliation: Department of Statistics, University of Michigan Affiliation: and University of Michigan Affiliation: Snigdha Panigrahi Affiliation: Department of Statistics, University of Michigan

Abstract

Modern AI models such as tabular foundation models and gradient-boosted ensembles can outpredict classical methods, but provide little basis for reasoning about their predictions. High-stakes decisions call for models that are both accurate and interpretable as built. Local linear modeling offers a path forward: a smooth regression function is locally well approximated by a linear one, allowing a linear fit near each query point to achieve high accuracy without sacrificing transparency. The challenges lie in learning what is “local” and developing statistical tools for interpretation.

Here, we propose local distillation, in which a black-box “teacher” guides a regularized linear “student” model at each query point. The teacher (1) defines locality by upweighting training observations with similar predicted outcomes, and (2) anchors the fit with its prediction at the query point, included as a pseudo-observation whose weight is estimated from the data. For interpretation, we add a small amount of Gaussian randomization to the local objective and use refits to assess stability: selection frequencies identify reliable features at a query point, and clustering the randomized fits identifies stable subgroups across the data. Under the lasso penalty, we prove that this randomization yields feature-selection probabilities that are stable under small perturbations of the training responses. Across 17 benchmark datasets, local distillation nearly matches its AI teacher’s accuracy while producing a sparse linear model at each test point. In a high-dimensional cancer gene expression example, the framework identifies patient subgroups whose local models use different genes; this heterogeneity is invisible to a global linear model, and difficult to surface in a black-box model.

Keywords: Local regression; Interpretable AI; Distillation; Randomization; Stability; Tabular foundation models

1 Introduction

Advances in artificial intelligence (AI) are rapidly changing what is predictable [1, 2]. Modern black-box models, such as foundation models and gradient-boosted ensembles, sometimes outpredict simpler classical models even in their traditional strongholds, such as small-sample tabular data [3, 4, 5, 6, 7]. Yet, acting on predictions requires more than accuracy alone: decision-makers must reason about the model’s predictions. Black-box models offer little basis for such reasoning, which limits the value of their predictions for decision-making.

The standard response, when opting for a black box, is to explain its predictions post hoc. Popular tools such as LIME [8] and SHAP [9] fit per-observation attributions that quantify how much each feature contributes to a given prediction. However, because such explanations are constructed separately from the model, their fidelity to it is not guaranteed. Moreover, the attributions can be unstable, changing with choices external to the model and data, such as LIME’s perturbation scheme or SHAP’s reference distribution [10, 11, 12, 13].

By contrast, our work is guided by the view that transparency and reasoning should come from the predictive model itself, i.e., predictions should be interpretable as produced, rather than explained through post hoc tools [14]. To achieve this goal, we propose local distillation, a method in which a high-performing black-box “teacher” guides a local linear “student” fit at each observation or query point. Because a smooth regression function is locally well approximated by a linear function, a locally fit simple model can approach the accuracy of a flexible black box while remaining transparent at the query point. Figure 1 illustrates local distillation on the Auto MPG dataset [15], where we predict fuel economy (miles per gallon, MPG). In the test set, each car receives its own sparse linear model; together, they improve on the global lasso’s prediction squared error (PSE) by 48% while nearly matching their foundation model teacher (PSE 5.59 vs. TabPFN 5.25; global lasso 10.81). The local coefficients show that the number of cylinders predicts fuel economy among the least efficient cars but is not predictive among the most efficient, where engine displacement is more useful. A global linear model averages over this difference and assigns both a coefficient of zero. We include a high-dimensional gene expression example in Section 6 and benchmarks across 17 datasets in Section 7 and Appendix B.

Refer to caption — Figure 1: Car-specific miles-per-gallon models from local distillation. Local coefficients for the Auto MPG data using local distillation with a tabular foundation model teacher, TabPFN, and a lasso student. Cars are ordered left to right by TabPFN’s predicted MPG; coefficients are per standard deviation of each feature. The column on the left shows the global lasso coefficients. On a held-out test set, local distillation improves over the global lasso’s PSE by 48% (global lasso regression PSE 10.81; TabPFN 5.25; local distillation 5.59).

The central question, then, is how to define “local”: which observations should inform the fit at a given query point? Classical local regression defines locality through unsupervised similarity in input space (LOESS [16]), which ignores the response variable and is known to be vulnerable to the curse of dimensionality. Recent work instead defines locality through supervised weights derived from random forests [17, 18, 19, 20]. Local distillation, as proposed in our work, takes this idea further, defining locality in a supervised and tuning-free way: the teacher plays two roles, (1) identifying which training observations are informative via similarity of its predictions, and (2) pulling the student’s prediction toward its own, with the strength of the teacher’s influence determined from the data by the cross-validated student-to-teacher loss ratio. If the teacher does not outperform the student, the method simply reverts to a global linear fit.

Using the teacher’s predicted response to define similarity collapses the $p$ –dimensional feature space into a single interpretable axis along which the model is localized. We define locality this way for two reasons: (1) it is an axis along which the feature–outcome relationship often varies, and (2) it is an axis of scientific and clinical interest. This construction gives each local model a natural interpretation: in medicine, for example, a patient’s model describes the feature–outcome relationship among patients within a similar risk group. Moreover, when the teacher is a pretrained foundation model such as TabPFN [3], the student inherits its prior knowledge.

For decision-makers, the ultimate goal is not only accurate and interpretable prediction, but also trustworthy conclusions drawn from those predictions. Trust is closely tied to the stability of such conclusions under small perturbations of the training data, as emphasized by the PCS framework [21, 22, 23]. Motivated by this perspective, we propose a simple modification to local distillation based on external randomization, treating a conclusion as interpretable to the extent that it can be shown to be stable. The resulting randomized distillation fits produce stable local conclusions—such as identifying which features are important for its prediction at a query point—without breaking the inherent dependence among features or requiring teacher predictions to be recomputed. Our choice of the randomization scheme is supported by stability guarantees for feature selection, and it preserves fidelity to the notion of locality learned from the training data. Beyond localized conclusions, the randomized fits also reveal how the local feature–outcome relationship varies across observations: we cluster observations using their randomized local coefficients to learn subgroups that are stable rather than artifacts of any single fit.

This work contributes (1) local distillation, a predictive method that uses a black-box teacher to construct accurate, sparse local linear fits (Sections 2 and 7), and (2) a randomized stability framework for interpreting these fits, both at individual query points and in aggregate, with theoretical guarantees (Sections 4 and 5). In two case studies, the framework reveals heterogeneity in the feature–outcome relationship that a global linear model cannot express (Sections 4.3 and 6), a central but elusive goal of prediction methods, particularly in personalized medicine.

2 Local distillation

Suppose we have training data $\bm{X}\in\mathbb{R}^{n\times p}$ (standardized), a continuous response $\bm{y}\in\mathbb{R}^{n}$ , and a test point $\bm{x}^{*}$ . We also have a teacher $\hat{\phi}:\mathbb{R}^{p}\to\mathbb{R}$ , a fitted model that predicts $y$ from $\bm{x}$ . Our goal is to predict $y$ at $\bm{x}^{*}$ with a sparse linear model fit locally to the training data.

We describe local distillation using squared-error loss in Section 2.1 and illustrate it with an example in Section 2.2. The method is general, and Section 2.3 covers its extension to other loss functions and forms of regularization.

2.1 Model fitting and prediction

Our work is inspired in part by knowledge distillation [24], in which a simpler “student” model is trained under the guidance of a more complex “teacher” model. The student is usually a smaller neural network, and the goal is compression: a fast, lightweight model that approximates a costly one. Applying the same principle to regression with a linear student would yield

\hat{\bm{\beta}}=\argmin_{\bm{\beta}}\;\frac{1}{2n}\left[\sum_{j=1}^{n}\bigl(y_{j}-\bm{x}_{j}^{\top}\bm{\beta}\bigr)^{2}+\mu\sum_{j=1}^{n}\bigl(\hat{\phi}(\bm{x}_{j})-\bm{x}_{j}^{\top}\bm{\beta}\bigr)^{2}\right],

(1)

where hyperparameter $\mu$ determines the influence of the teacher.

When $\mathbb{E}[y\mid\bm{x}]$ is highly nonlinear, however, a single linear student cannot match a flexible teacher’s predictive performance. We therefore introduce local distillation, summarized in Algorithm 1, in which a separate student model is fit for each test observation. The teacher guides the definition of the local neighborhood around the query point and anchors the fit through its prediction at that point.

More precisely, our proposed method in Algorithm 1 modifies the standard knowledge distillation objective in (1) in three ways:

1.

First, we fit a separate local model for each test observation $\bm{x}^{*}$ , replacing the global distillation sum with a term that pulls the student’s prediction ${\bm{x}^{*}}^{\top}\bm{\beta}$ toward the teacher’s prediction; see (4). This anchors the local fit to the teacher’s prediction at the query point.
2.

Second, we replace uniform training weights with similarity weights ${\{\hat{S}_{j}(\bm{x}^{*}):j\in\{1,\ldots,n\}\}}$ that upweight training points whose teacher prediction is close to $\hat{\phi}(\bm{x}^{*})$ ; see (3). These weights define locality around the query point by determining which training observations belong to its local neighborhood.
3.

Third, we estimate $\hat{\mu}$ from the data as the cross-validated student-to-teacher loss ratio, rather than tuning it; see (2). We scale it further by $1/\sqrt{\hat{n}_{\text{eff}}}$ , where $\hat{n}_{\text{eff}}$ is the effective sample size defined in (3). The teacher’s influence therefore grows both when the teacher outperforms the student globally (larger $\hat{\mu}$ ) and when the local neighborhood is sparse (smaller $\hat{n}_{\text{eff}}$ ), where the local fit most needs anchoring.

When $\hat{\mu}\leq 1$ , the teacher offers no improvement and we instead return the global linear fit for all test observations.¹¹ 1 This choice is deliberately conservative: a teacher that offers no global improvement may still help locally, but estimating where requires localized loss ratios, which we found too variable to be reliable (see “Further localization”, Section 2.3).

In an ablation study across our benchmark datasets (Appendix C), we find that using both the similarity weights and the teacher prediction anchor yields better predictions than using either alone.

In Algorithm 1 and throughout, we write $\hat{\phi}^{(-j)}(\bm{x}_{j})$ for the teacher’s prediction at training point $j$ computed without access to $(\bm{x}_{j},y_{j})$ : if the teacher is fit to the training data, or uses it as context at prediction time (as with the in-context learning of tabular foundation models), $\hat{\phi}^{(-j)}$ is the out-of-fold (OOF) prediction; if the teacher makes no use of the training data, $\hat{\phi}^{(-j)}=\hat{\phi}$ . We additionally include an unpenalized intercept, which we suppress in the notation for clarity.

Algorithm 1 Local distillation for regularized linear models

Input: Training data $(\bm{X},\bm{y})$ with $n$ observations; test observation $\bm{x}^{*}$ ; teacher $\hat{\phi}$ ; elastic-net parameter $\alpha\in[0,1]$ .
Output: Prediction $\hat{y}^{*}$ and local coefficients $\hat{\bm{\beta}}^{*}$ at $\bm{x}^{*}$ .

Estimate distillation strength as the cross-validated student-to-teacher loss ratio:

\hat{\mu}=\frac{\hat{L}(f_{\hat{\bm{\beta}}})}{\hat{L}(\hat{\phi})},\qquad\hat{L}(g)=\frac{1}{n}\sum_{j=1}^{n}\left(y_{j}-g^{(-j)}(\bm{x}_{j})\right)^{2},

(2)

where $f_{\hat{\bm{\beta}}}$ is the global elastic-net fit with shrinkage parameter $\hat{\lambda}$ chosen by cross-validation, and $g^{(-j)}$ denotes prediction at $\bm{x}_{j}$ without access to $(\bm{x}_{j},y_{j})$ .
If $\hat{\mu}\leq 1$ , return $\hat{y}^{*}=f_{\hat{\bm{\beta}}}(\bm{x}^{*})$ and $\hat{\bm{\beta}}^{*}=\hat{\bm{\beta}}$ , and stop.

Compute similarity weights and effective sample size:

\hat{S}_{j}=\hat{S}_{j}(\bm{x}^{*})=\frac{\exp(-d_{j})}{\sum_{k=1}^{n}\exp(-d_{k})},\qquad d_{j}=\frac{\bigl(\hat{\phi}^{(-j)}(\bm{x}_{j})-\hat{\phi}(\bm{x}^{*})\bigr)^{2}}{\hat{\sigma}^{2}_{\phi}},\qquad\hat{n}_{\text{eff}}=\biggl(\sum_{j=1}^{n}\hat{S}_{j}^{2}\biggr)^{-1},

(3)

where $\hat{\sigma}^{2}_{\phi}$ is the empirical variance of $\{\hat{\phi}^{(-j)}(\bm{x}_{j})\}_{j=1}^{n}$ .

Fit the local model and predict: $\hat{y}^{*}=\bm{x}^{*\top}\hat{\bm{\beta}}^{*}$ , where

\hat{\bm{\beta}}^{*}=\argmin_{\bm{\beta}}\;\frac{1}{2}\sum_{j=1}^{n}\hat{S}_{j}\bigl(y_{j}-\bm{x}_{j}^{\top}\bm{\beta}\bigr)^{2}+\frac{\hat{\mu}}{2\sqrt{\hat{n}_{\text{eff}}}}\bigl(\hat{\phi}(\bm{x}^{*})-\bm{x}^{*\top}\bm{\beta}\bigr)^{2}+\hat{\lambda}\bigl[\alpha\|\bm{\beta}\|_{1}+\frac{(1-\alpha)}{2}\|\bm{\beta}\|_{2}^{2}\bigr].

(4)

Note: For a set of test observations, run step 1 once and steps 2–3 independently (in parallel) for each $\bm{x}^{*}$ .

2.1.1 The teacher as a Bayesian prior

Our approach has a natural Bayesian interpretation: the teacher’s prediction at the test observation acts as a Gaussian predictive prior on $\bm{x}^{*\top}\bm{\beta}$ , centered at $\hat{\phi}(\bm{x}^{*})$ with precision $\dfrac{\hat{\mu}}{\sqrt{\hat{n}_{\text{eff}}}}$ . This connects to data augmentation priors [25, 26], in which prior beliefs are encoded via pseudo-observations rather than parameter distributions. These priors express beliefs about observable quantities $y\mid\bm{x}$ , which are typically more interpretable than beliefs about regression coefficients, and yield posterior inference via weighted least squares on a modified dataset.

Our method differs in where the prior comes from and how it is deployed: rather than using a fixed set of elicited prior locations shared by one global fit, each test observation gets its own local model with a single pseudo-observation at the query point, centered at the teacher’s prediction. Additionally, we estimate the prior precision $\dfrac{\hat{\mu}}{\sqrt{\hat{n}_{\text{eff}}}}$ from the data, in the spirit of empirical Bayes. This parallels the power prior of Ibrahim and Chen [27], where historical data enters the likelihood with a tunable weight; $\hat{\mu}$ plays the analogous role and is estimated directly from the student-to-teacher loss ratio.

2.2 A worked example of local distillation

We illustrate the three steps of Algorithm 1 using the Auto MPG data [15], a set of $n=392$ vehicles from the 1983 American Statistical Association Exposition with $p=8$ predictors after encoding, including engine characteristics, vehicle weight, model year, and region of manufacture. Our goal is to predict fuel economy in miles per gallon.

We use the lasso as our student model, and TabPFN as the teacher. Using a 60/40 train/test split, we proceed as follows:

1.

Estimate distillation strength. The global lasso has CV PSE $12.81$ ; the TabPFN teacher $6.31$ ; therefore, $\hat{\mu}=\dfrac{12.81}{6.31}=2.03$ . Since $\hat{\mu}>1$ , the teacher improves on the global lasso and we proceed; had $\hat{\mu}\leq 1$ , we would have returned the global fit.
2.

Compute similarity weights. For each test vehicle, we weight the training set by similarity of TabPFN’s predicted MPG as in (3). The effective sample sizes $\hat{n}_{\text{eff}}$ range from 39 to 172 (median 141) out of $n_{\text{train}}=235$ .
3.

Fit and predict. We fit a locally distilled model for every test vehicle by solving (4). This produces one model and prediction per car. On this test set, the global lasso has PSE $10.81$ , TabPFN $5.25$ , and local distillation $5.59$ : the student performance is close to that of its teacher. In Section 4, we visualize and interpret the models.

2.3 Computability, generalizations and further localization

Computability.

For a single test observation $\bm{x}^{*}$ , Equation (4) is a weighted elastic-net problem with fixed $\hat{\lambda}$ and $\alpha$ : append $\bm{x}^{*}$ to $\bm{X}$ and $\hat{\phi}(\bm{x}^{*})$ to $\bm{y}$ , and use weights $\{\hat{S}_{1},\dots,\hat{S}_{n},\hat{\mu}/\sqrt{\hat{n}_{\text{eff}}}\}$ . Any solver that supports observation weights (e.g. glmnet [28] or adelie [29]) can fit it directly. Fitting many models scales naturally: $\hat{\lambda}$ is estimated once on the training set, and the local fits are then independent single- $\lambda$ problems, parallelizable across query points.

Other forms of regularization, and nonlinearities in the student.

We have focused on the elastic-net penalty, which spans lasso through ridge, but the pseudo-observation construction above is agnostic to the penalty: any regularizer supported by the solver, such as the group or fused lasso, can be used in its place.

We have also chosen the student to be locally linear, but the feature map is a modeling choice: replacing $\bm{x}$ with a basis expansion $\Phi(\bm{x})$ (e.g. pairwise interactions) yields a student that is linear in $\Phi(\bm{x})$ and fit identically, at the cost of interpreting coefficients in the expanded basis. Natural choices here are glinternet [30], which selects pairwise interactions under a strong-hierarchy constraint, or reluctant interaction modeling [31], which adds interactions only where main effects leave residual signal; either keeps the local model sparse and interpretable while capturing interaction structure.

Generalization to other losses.

We have illustrated our method with squared-error loss, but it extends to any loss $\ell(y,\eta)$ convex in $\eta=\bm{x}^{\top}\bm{\beta}$ : the local objective of Algorithm 1 becomes

\hat{\bm{\beta}}^{*}=\argmin_{\bm{\beta}}\sum_{j=1}^{n}\hat{S}_{j}\,\ell(y_{j},\bm{x}_{j}^{\top}\bm{\beta})+\frac{\hat{\mu}}{\sqrt{\hat{n}_{\text{eff}}}}\,\ell(\hat{\phi}(\bm{x}^{*}),\bm{x}^{*\top}\bm{\beta})+\hat{\lambda}[\alpha\|\bm{\beta}\|_{1}+\tfrac{(1-\alpha)}{2}\|\bm{\beta}\|_{2}^{2}],

(5)

where $\hat{\mu}$ is the cross-validated student-to-teacher loss ratio with $\ell$ in place of squared error. Taking $\ell(y,\eta)=\tfrac{1}{2}(y-\eta)^{2}$ recovers Equation (4).

For logistic regression, the distillation term is the cross-entropy between the teacher’s and student’s predicted probabilities, matching the soft-target distillation of Hinton et al. [24]; similarity distances $d_{j}$ are computed on the logit scale.

A survival response requires one change: the loss is the weighted Cox partial likelihood, and since the teacher has no natural place in its risk-set structure, it enters instead as a squared penalty $(\hat{\phi}(\bm{x}^{*})-\bm{x}^{*\top}\bm{\beta})^{2}$ on the log-relative-hazard scale.

Further localization.

For each test observation, one could select $\lambda$ by local cross-validation and estimate $\mu$ from locally weighted losses. We instead reuse the global $\hat{\lambda}$ : the local loss in (4) is a weighted mean on the same scale as the cross-validation loss used to select $\hat{\lambda}$ , so the same value is a sensible default, and refitting locally is expensive and in our experiments rarely improved prediction. Local estimates of $\mu$ were highly variable and did not reliably improve performance either.

Teacher selection.

Given several candidate teachers, we propose selecting the one minimizing the cross-validated loss $\hat{L}(\hat{\phi})$ in (2); this adds little cost beyond Step 1 of Algorithm 1, and in our benchmarks, it worked well (Appendix B).

Cross-modal and cross-domain distillation.

The teacher enters Algorithm 1 only through its predictions: out-of-fold predictions on the training data, and the prediction at $\bm{x}^{*}$ . There is therefore no requirement that it use the same features as the student. A teacher built on a richer modality (e.g. images, text, or additional clinical measurements) can be distilled into a local linear model on the features we wish to interpret; conversely, the teacher may perform best with fewer features than the student, as in our gene expression example (Section 6), where the teacher screens to $500$ genes while the student is fit over all $17{,}322$ . Nor is the teacher required to be trained on the data at hand. Most of our examples use TabPFN, a foundation model pretrained on synthetic data; a model fit to an external cohort is analogous. In either case, $\hat{\mu}$ guards against domain shift: if the teacher does not transfer well, $\hat{\mu}\leq 1$ and the method reverts to the global fit.

3 Related work

Local linear modeling.

Local modeling methods fit a separate model for each test observation rather than a single global model. The classical example is LOESS [16], which at each test observation fits a linear or quadratic regression weighted by kernel proximity in input space. LOESS produces smooth, adaptive predictions without committing to a global functional form, and its local coefficients are interpretable. Its weighting is defined by a user-specified kernel with a bandwidth tuning parameter, and it does not use a teacher.

Closer to ours are methods that fit a local linear model with supervised weights derived from a forest. Generalized random forests [20] and local linear forests [19] use a forest’s leaf co-membership, and the attention lasso [32] similarly uses random forest proximity, then blends each local fit with a global baseline via a mixing parameter tuned by cross-validation. Local distillation instead defines locality through similarity of predicted response, which allows the use of any accurate regressor as teacher. The teacher’s influence is estimated from the cross-validated loss ratio rather than tuned, and the method reverts to the global fit when the teacher offers no improvement. Moreover, the teacher enters as a pseudo-observation within a single fit, rather than blending with a second model: each local model is then an elastic-net fit on a weighted, augmented dataset with a single active set. When the local models use lasso regularization (elastic-net $\alpha=1$ ), the stability guarantees of Section 5 apply to their feature-selection probabilities.

Many recent local linear methods are designed to produce local explanations of black-box predictors; they differ in what they fit and in how the black box enters. The most widely used, LIME [8], fits to the black box: a sparse linear model at each test observation is fit to proximity-weighted perturbations of the input, with the teacher’s predictions as the response, so LIME approximates the teacher rather than the data; like SHAP [9], its explanations are sensitive to choices external to the model and data [10, 11, 12, 13]. MAPLE [33] is more flexible in its response: a random forest fit to $y$ supplies global feature selection and a proximity kernel, which together weight a local regression at each test observation. This local model can regress on either the observed labels $y$ or a black-box teacher’s predictions. As a predictor, MAPLE is closely related to local linear forests, which we include in our benchmarks as a representative of this family (Section 7).

Local distillation shares this per-test-point linear structure and fits the observed response $y$ ; the teacher defines locality and anchors the fit rather than serving as the target. It is a predictive method in its own right, and in contrast to the local methods above, it comes with stability guarantees for feature selection (Section 5).

Prediction-powered inference.

Prediction-powered inference (PPI, PPI++) [34, 35] augments classical statistical analyses with predictions from a powerful machine learning (ML) model when labeled data are scarce. Given a small labeled dataset and a large unlabeled dataset, it imputes the missing labels and builds confidence intervals for population-level estimands that account for the model’s imputation error and recover the classical estimator when the ML model adds nothing. PPI and local distillation share a philosophy: both use a powerful ML model to strengthen a simpler statistical procedure, and revert to the simpler procedure—in our case, the well-understood global linear model—when the ML model is unreliable. They differ in target and mechanism. PPI targets population-level inference and corrects for prediction error in its confidence intervals; local distillation targets pointwise prediction and guards against an unhelpful teacher through the data-driven $\hat{\mu}$ and reversion to the global fit. The power-tuning parameter $\lambda_{\text{PPI}}\in[0,1]$ of PPI++ (chosen from data to minimize variance) mirrors the role of $\hat{\mu}$ , though $\hat{\mu}$ enters as the precision of a pseudo-observation rather than a mixing weight.

4 Interpretability through stability

Local distillation, as described in Algorithm 1, predicts each query point through a sparse local linear fit anchored at that point. Building on the concerns raised in Section 1, conclusions drawn from predictions may inspire little trust unless they can be shown to be stable, motivating the development of a stability-based framework for interpretation. To this end, we propose randomized local distillation, a simple modification of Step 3 in Algorithm 1. The resulting framework assesses stability by examining how the local coefficients vary under random perturbations of the local distillation optimization, while holding the query point $\bm{x}^{*}$ and features $\bm{X}$ fixed. The stability analysis yields interpretations of predictions at two levels, as demonstrated in this section: (i) individually, through the selected features in each sparse local model, and (ii) in aggregate, by characterizing heterogeneity across query points to identify similar and dissimilar regions of the dataset.

4.1 Randomized local distillation

Let $\bm{w}=(w_{1},\dots,w_{n+1})$ denote a vector of $(n+1)$ independent and identically distributed randomization variables, with $w_{j}\sim\mathcal{N}(0,\tau^{2})$ for $j=1,\dots,n+1$ . Rather than relying on a single local fit, we generate repeated randomized local fits by solving

	$\displaystyle\hat{\bm{\beta}}^{\bm{w}}=\argmin_{\bm{\beta}}\;\frac{1}{2}\sum_{j=1}^{n}$	$\displaystyle\hat{S}_{j}\bigl(y_{j}-\bm{x}_{j}^{\top}\bm{\beta}\bigr)^{2}+\frac{\hat{\mu}}{2(\hat{n}_{\text{eff}})^{1/2}}\bigl(\hat{\phi}(\bm{x}^{})-\bm{x}^{\top}\bm{\beta}\bigr)^{2}+\hat{\lambda}\bigl[\alpha\\|\bm{\beta}\\|_{1}+\frac{(1-\alpha)}{2}\\|\bm{\beta}\\|_{2}^{2}\bigr]$		(6)
		$\displaystyle-\left(\sum_{j=1}^{n}\hat{S}_{j}^{1/2}w_{j}\bm{x}_{j}+\frac{\hat{\mu}^{1/2}}{n_{\text{eff}}^{1/4}}w_{n+1}\bm{x}^{*}\right)^{\top}\bm{\beta},$		(6)

where $\hat{\bm{\beta}}^{\bm{w}}$ denotes the coefficient vector obtained at $\bm{x}^{*}$ for a given randomization draw $\bm{w}$ .

Randomized local distillation modifies (4) of Algorithm 1 by introducing a linear perturbation through normally distributed noise variables, while leaving the loss and penalty terms unchanged. Thus, as with local distillation, the randomized fit naturally extends to other loss functions and regularizers. Analogous to the weighting of the summands in the loss, each randomization variable $w_{j}$ in the linear perturbation term is scaled by the square root of its corresponding weight, $\hat{S}_{j}$ . Consequently, if $\hat{S}_{j}$ is zero or close to zero, $w_{j}$ has no or little effect on the optimization.

The amount of randomization in each randomized local fit is controlled by $\tau^{2}$ , the variance of the normal randomization variables in $\bm{w}$ . In practice, we set $\tau^{2}=t\;*\dfrac{\hat{\sigma}^{2}}{\hat{n}_{\text{eff}}}$ , where $\hat{n}_{\text{eff}}=\frac{1}{\hat{S}_{1}^{2}+\ldots+\hat{S}_{n}^{2}}$ is the effective sample size (Equation (3)), $\hat{\sigma}^{2}$ is the variance of the teacher’s out-of-fold residuals, and $t$ is a constant chosen by the analyst. Choosing the randomization variance as a fixed fraction of $\dfrac{\hat{\sigma}^{2}}{\hat{n}_{\text{eff}}}$ ensures that the scale of the randomization is comparable to the scale of variability contributed by the data to the optimization objective. We recommend using data to guide the choice of $t$ . The sensitivity bound established in Theorem 1, from which the stability bound in Corollary 1 follows, improves as the randomization sd $\tau$ increases, and thus as $t$ increases. But too much randomization can reduce the predictive accuracy of the randomized fits. Therefore, we take the largest $t$ whose out-of-fold median prediction error is within a small tolerance of that from the unperturbed fits ( $t=0$ ); the tolerance reflects how much prediction accuracy the analyst is willing to trade for stability, and we use 5% throughout. Figure 2 (left) shows the selection process for the Auto MPG data.

Remark 1.

The randomization in (6) is similar in form to that used in the post-selection inference literature, where linear perturbation terms are added to penalized M-estimation problems to preserve information that can later be leveraged for valid inference in selected models. See, for example, the randomized inference methods developed in [36, 37, 38, 39, 40]. Here, however, the added randomization serves a different role: in Section 5, we prove that it provides stability guarantees for feature–selection probabilities under small perturbations of the training data.

4.2 Interpreting a single local model

At a given query point $\bm{x}^{*}$ , the local model characterizes the feature–response relationship among observations with similar predicted outcomes. More specifically, a single local distillation fit yields a sparse set of features, $\{j:\hat{\beta}_{j}\neq 0\}$ (for $\alpha>0$ ), providing an interpretation of its prediction at $\bm{x}^{*}$ . But, as with the standard lasso, this selected set may be unstable: predictors near the selection boundary may enter or leave the model under even slight perturbations in the outcomes, ultimately yielding different local interpretations at $\bm{x}^{*}$ .

To address this instability, we compute for each feature its empirical selection frequency across the randomized local fits

\hat{\pi}_{j}=\frac{1}{B}\sum_{b=1}^{B}\mathbb{I}\left\{\hat{\beta}^{\bm{w}_{b}}_{j}\neq 0\right\},

where $B$ is the number of refits of (6) and $\bm{w}_{b}$ denotes the randomization used in the $b$ -th fit. That is, $\hat{\pi}_{j}$ is the proportion of randomized local refits in which that feature is selected. In particular, a low selection frequency signals uncertainty about a feature’s contribution to the prediction at $\bm{x}^{*}$ .

The table in Figure 2 shows one such model for a single test car. Alongside each coefficient we report its selection frequency, the fraction of randomized local refits retaining that feature. In this example, five of the seven selected features are retained in over 90% of refits; the least stable are cylinders and horsepower, retained in 75% and 70% respectively. Randomization flags these as uncertain.

Feature	Coef.	Sel. prob.
weight	$-5.11$	100%
model year	$2.28$	100%
origin: European	$0.40$	99%
acceleration	$0.52$	99%
origin: Japanese	$0.41$	98%
cylinders	$\bm{0.24}$	75%
horsepower	$\bm{-0.17}$	70%

Figure 2: Left: choosing the randomization scale

t

. We take the largest

t

whose out-of-fold median squared error remains within 5% of the unperturbed fit (shaded); the selected value is circled. Right: a local model for a single test car with mid-range fuel economy (distillation predicted

23.7

mpg, observed

23.9

mpg), ordered by decreasing selection probability. The lasso penalty selected 7 of the 8 features in the unperturbed fit. Selection probability is the fraction of 100 randomized refits retaining the feature at the selected scale

\hat{t}

(circled at left), and features retained in under 90% are shown in bold. Region coefficients are contrasts against American manufacture.

Here, with only $8$ features, most selections are stable across refits. This is not the case in higher dimensions: in the gene expression example of Section 6 ( $p=17{,}322$ ), the median local fit selects $94$ genes, of which typically only $15$ are retained in over $90\%$ of refits. In simulation (Appendix D), we find that the selection frequencies distinguish the true support while yielding a smaller, more interpretable set of features.

Using randomized refits to address the instability of lasso-selected features is not new; a prominent example is stability selection [41] based on subsampling. The randomization mechanism in our approach, however, is deliberately different from subsampling. This distinction has consequences for both the computational cost of local refitting and the interpretation of the resulting stability measures, which we discuss below.

Remark 2.

Across randomized refits of local distillation, as proposed in (6), the similarity weights and teacher predictions are computed once from the training data and then held fixed. The added randomization perturbs only the optimization objective, while keeping the design fixed and using all $n$ observations in each refit. From a computational perspective, this avoids the potentially substantial additional cost of subsampling-based refitting procedures such as stability selection. Under subsampling, Steps 1 and 2 of Algorithm 1 would need to be recomputed for each refit, as would the teacher predictions whenever they depend on the sampled training data; take, for example, predictions from tabular foundation models deploying in-context learning.

Remark 3.

The choice of randomization determines the notion of stability being assessed and, consequently, the interpretation of the resulting predictions. Beyond its computational advantages, our randomization mechanism is designed to remain faithful to the locality learned from the training data: each refit preserves the same similarity weights and the teacher prediction anchor. By contrast, under subsampling-based refits, even small changes in sample composition can potentially alter the local structure in data, especially when the effective sample size $\hat{n}_{\text{eff}}$ is small, i.e., when the weights $\hat{S}_{j}$ are concentrated on only a few observations, or when outcomes have heterogeneous noise levels. In such settings, variation across refits may reflect changes in the local structure rather than instability in the local feature–response relationship itself, making stability a less reliable characterization of that relationship.

4.3 Interpreting the local models in aggregate

Beyond individual models, we are also interested in understanding our sample and the heterogeneity within it. A natural approach would be to fit locally distilled models on a test set with $n_{\text{test}}$ observations (or in a leave-one-out setting on the training data), and then cluster the fitted coefficient vectors into $k$ clusters. However, our goal is to identify and interpret subgroups that are stable under small perturbations of the response $\bm{y}$ rather than artifacts of a single fit. We therefore aggregate the clustering obtained across the randomized refits from Section 4.1 using evidence accumulation clustering [42], allowing the same refits used to characterize local feature–response relationships to also inform stable subgroup structure in the data. We proceed as follows:

1.

Fit unperturbed local models: run local distillation to obtain $n_{\text{test}}$ fitted local models, with coefficient vectors in $\mathbb{R}^{p}$ .
2.

Choose the number of clusters: cluster the $n_{\text{test}}$ unperturbed coefficient vectors and select the appropriate $k$ .
3.

Cluster the randomized refits: draw $B$ independent randomization vectors $\bm{w}_{1},\ldots,\bm{w}_{B}$ , as described in Section 4.1. For each draw $\bm{w}_{b}$ , run randomized local distillation (solve (6)) to obtain $n_{\text{test}}$ randomized coefficient vectors, and cluster these into $k$ groups, yielding $B$ cluster assignments of the test observations.
4.

Cluster the co-occurrence matrix: define the similarity between two test observations as the fraction of the $B$ clusterings in which they fall in the same group, and use this metric for a final clustering into $k$ groups.

We illustrate this using the Auto MPG data. We (1) fit local models for the ${n_{\text{test}}=157}$ observations and (2) select the number of clusters $\mathbf{k}$ using $k$ -means clustering of the unperturbed local coefficient vectors; the silhouette score selects $k=3$ . We then (3) cluster the randomized refits: for each of $B=100$ randomized local distillation runs, we cluster the refitted coefficient vectors using $k$ -means with $k=3$ . Finally, we (4) cluster the co-occurrence matrix: the similarity between two cars is the fraction of the $100$ runs in which they are assigned to the same cluster, and clustering this matrix by average linkage yields the three groups shown in Figure 3. To visualize the result, we display the unperturbed coefficients in a heatmap, either ordered by cluster, or averaged within cluster. The three clusters use different features: cylinders is predictive among the least efficient cars, and displacement among the most efficient.

5 Theoretical analysis of stability

In this section, we establish stability guarantees for feature–selection probabilities under randomized local distillation (6), whose local coefficients enable interpretation of the distilled fits both individually and in aggregate. Throughout this section, we treat the similarity weights $\hat{S}_{j}$ and the regularization parameters $\hat{\mu}$ and $\hat{\lambda}$ as fixed. To avoid confusion, we drop the hats and write these quantities as $S_{j}$ , $\mu$ , and $\lambda$ , respectively; similarly we write $n_{\mathrm{eff}}$ for the effective sample size. We slightly modify our notation for the teacher’s prediction at the query point, denoting it by $\hat{\phi}(\bm{x}^{*};\bm{y})$ , to make explicit that it may depend on the training response, as is the case in tabular foundation models, for example. We focus on the lasso penalty, i.e., $\alpha=1$ , while deferring an extension to the elastic-net penalty for future work to keep the theoretical development streamlined.

Our main result, Theorem 1, characterizes how the feature–selection probabilities are sensitive to changes in the input responses, leading to the uniform stability guarantee in Corollary 1 under small perturbations of the training response. To prove this result, we first establish a stability guarantee for randomized lasso regression of a perturbed response on the design matrix; see Theorem 2 in Appendix A.2. To the best of our knowledge, this is the first such guarantee for Gaussian randomization introduced through linear perturbations of the optimization objective, providing a theoretical basis for the stability-based interpretation of feature–selection probabilities under the lasso penalty.

5.1 Notation and preliminaries

Before proceeding, we introduce notation used to develop the theory.

First, let

\bm{\overline{y}}=\begin{pmatrix}y_{1}\\ \vdots\\ y_{n}\\ \hat{\phi}(\bm{x}^{*};\bm{y})\end{pmatrix}\in\mathbb{R}^{n+1},\qquad\bm{\overline{X}}=\begin{pmatrix}\bm{x}_{1}^{\top}\\ \vdots\\ \bm{x}_{n}^{\top}\\ (\bm{x}^{\ast})^{\top}\end{pmatrix}\in\mathbb{R}^{(n+1)\times p},

denote the augmented response vector and design matrix, obtained by appending the teacher’s prediction $\hat{\phi}(\bm{x}^{*};\bm{y})$ and the query point $\bm{x}^{*}$ to the response vector and design matrix, respectively.

Let

\bm{\Omega}\;=\;\operatorname{diag}\!\left(S_{1},\dots,S_{n},S_{n+1}\right)\in\mathbb{R}^{(n+1)\times(n+1)},

denote the diagonal matrix of weights for the $n+1$ observations in the augmented dataset, where $S_{n+1}=\frac{\mu}{\sqrt{n_{\mathrm{eff}}}}$ is the weight assigned to the pseudo-response at the query point. We then define

\bm{z}=\bm{z}(y_{1},\ldots,y_{n})\;=\;\bm{\Omega}^{1/2}\,\bm{\overline{y}}\in\mathbb{R}^{n+1},\qquad\bm{V}\;=\;\bm{\Omega}^{1/2}\,\bm{\overline{X}}\in\mathbb{R}^{(n+1)\times p},

(7)

to be the weight-adjusted response and design, respectively, i.e., coordinatewise, $Z_{k}=\sqrt{S_{k}}\,y_{k}\quad(k\leq n)$ , and $Z_{n+1}=\sqrt{S_{n+1}}\,\hat{\phi}(\bm{x}^{*};\bm{y})=\frac{\sqrt{\mu}}{n_{\mathrm{eff}}^{1/4}}\,\hat{\phi}(\bm{x}^{*};\bm{y})$ .

Under this notation, when $\alpha=1$ and $\lambda$ denotes the tuning parameter for the lasso penalty, the randomized local distillation problem (6) is equivalent to solving

\bm{\widehat{\beta}}^{\,\bm{\omega}}\;=\;\operatorname*{argmin}_{\bm{\beta}\in\mathbb{R}^{p}}\;\tfrac{1}{2}\,\lVert\bm{z}+\bm{\omega}-\bm{V}\bm{\beta}\rVert_{2}^{2}\;+\;\lambda\lVert\bm{\beta}\rVert_{1},

(8)

where $\bm{\omega}\sim\mathcal{N}(0,\tau^{2}\bm{I}_{n+1})$ is drawn independently of the training response. The equivalent optimization problem, which follows from straightforward algebra, is used throughout this section. An advantage of this formulation is that it allows us to first establish guarantees for randomized lasso regression of a perturbed response on a design matrix. These results then yield the desired guarantees for local distillation, while also being of potential independent interest beyond the present setting.

Let $\widehat{E}_{V}(\bm{z},\bm{\omega})=\operatorname{supp}\!\left(\widehat{\bm{\beta}}^{\bm{\omega}}(\bm{z})\right)$ denote the active set obtained by solving the randomized lasso in (8). The quantity of interest is the feature-selection probability for predictor $j$ , evaluated as a function of $\bm{y}$ and denoted by

\pi_{j}(\bm{y})\;=\;\mathbb{P}_{\bm{\omega}}\!\left[\,j\in\widehat{E}_{V}(\bm{z},\bm{\omega})\right],\quad\text{ for }j\in[p],

(9)

where the probability is over the randomization $\bm{\omega}$ , with $\bm{y}$ and hence $\bm{z}=\bm{z}(y_{1},\ldots,y_{n})$ treated as fixed, and $[p]=\{1,2,\ldots,p\}$ . The subscript $\bm{\omega}$ on the probability, here and throughout, indicates that the probability is taken only with respect to the randomization.

5.2 Stability of randomized local distillation

Because the selection probabilities in (9) are obtained by convolving the discontinuous feature–selection indicator with the Gaussian density of $\bm{w}$ , they are smooth provided that the teacher’s prediction is a smooth function of the input $\bm{y}$ . Lemma 1 formalizes this observation before we turn to the stability guarantee for randomized local distillation; the proof is provided in Appendix A.1.

Lemma 1 (Smoothness of feature–selection probabilities).

Fix $k\in\mathbb{N}\cup\{\infty\}$ and suppose that $\bm{y}\mapsto\hat{\phi}(\bm{x}^{*};\bm{y})$ belongs to $\mathcal{C}^{k}(\mathbb{R}^{n})$ . Then for every $j\in[p]$ , the selection probability $\pi_{j}(\bm{y})$ , defined in (9), also belongs to $\mathcal{C}^{k}(\mathbb{R}^{n})$ .

We now state Theorem 1, which bounds the sensitivity of the feature–selection probabilities to perturbations of the training response $\bm{y}$ , leveraging the smoothness of these probabilities established in Lemma 1. We derive this result under the following assumptions on the design matrix $\bm{X}$ , the similarity weights $S_{j}$ and the weight-adjusted design matrix $\bm{V}$ .

Assumption 1 (Weight-adjusted design in general position).

The columns of $\bm{V}$ are in general position [43].

Assumption 2 (Bounded weight concentration).

We assume that there exists a constant $C_{S}<\infty$ such that $n_{\mathrm{eff}}S_{\max}\leq C_{S}$ .

Assumption 3 (Bounded design and local distillation weight).

For some constants $B_{X},C_{\mu}>0$ , we assume that $\max\left\{\max_{i\in[n]}\|\bm{x}_{i}\|_{\infty},\|\bm{x}^{*}\|_{\infty}\right\}\leq B_{X}$ , and $0<\mu\leq C_{\mu}$ .

Assumption 4 (Non-degeneracy of restricted design).

For some constant $\kappa_{X}>0$ , we assume that, uniformly over $j\in[p]$ and $E\in\mathcal{E}_{-j}^{V}$ , where $\mathcal{E}_{-j}^{V}$ denotes the collection of essential active sets from lasso regression on $\bm{V}_{-j}$ (i.e., active sets from the leave- $j$ -covariate-out lasso fit whose corresponding selection regions have positive Lebesgue measure),

\sigma_{\min}\left(\frac{1}{\sqrt{|\mathcal{I}_{S}|}}\,\bm{X}_{\mathcal{I}_{S},E\cup\{j\}}\right)\geq\kappa_{X},

where $\mathcal{I}_{S}=\left\{i\in[n]:S_{i}\geq\frac{1}{2n_{\mathrm{eff}}}\right\}$ collects the training observations whose similarity weights are at least one half of the effective uniform weight.

Under Assumption 1, the randomized local distillation problem (8) admits a unique solution. Assumption 2 prevents the similarity weights from concentrating too heavily on a small number of observations. Assumption 3 requires both the design covariates and the distillation-weight parameter to be bounded. Assumption 4 imposes a lower singular-value condition on the design restricted to the effective similarity neighborhood $\mathcal{I}_{S}$ , preventing the relevant design columns from becoming nearly linearly dependent among observations receiving non-negligible weights.

Theorem 1 (Sensitivity bound for selection probabilities).

Suppose that $\bm{y}\mapsto\hat{\phi}(\bm{x}^{*};\bm{y})$ belongs to $\mathcal{C}^{1}(\mathbb{R}^{n})$ with $L_{\hat{\phi}}(\bm{x}^{*})=\sup_{\bm{y}\in\mathbb{R}^{n}}\|\nabla\hat{\phi}(\bm{x}^{*};\bm{y})\|_{\infty}<\infty$ , and that Assumptions 1, 2, 3 and 4 hold. Then there exists a constant $C<\infty$ such that, for every $j\in[p]$ ,

\sup_{\bm{y}\in\mathbb{R}^{n}}\|\nabla\pi_{j}(\bm{y})\|_{\infty}\;\leq\;\sqrt{\frac{2}{\pi}}\frac{C}{\tau}\left(\frac{1}{\sqrt{n_{\mathrm{eff}}}}\;+\;\frac{L_{\hat{\phi}}(\bm{x}^{*})}{n_{\mathrm{eff}}^{1/4}}\right).

A proof of Theorem 1 is provided in Appendix A.1. The proof consists of two main components. First, we analyze the stability of the randomized lasso obtained by regressing a perturbed response on a $p$ -dimensional feature matrix, as detailed in Appendix A.2. Second, we apply the chain rule to characterize the additional contribution from local distillation, as detailed in Appendix A.3.

We make a few observations on the above-stated result.

(i)

It follows directly from the bound in Theorem 1 that the contribution of each training observation to the sensitivity of a feature’s selection probability has two components: the first is a direct contribution, excluding the effect of distillation, which is of order $n_{\mathrm{eff}}^{-1/2}$ , and the second contribution arises through the teacher’s prediction, which is of order $n_{\mathrm{eff}}^{-1/4}L_{\hat{\phi}}(\bm{x}^{*})$ .
(ii)

Although we make no assumption on $L_{\hat{\phi}}(\bm{x}^{*})$ , which reflects how the sensitivity of the teacher’s prediction scales with the effective sample size, it may decrease as the effective sample size grows, leading to a sharper rate for the second contribution in the bound.
(iii)

When the teacher is a fully pretrained model (i.e., the teacher’s prediction does not depend on the training samples), the bound in Theorem 1 simplifies to

$\sup_{\bm{y}\in\mathbb{R}^{n}}\|\nabla\pi_{j}(\bm{y})\|_{\infty}\;\leq\;\sqrt{\frac{2}{\pi}}\frac{C}{\tau}n_{\mathrm{eff}}^{-1/2},$

since $L_{\hat{\phi}}(\bm{x}^{*})=0$ . As a further special case, for the global lasso fit without using the teacher as a Bayesian prior to shrink the linear predictor toward its predictions, and with uniform similarity weights $S_{j}=1/n$ , the same bound holds with $n_{\mathrm{eff}}=n$ .

Theorem 1 immediately yields the following corollary, which quantifies the stability of the feature–selection probabilities in response to a perturbation of a single response value, paralleling notions of algorithmic stability under single-observation perturbations [44].

Corollary 1 (Stability guarantee).

Under the assumptions of Theorem 1, there exists a constant $C<\infty$ such that, for every $j\in[p]$ , $i\in[n]$ , and for any $\bm{y},\bm{y}^{\prime}\in\mathbb{R}^{n}$ satisfying $y_{k}=y_{k}^{\prime}$ for all $k\neq i$ (equivalently, differing only in the $i$ -th coordinate),

\bigl\lvert\pi_{j}(\bm{y})-\pi_{j}(\bm{y}^{\prime})\bigr\rvert\;\leq\sqrt{\frac{2}{\pi}}\frac{C}{\tau}\left(\frac{1}{\sqrt{n_{\mathrm{eff}}}}\;+\;\frac{L_{\hat{\phi}}(\bm{x}^{*})}{n_{\mathrm{eff}}^{1/4}}\right)|y_{i}-y^{\prime}_{i}|.

The proof of Corollary 1 follows directly from the sensitivity bound in Theorem 1 and is therefore omitted.

The stability bound presented in Corollary 1 improves with the randomization standard deviation $\tau$ . However, this does not imply that arbitrarily increasing the amount of randomization is desirable, since doing so can compromise predictive accuracy. Accordingly, as described in Section 4.1, we choose $\tau$ to balance stability against a user-specified tolerance for loss in predictive accuracy.

Extension to stability over a regularization-grid.

Our stability-based approach could be extended to base decisions on a grid $\Lambda$ of lasso regularization parameters $\lambda$ , rather than on a single fixed value of $\lambda$ . As proposed in the stability selection framework of [41], a natural quantity on which to base decisions is $\displaystyle\max_{\lambda\in\Lambda}\pi_{j}(\bm{y};\lambda)$ , where $\pi_{j}(\bm{y};\lambda)$ denotes the selection probability at $\lambda$ (with the dependence on $\lambda$ made explicit in the notation), and the maximum taken over the grid of regularization parameters. However, our proof techniques for controlling the sensitivity of this quantity to perturbations of the training input, and hence for establishing stability guarantees, may not extend directly to the maximum operation because it is not smooth.

Instead, one can consider a smooth approximation to $\displaystyle\max_{\lambda\in\Lambda}\pi_{j}(\bm{y};\lambda)$ , namely $\Pi_{\Lambda}(\bm{y})=\frac{1}{\eta}\log\left(\frac{1}{|\Lambda|}\sum_{\lambda\in\Lambda}\exp\left(\eta\pi_{j}(\bm{y};\lambda)\right)\right)$ , for fixed $\eta\in\mathbb{R}^{+}$ . Like the selection probability for any fixed $\lambda$ , this quantity also takes values in $[0,1]$ . In practice, using its empirical counterpart $\widehat{\Pi}_{\Lambda}(\bm{y})$ , one can then determine the set of selected features as $\{j:\widehat{\Pi}_{\Lambda}(\bm{y})\geq p_{\text{thr}}\}$ , for a chosen threshold $p_{\text{thr}}$ . Furthermore, the gradient of $\Pi_{\Lambda}(\bm{y})$ can be written as

\nabla\Pi_{\Lambda}(\bm{y})=\sum_{\lambda\in\Lambda}w_{\lambda}(\bm{y})\,\nabla\pi_{j}(\bm{y};\lambda),\quad\text{with }\ w_{\lambda}(\bm{y})=\dfrac{\exp\!\left(\eta\pi_{j}(\bm{y};\lambda)\right)}{\sum_{\lambda^{\prime}\in\Lambda}\exp\!\left(\eta\pi_{j}(\bm{y};\lambda^{\prime})\right)},

a convex combination of the derivatives of the selection probabilities. Consequently, a fairly straightforward extension of our results to establish a uniform stability bound over the grid $\Lambda$ would yield a corresponding stability guarantee for $\Pi_{\Lambda}(\bm{y})$ under perturbations to a single coordinate of the input response, analogous to the bound in Corollary 1. We omit the details of this extension from the present work.

6 A high-dimensional example: predicting gene expression

Here, we turn to a breast cancer gene expression study from The Cancer Genome Atlas (TCGA) [45], distributed by Breheny and Huang [46]. Tumor samples from $n=536$ patients were assayed on Agilent mRNA expression microarrays, and measurements are on the log scale. Following the example in Breheny and Huang [46], we treat BRCA1 expression as the response and the remaining $p=17{,}322$ genes as predictors, excluding 491 genes with missing data. BRCA1 is the first gene identified whose mutations increase the risk of early onset breast cancer, and because it is likely to interact with many others, those whose expression is related to BRCA1 are candidates for further study.

We divide the data into a 60/40 train/test split, with $n_{\text{train}}=321$ and $n_{\text{test}}=215$ . The student model is lasso regression. For the teacher, we considered two candidates and chose between them using cross-validated error as described in Section 2.3: TabPFN applied to all $17{,}322$ genes, and TabPFN applied to the $500$ genes most correlated with BRCA1 expression, screened within each training fold. The screened teacher performed far better and was selected. (The global lasso did not improve when restricted to the $500$ gene subset.) Note that in this setting, the teacher and student have different feature representations, as in cross-modal distillation: the teacher uses $500$ genes while the student uses all $17{,}322$ .

On this split, the global lasso had test PSE $0.189$ ( $R^{2}=0.579$ ), TabPFN $0.150$ ( $R^{2}=0.666$ ), and local distillation $0.148$ ( $R^{2}=0.670$ ): a 22% reduction relative to the global lasso, matching (here, slightly exceeding) its teacher, while retaining transparency. Additionally, the local fits were sparser than the global lasso fit: the global fit selected $123$ genes, while the median local fit selected $94$ .²² 2 The PSE improvement was robust across 100 random 60/40 train/test splits, where local distillation improved on the global lasso in 96% of runs, with a median PSE reduction of 22%. The sparsity comparison was also stable: the median local fit was sparser than the global fit on 97 of 100 splits, with a median reduction of 15% in the number of selected genes. Applying the stability screen of Section 4 ( $\hat{\pi}_{j}>0.9$ across $100$ randomized refits) reduces this further: the local models retained a median of $15$ stably selected genes (IQR $13$ – $17$ ).

We then clustered the local models using 100 randomized fits with randomization scale $t=0.1$ and $k=5$ clusters, with all parameters selected as described and exemplified in Section 4.3; the resulting clusters are visualized in Figure 4. The clusters reveal heterogeneity across patients: FAM107A is selected almost only in cluster one, KLF14 almost only in cluster five, both with negative coefficients, and both given a coefficient of zero by the global lasso. These are reported tumor suppressors down-regulated in cancer [47, 48]. This heterogeneity is not visible to a global linear model, and is difficult to surface in black-box models.

7 Benchmark comparisons

Here, we use common machine learning benchmark datasets to compare local distillation to (1) the student model (global lasso or ridge regression), (2) the teacher model (TabPFN or XGBoost [49]) and (3) two local linear models (LOESS and local linear forests). We evaluate on 17 regression datasets spanning sample sizes $n\in[159,4177]$ and feature counts $p\in[5,51]$ . The datasets are from the UCI Machine Learning Repository [50] (automobile, servo, liver disorders, auto MPG, real estate valuation, infrared thermography temperature, student performance) and the OpenML-CTR23 regression benchmark [51, 52] (cars, QSAR fish toxicity, concrete compressive strength, socmob, airfoil self-noise, red wine, auction verification, space ga, abalone, white wine).

For each dataset, we generate 20 random 80/20 train/test splits. We use a complete-case design and perform one-hot encoding for categorical variables. Within each train/test split, we normalize features using the mean and standard deviation of the training set. Per-dataset sample size and feature counts are given in Appendix Table 1.

We compare methods using test $R^{2}$ , and we find that local distillation closely matches the predictive performance of its teacher across a wide range of datasets. See Figure 5 for representative examples, where local distillation is labeled as “LD (teacher, regularization)”. For teacher models, we use TabPFN and XGBoost, and student regularization comparators are lasso (L) and ridge (R). Appendix B shows complete results.

Figure 5: Predictive performance (test

R^{2}

, median

\pm

1 standard error) for six datasets from the UCI ML and OpenML repositories. Local distillation is labeled as “LD (teacher, regularization)”, using “L” for lasso and “R” for ridge. When the teacher outperforms the global linear model, local distillation usually approaches the teacher’s performance. Plots for the remaining 11 datasets are included in Appendix B.

8 Discussion

Modern black-box models are presenting new opportunities for predictive modeling across domains and data types, often with performance that classical methods cannot easily match without significant feature engineering [14]. However, predictive models are only useful to the extent that decision-makers can draw trustworthy conclusions from them. We posit that well-constructed local linear models are in a “sweet spot” between interpretable modeling and modern AI: they retain the benefits of linear models (transparent, easy to understand, computationally simple) while rivaling the performance of black-box models. The statistical principle underlying this intuition is familiar: a smooth regression surface is locally well approximated by a linear model. But it is nontrivial to determine what constitutes local, how the local models should be fit, and how the resulting collection of fits should be interpreted.

Local distillation, as proposed in this work, relies on a black-box “teacher” for prediction and on randomized refits for interpretation. For prediction, the teacher plays two roles: its predictions (1) define locality by determining which training observations inform the fit at each query point, and (2) provide anchoring pseudo-observations. Across 17 benchmark datasets and a high-dimensional gene expression example, and across a range of teacher models (TabPFN, TabFM, XGBoost), local distillation consistently matches or approaches the predictive accuracy of its teacher.

For interpretation, we apply a small amount of Gaussian randomization to the local distillation optimization, leaving both the loss and penalty unchanged. The randomized refits identify which interpretations are sufficiently stable and therefore reliable, both at an individual test point (through selection frequencies) and across the test dataset as a whole (through clustering into stable subgroups). Under the lasso penalty, we established theoretical guarantees showing that this randomization yields stability under small perturbations of the training responses. These results are of independent interest and extend beyond local distillation, providing a general mechanism for stabilizing feature–selection probabilities under lasso penalization.

Extensions.

Section 2.3 suggests several extensions of local distillation. For example, when a linear student fails to recover the teacher’s accuracy, using a richer student class (with interactions or transformations of the covariates) may narrow the gap. And, when there are many candidate teacher models, the teacher may be selected through cross-validation with the training data. Finally, local distillation can incorporate external datasets or other data modalities through cross-modal distillation, in which the teacher is built on a different feature set than the student. Our gene expression example provides one instance of this approach, and we view distillation across genuinely different modalities as a promising avenue.

The distillation strength $\hat{\mu}$ is estimated using a ratio of losses; this rule is intuitive and it performed well in our experiments, but we have not made any claims about optimality. We additionally use a hard cutoff to decide when to revert to the global linear model (when $\hat{\mu}\leq 1$ , the estimated teacher error is worse than that of the student). We considered a continuous alternative, weighting the teacher by its excess performance $(\hat{\mu}-1)_{+}$ , but found that this reduced predictive performance: when the teacher is stronger than the student, the local fits benefit from a strong teacher weight. A similar observation was made in the distillation paper from Hinton et al. [24], where they found the best results placed most of the weight on the teacher’s soft targets rather than the true labels. Alternative approaches to estimating $\hat{\mu}$ and determining when to revert to a simpler model may nevertheless be worth exploring.

Local regression more generally.

We view local linear modeling as a general and flexible framework that can compete with modern predictive methods, built from modular components that can be chosen to suit the problem: the weights, which define locality (kernels in LOESS, forest proximities, or, here, similarity of the teacher’s predictions); the penalty, which defines structure (e.g., lasso, elastic-net, group lasso); and pseudo-observations, which carry external information to anchor the fit (elicited priors, or, here, the teacher’s prediction at the query point). The choice of weights deserves particular care, because the definition of “local” determines which heterogeneity the local models can express. This is analogous to unsupervised clustering, where many clusterings may be equally valid though not all are equally informative; we expect different definitions of locality to likewise have different virtues. Interpretation methods for local regression also require careful consideration: conclusions drawn from the local fits are only reliable when they are stable, and the appropriate notion of stability depends on the choice of locality. The stability theory under Gaussian randomization is agnostic to the specific choice of weights in our approach, allowing the guarantees to extend beyond our construction of local linear models. Their specific form, however, may offer additional structure that can be exploited for deriving different theoretical guarantees, which we leave for future investigation.

Local distillation is one instantiation of local linear modeling: a black-box teacher defines locality and anchors each fit, and randomization assesses stability. Our results suggest that this framework is a promising path toward predictive, interpretable, and trustworthy statistical modeling.

Acknowledgements. We would like to thank Robert Tibshirani, Trevor Hastie and Shihan Khan for helpful comments. The authors used Claude Opus 4.5 (model ID: claude-opus-4-5-20251101) and ChatGPT 5.0 for coding support and text editing. The gene expression example shown here is based upon data generated by the TCGA Research Network: https://www.cancer.gov/tcga.

Funding. S.P. was supported by NSF CAREER Award DMS-2337882.

Disclosure. The authors report there are no competing interests to declare.

Data availability. The data that support the findings of this study are public and cited throughout. The scripts to download data and run simulations are published on Github at https://github.com/erincr/local-distillation-benchmark. The breast cancer expression data are from The Cancer Genome Atlas (https://www.cancer.gov/tcga); we use the processed version distributed with the hdrm R package (https://github.com/pbreheny/hdrm).

References

[1] J. Jumper, R. Evans, A. Pritzel, T. Green, M. Figurnov, O. Ronneberger, K. Tunyasuvunakool, R. Bates, A. Žídek, A. Potapenko, et al. (2021) Highly accurate protein structure prediction with AlphaFold. Nature 596 (7873), pp. 583–589. Cited by: §1.
[2] R. Lam, A. Sanchez-Gonzalez, M. Willson, P. Wirnsberger, M. Fortunato, F. Alet, S. Ravuri, T. Ewalds, Z. Eaton-Rosen, W. Hu, et al. (2023) Learning skillful medium-range global weather forecasting. Science 382 (6677), pp. 1416–1421. Cited by: §1.
[3] N. Hollmann, S. Müller, L. Purucker, A. Krishnakumar, M. Körfer, S. B. Hoo, R. T. Schirrmeister, and F. Hutter (2025) Accurate predictions on small data with a tabular foundation model. Nature 637 (8045), pp. 319–326. Cited by: §1, §1.
[4] C. Kolberg, J. Kreuer, J. Huurdeman, S. Ouaari, K. Eggensperger, and N. Pfeifer (2025) TabPFN-wide: continued pre-training for extreme feature counts. arXiv preprint arXiv:2510.06162. Cited by: §1.
[5] J. Ma, V. Thomas, R. Hosseinzadeh, A. Labach, J. Cresswell, K. Golestan, G. Yu, A. L. Caterini, and M. Volkovs (2026) TabDPT: scaling tabular foundation models on real data. Advances in Neural Information Processing Systems 38. Cited by: §1.
[6] W. Kong and A. Das (2026) Introducing TabFM: a zero-shot foundation model for tabular data. Note: Google Research BlogAccessed July 7, 2026 External Links: Link Cited by: Appendix B, §1.
[7] J. Qu, D. Holzmüller, G. Varoquaux, and M. L. Morvan (2025) TabICL: a tabular foundation model for in-context learning on large data. arXiv preprint arXiv:2502.05564. Cited by: §1.
[8] M. T. Ribeiro, S. Singh, and C. Guestrin (2016) “Why should I trust you?”: explaining the predictions of any classifier. In Proceedings of the 22nd ACM SIGKDD International Conference on Knowledge Discovery and Data Mining, pp. 1135–1144. Cited by: §1, §3.
[9] S. M. Lundberg and S. Lee (2017) A unified approach to interpreting model predictions. Advances in Neural Information Processing Systems 30. Cited by: §1, §3.
[10] D. Garreau and U. von Luxburg (2020) Explaining the explainer: a first theoretical analysis of LIME. In International Conference on Artificial Intelligence and Statistics, pp. 1287–1296. Cited by: §1, §3.
[11] D. Slack, S. Hilgard, E. Jia, S. Singh, and H. Lakkaraju (2020) Fooling LIME and SHAP: adversarial attacks on post hoc explanation methods. In Proceedings of the AAAI/ACM Conference on AI, Ethics, and Society, pp. 180–186. Cited by: §1, §3.
[12] I. E. Kumar, S. Venkatasubramanian, C. Scheidegger, and S. Friedler (2020) Problems with Shapley-value-based explanations as feature importance measures. In International Conference on Machine Learning, pp. 5491–5500. Cited by: §1, §3.
[13] K. Aas, M. Jullum, and A. Løland (2021) Explaining individual predictions when features are dependent: more accurate approximations to Shapley values. Artificial Intelligence 298, pp. 103502. Cited by: §1, §3.
[14] C. Rudin (2019) Stop explaining black box machine learning models for high stakes decisions and use interpretable models instead. Nature Machine Intelligence 1 (5), pp. 206–215. Cited by: §1, §8.
[15] R. Quinlan (1993) Auto mpg. Note: UCI Machine Learning Repositoryhttps://doi.org/10.24432/C5859H. Dataset from the 1983 American Statistical Association Data Exposition Cited by: §1, §2.2.
[16] W. S. Cleveland and S. J. Devlin (1988) Locally weighted regression: an approach to regression analysis by local fitting. Journal of the American Statistical Association 83 (403), pp. 596–610. Cited by: §1, §3.
[17] R. Qiu, Z. Yu, and R. Zhu (2024) Random forest weighted local Fréchet regression with random objects. Journal of Machine Learning Research 25 (107), pp. 1–69. Cited by: §1.
[18] A. Bloniarz, A. Talwalkar, B. Yu, and C. Wu (2016) Supervised neighborhoods for distributed nonparametric regression. In Artificial Intelligence and Statistics, pp. 1450–1459. Cited by: §1.
[19] R. Friedberg, J. Tibshirani, S. Athey, and S. Wager (2020) Local linear forests. Journal of Computational and Graphical Statistics 30 (2), pp. 503–517. Cited by: §1, §3.
[20] S. Athey, J. Tibshirani, and S. Wager (2019) Generalized random forests. The Annals of Statistics 47 (2), pp. 1148–1178. Cited by: §1, §3.
[21] B. Yu (2013) Stability. Bernoulli 19 (4), pp. 1484–1500. Cited by: §1.
[22] B. Yu and K. Kumbier (2020) Veridical data science. Proceedings of the National Academy of Sciences 117 (8), pp. 3920–3929. Cited by: §1.
[23] Z. T. Rewolinski and B. Yu (2025) PCS workflow for veridical data science in the age of AI. arXiv preprint arXiv:2508.00835. Cited by: §1.
[24] G. Hinton, O. Vinyals, and J. Dean (2015) Distilling the knowledge in a neural network. arXiv preprint arXiv:1503.02531. Cited by: §2.1, §2.3, §8.
[25] J. B. Kadane, J. M. Dickey, R. L. Winkler, W. S. Smith, and S. C. Peters (1980) Interactive elicitation of opinion for a normal linear model. Journal of the American Statistical Association 75 (372), pp. 845–854. Cited by: §2.1.1.
[26] E. J. Bedrick, R. Christensen, and W. Johnson (1996) A new perspective on priors for generalized linear models. Journal of the American Statistical Association 91 (436), pp. 1450–1460. Cited by: §2.1.1.
[27] J. G. Ibrahim and M. Chen (2000) Power prior distributions for regression models. Statistical Science 15 (1), pp. 46–60. Cited by: §2.1.1.
[28] J. Friedman, T. Hastie, and R. Tibshirani (2010) Regularization paths for generalized linear models via coordinate descent. Journal of Statistical Software 33 (1), pp. 1–22. Cited by: §2.3.
[29] J. Yang and T. Hastie (2024) A fast and scalable pathwise-solver for group lasso and elastic net penalized regression via block-coordinate descent. arXiv preprint arXiv:2405.08631. Cited by: §2.3.
[30] M. Lim and T. Hastie (2015) Learning interactions via hierarchical group-lasso regularization. Journal of Computational and Graphical Statistics 24 (3), pp. 627–654. Cited by: §2.3.
[31] G. Yu, J. Bien, and R. Tibshirani (2019) Reluctant interaction modeling. arXiv preprint arXiv:1907.08414. Cited by: §2.3.
[32] E. Craig and R. Tibshirani (2025) Supervised learning pays attention. arXiv preprint arXiv:2512.09912. Cited by: §3.
[33] G. Plumb, D. Molitor, and A. S. Talwalkar (2018) Model agnostic supervised local explanations. Advances in Neural Information Processing Systems 31. Cited by: §3.
[34] A. N. Angelopoulos, S. Bates, C. Fannjiang, M. I. Jordan, and T. Zrnic (2023) Prediction-powered inference. Science 382 (6671). External Links: Document Cited by: §3.
[35] A. N. Angelopoulos, J. C. Duchi, and T. Zrnic (2023) PPI++: efficient prediction-powered inference. arXiv preprint arXiv:2311.01453. Cited by: §3.
[36] X. Tian and J. Taylor (2018) Selective inference with a randomized response. The Annals of Statistics 46 (2), pp. 679–710. Cited by: Remark 1.
[37] S. Panigrahi, J. Zhu, and C. Sabatti (2021) Selection-adjusted inference: an application to confidence intervals for cis-eqtl effect sizes. Biostatistics 22 (1), pp. 181–197. Cited by: Remark 1.
[38] S. Bakshi, W. Dempsey, and S. Panigrahi (2024) Selective inference for time-varying effect moderation. arXiv preprint arXiv:2411.15908. Cited by: Remark 1.
[39] Y. Huang, S. Pirenne, S. Panigrahi, and G. Claeskens (2025) Selective inference using randomized group lasso estimators for general models. Electronic Journal of Statistics 19 (2), pp. 3489–3531. Cited by: Remark 1.
[40] R. Perry, S. Panigrahi, and D. Witten (2026) Post-selection inference for penalized m-estimators via score thinning. arXiv preprint arXiv:2601.13514. Cited by: Remark 1.
[41] N. Meinshausen and P. Bühlmann (2010) Stability selection. Journal of the Royal Statistical Society: Series B (Statistical Methodology) 72 (4), pp. 417–473. Cited by: §4.2, §5.2.
[42] A. L. Fred and A. K. Jain (2005) Combining multiple clusterings using evidence accumulation. IEEE Transactions on Pattern Analysis and Machine Intelligence 27 (6), pp. 835–850. Cited by: §4.3.
[43] R. J. Tibshirani (2013) The lasso problem and uniqueness. Electronic Journal of Statistics 7, pp. 1456–1490. Cited by: Assumption 1.
[44] O. Bousquet and A. Elisseeff (2002) Stability and generalization. Journal of Machine Learning Research 2 (Mar), pp. 499–526. Cited by: §5.2.
[45] T. C. G. A. Network (2012) Comprehensive molecular portraits of human breast tumours. Nature 490 (7418), pp. 61–70. External Links: Document Cited by: §6.
[46] P. Breheny and J. Huang (2025) Hdrm: high-dimensional regression modeling. Note: R package version 0.17.1 External Links: Link Cited by: §6.
[47] D. Ou, Z. Zhang, Z. Wu, P. Shen, Y. Huang, S. She, S. She, and M. Lin (2022) Identification of the putative tumor suppressor characteristics of FAM107A via pan-cancer analysis. Frontiers in Oncology 12, pp. 861281. Cited by: §6.
[48] J. Chu, X. Hu, C. Li, T. Li, H. Fan, and G. Jiang (2022) KLF14 alleviated breast cancer invasion and m2 macrophages polarization through modulating SOCS3/RhoA/Rock/STAT3 signaling. Cellular Signalling 92, pp. 110242. Cited by: §6.
[49] T. Chen and C. Guestrin (2016) XGBoost: a scalable tree boosting system. In Proceedings of the 22nd ACM SIGKDD International Conference on Knowledge Discovery and Data Mining, pp. 785–794. Cited by: §7.
[50] M. Kelly, R. Longjohn, and K. Nottingham (2025) The uci machine learning repository. Note: https://archive.ics.uci.edu Cited by: §7.
[51] B. Bischl, G. Casalicchio, M. Feurer, P. Gijsbers, F. Hutter, M. Lang, R. G. Mantovani, J. N. van Rijn, and J. Vanschoren (2017) OpenML benchmarking suites. arXiv preprint arXiv:1708.03731. Cited by: §7.
[52] S. F. Fischer, M. Feurer, and B. Bischl (2023) OpenML-CTR23–a curated tabular regression benchmarking suite. In AutoML Conference 2023 (Workshop), Cited by: §7.

Appendix A Details of theoretical results

A.1 Proof of main results

Proof of Lemma 1.

Let $\pi_{j}^{\circ}(\bm{z})=\mathbb{P}_{\bm{\omega}}[\,j\in\widehat{E}_{V}(\bm{z},\bm{\omega})]$ denote the selection probability of predictor $j$ under lasso regression of $\bm{z}+\bm{\omega}$ on $\bm{V}$ , evaluated as a function of the weight-adjusted response $\bm{z}$ . Then, by definition,

\pi_{j}(\bm{y})=\pi_{j}^{\circ}\bigl(\bm{z}(\bm{y})\bigr)=\pi_{j}^{\circ}\circ\bm{z}(\bm{y}).

The map $\bm{y}\mapsto\bm{z}(\bm{y})$ has entries $\sqrt{S_{k}}\,y_{k}$ for $k\leq n$ and $\sqrt{\mu}\,n_{\mathrm{eff}}^{-1/4}\hat{\phi}(\bm{x}^{*};\bm{y})$ for $k=n+1$ ; the first $n$ entries are linear in $\bm{y}$ , while the final entry is $k$ times differentiable by assumption, and hence, $\bm{z}:\mathbb{R}^{n}\to\mathbb{R}^{n+1}$ is $k$ times differentiable.

Observe that we can write

\pi_{j}^{\circ}(\bm{z}):=\mathbb{P}_{\bm{\omega}}\bigl[\,j\in\widehat{E}_{V}(\bm{z},\bm{\omega})\bigr]=\int_{\mathbb{R}^{n+1}}\chi_{j}(\bm{z}+\bm{w})\,\varphi_{\tau}(\bm{w})\,d\bm{w}=(\chi_{j}*\varphi_{\tau})(\bm{z}),

where $\chi_{j}(\bm{t})=\mathbf{1}\{\,j\in\widehat{E}(\bm{t})\,\}$ , and $\varphi_{\tau}(\bm{u})$ is the $N(\bm{0},\tau^{2}\bm{I}_{n+1})$ density function at $\bm{u}$ .

Since $\chi_{j}$ is bounded and $\varphi_{\tau}$ has integrable derivatives of all orders, differentiation under the integral sign yields $\pi_{j}^{\circ}\in C^{\infty}(\mathbb{R}^{n+1})$ , with $\partial^{\alpha}\pi_{j}^{\circ}=\chi_{j}*\partial^{\alpha}\varphi_{\tau}$ for every multi-index $\alpha$ . Applying the chain rule to $\pi_{j}=\pi_{j}^{\circ}\circ\bm{z}$ shows that, for each $r\leq k$ , every derivative of $\pi_{j}$ of order $r$ can be expressed in terms of derivatives of $\pi_{j}^{\circ}$ and derivatives of $\bm{z}$ of order at most $r$ . Since all such derivatives exist, $\pi_{j}$ is $k$ times differentiable on $\mathbb{R}^{n}$ . ∎

Proof of Theorem 1.

Combining Proposition 1 with Lemma 9 (both stated in Appendix A.3) yields, for any $y\in\mathbb{R}^{n}$ and $i\in[n]$ , that

	$\displaystyle\|\partial_{i}\pi_{j}(\bm{y})\|$	$\displaystyle\leq\sqrt{\frac{2}{\pi}}\,\frac{1}{\tau}\left\{\sqrt{S_{i}}\,\rho^{{V}}_{ij}+\sqrt{S_{n+1}}\,\|\partial_{i}\hat{\phi}(\bm{x}^{*};\bm{y})\|\,\rho^{{V}}_{n+1,j}\right\}$
		$\displaystyle\leq\sqrt{\frac{2}{\pi}}\,\frac{C}{\tau}\left\{\frac{1}{\sqrt{n_{\mathrm{eff}}}}+\frac{\|\partial_{i}\hat{\phi}(\bm{x}^{*};\bm{y})\|}{n_{\mathrm{eff}}^{1/4}}\right\}.$

Taking supremum over $y$ and $i\in[n]$ proves our claim. ∎

A.2 Theoretical analysis of stability under randomized lasso

Throughout this section, we study the stability of the feature–selection probabilities

\pi_{j}^{\circ}(\bm{z})=\mathbb{P}_{\bm{\omega}}\left[j\in\widehat{E}_{V}(\bm{z},\bm{\omega})\right].

with respect to perturbations in $\bm{z}$ . The main result of this section, Theorem 2, characterizes the sensitivity of feature-selection probabilities to perturbations in the response under randomized lasso regression.

A.2.1 Leave- $j$ -out score and lasso geometry

For a feature $j\in[p]$ , define the leave- $j$ -out lasso

\widehat{\bm{\beta}}_{V}^{(-j)}(\bm{z})=\argmin_{\bm{b}\in\mathbb{R}^{p-1}}\left\{\frac{1}{2}\|\bm{z}-\bm{V}_{-j}\bm{b}\|_{2}^{2}+\lambda\|\bm{b}\|_{1}\right\}.

(10)

Furthermore, let $\bm{R}_{-j}^{V}(\bm{z})=\bm{z}-\bm{V}_{-j}\widehat{\bm{\beta}}_{V}^{(-j)}(\bm{z})$ be the leave- $j$ -out residual, and define

T_{j}^{V}(\bm{z})=\bm{V}_{j}^{\top}\bm{R}_{-j}^{V}(\bm{z}).

(11)

Lemma 2 (KKT characterization).

Under Assumption 1, it holds that

j\notin\widehat{E}_{V}(\bm{z},\bm{\omega})\quad\Longleftrightarrow\quad|T_{j}^{V}(\bm{z}+\bm{\omega})|\leq\lambda,

where $T_{j}^{V}(\cdot)$ is as defined in (11).

Proof.

Set $\bm{t}=\bm{z}+\bm{\omega}$ . By definition, $\widehat{E}_{V}(\bm{z},\bm{\omega})$ is the support of the full Lasso solution with response $\bm{t}$ . The KKT conditions for the full Lasso problem state that a vector $\bm{\beta}\in\mathbb{R}^{p}$ is optimal if and only if there exists $\bm{\gamma}\in\mathbb{R}^{p}$ such that

\bm{V}^{\top}(\bm{t}-\bm{V}\bm{\beta})=\lambda\bm{\gamma},

(12)

where, for every $k\in[p]$ , $\bm{\gamma}_{k}=\begin{cases}\operatorname{sign}(\bm{\beta}_{k}),&\bm{\beta}_{k}\neq 0,\\ u_{k}\in[-1,1],&\bm{\beta}_{k}=0.\end{cases}$

Let $\widehat{\bm{\beta}}_{V}^{(-j)}(\bm{t})$ be the solution of the leave- $j$ -out Lasso, and define $\bm{R}_{-j}^{V}(\bm{t})=\bm{t}-\bm{V}_{-j}\widehat{\bm{\beta}}_{V}^{(-j)}(\bm{t}).$ The KKT conditions for the restricted problem imply that there exists $\bm{\gamma}_{-j}\in\mathbb{R}^{p-1}$ such that

\bm{V}_{-j}^{\top}\bm{R}_{-j}^{V}(\bm{t})=\lambda\bm{\gamma}_{-j},

(13)

with

(\bm{\gamma}_{-j})_{k}=\begin{cases}\operatorname{sign}\bigl(\widehat{\bm{\beta}}_{V,k}^{(-j)}(\bm{t})\bigr),&\widehat{\bm{\beta}}_{V,k}^{(-j)}(\bm{t})\neq 0,\\ u_{k}\in[-1,1],&\widehat{\bm{\beta}}_{V,k}^{(-j)}(\bm{t})=0.\end{cases}

Now construct the $p$ -dimensional candidate $\widetilde{\bm{\beta}}$ with

\widetilde{\bm{\beta}}_{j}=0,\qquad\widetilde{\bm{\beta}}_{-j}=\widehat{\bm{\beta}}_{V}^{(-j)}(\bm{t}).

Its residual in the full Lasso problem is exactly $\bm{t}-\bm{V}\widetilde{\bm{\beta}}=\bm{R}_{-j}^{V}(\bm{t}).$ Equation (13) verifies the full KKT conditions for every coordinate other than $j$ . Since $\widetilde{\bm{\beta}}_{j}=0$ , the remaining KKT condition is $\bm{V}_{j}^{\top}\bm{R}_{-j}^{V}(\bm{t})\in\lambda[-1,1].$ By definition, $T_{j}^{V}(\bm{t})=\bm{V}_{j}^{\top}\bm{R}_{-j}^{V}(\bm{t}),$ which gives $|T_{j}^{V}(\bm{t})|\leq\lambda.$ Therefore,

|T_{j}^{V}(\bm{t})|\leq\lambda\quad\Longleftrightarrow\quad\widetilde{\bm{\beta}}\text{ is a solution of the full Lasso problem}.

(14)

If $|T_{j}^{V}(\bm{t})|\leq\lambda$ , then (14) shows that $\widetilde{\bm{\beta}}$ is a full Lasso solution. By uniqueness of the full Lasso solution, guaranteed by Assumption 1, $\widehat{\bm{\beta}}(\bm{t})=\widetilde{\bm{\beta}}.$ Since $\widetilde{\bm{\beta}}_{j}=0$ , it follows that $j\notin\widehat{E}_{V}(\bm{z},\bm{\omega})$ .

Conversely, suppose that $j\notin\widehat{E}_{V}(\bm{z},\bm{\omega})$ , so that $\widehat{\bm{\beta}}_{j}(\bm{t})=0$ . The vector $\widehat{\bm{\beta}}_{-j}(\bm{t})$ must then solve the leave- $j$ -out problem. Indeed, if some $\bm{b}\in\mathbb{R}^{p-1}$ had a strictly smaller restricted objective, then the full vector $(\bm{b},0)$ would have a strictly smaller full Lasso objective than $\widehat{\bm{\beta}}(\bm{t})$ , contradicting its optimality. By uniqueness of the restricted solution, implied by Assumption 1, $\widehat{\bm{\beta}}_{-j}(\bm{t})=\widehat{\bm{\beta}}_{V}^{(-j)}(\bm{t}).$ The full KKT condition at the zero coefficient $\widehat{\bm{\beta}}_{j}(\bm{t})=0$ therefore gives

\left|\bm{V}_{j}^{\top}\left(\bm{t}-\bm{V}_{-j}\widehat{\bm{\beta}}_{V}^{(-j)}(\bm{t})\right)\right|\leq\lambda.

Equivalently, $|T_{j}^{V}(\bm{t})|\leq\lambda$ . We have thus proved

j\notin\widehat{E}_{V}(\bm{z},\bm{\omega})\quad\Longleftrightarrow\quad|T_{j}^{V}(\bm{t})|\leq\lambda.

∎

Having characterized the lasso active set through the leave- $j$ -out score, we now turn to the analysis of this score. To this end, we introduce some additional notation. For an active set $E\subseteq[p]\setminus\{j\}$ and sign vector $\bm{s}\in\{-1,1\}^{|E|}$ , define the polyhedral selection region

\mathcal{R}_{E,\bm{s}}^{V,-j}=\left\{\bm{z}\in\mathbb{R}^{n}:\operatorname{supp}\bigl(\widehat{\bm{\beta}}_{V}^{(-j)}(\bm{z})\bigr)=E,\ \operatorname{sign}\bigl(\widehat{\bm{\beta}}_{V,E}^{(-j)}(\bm{z})\bigr)=\bm{s}\right\},

and let $\mathcal{E}_{-j}^{V}=\left\{E:\operatorname{Leb}_{n}\bigl(\mathcal{R}_{E,\bm{s}}^{V,-j}\bigr)>0\text{ for some }\bm{s}\right\}$ be the set of essential active sets from lasso regression on $\bm{V}_{-j}$ , i.e., whose corresponding selection regions have positive Lebesgue measure.

Lemma 3 (Almost-everywhere gradient of the score).

For every $\bm{z}\in\cup_{E,s}\text{int}(\mathcal{R}_{E,\bm{s}}^{V,-j})$ , let $E=\operatorname{supp}\bigl(\widehat{\bm{\beta}}_{V}^{(-j)}(\bm{z})\bigr)$ . Then

\nabla T_{j}^{V}(\bm{z})=\bm{P}_{\bm{V}_{E}}^{\perp}\bm{V}_{j},

where $\bm{P}_{\bm{V}_{E}}=\bm{V}_{E}(\bm{V}_{E}^{\top}\bm{V}_{E})^{-1}\bm{V}_{E}^{\top}$ , and $\bm{P}_{\bm{V}_{E}}^{\perp}=\bm{I}-\bm{P}_{\bm{V}_{E}}.$

Proof.

Fix $\bm{z}\in\cup_{E,s}\text{int}(\mathcal{R}_{E,\bm{s}}^{V,-j})$ , and let $(E,\bm{s})$ denote the support and sign of the leave- $j$ -out lasso solution. Assumption 1 implies that $\bm{V}_{E}$ has full column rank, and therefore

\widehat{\bm{\beta}}_{V,E}^{(-j)}(\bm{z})=(\bm{V}_{E}^{\top}\bm{V}_{E})^{-1}(\bm{V}_{E}^{\top}\bm{z}-\lambda\bm{s})

throughout this neighborhood. It follows that

\bm{R}_{-j}^{V}(\bm{z})=\bm{z}-\bm{V}_{E}\widehat{\bm{\beta}}_{V,E}^{(-j)}(\bm{z})=\bm{P}_{\bm{V}_{E}}^{\perp}\bm{z}+\lambda\bm{V}_{E}(\bm{V}_{E}^{\top}\bm{V}_{E})^{-1}\bm{s},

and therefore,

\nabla T_{j}^{V}(\bm{z})=\nabla\left(\bm{V}_{j}^{\top}\bm{R}_{-j}^{V}(\bm{z})\right)=\bm{P}_{\bm{V}_{E}}^{\perp}\bm{V}_{j}.

Here, $\bm{z}\in\cup_{E,s}\text{int}(\mathcal{R}_{E,\bm{s}}^{V,-j})$ guarantees that $\bm{z}$ is in the interior of $\mathcal{R}_{E,\bm{s}}^{V,-j}$ , and therefore $(E,\bm{s})$ stay constant in a local neighborhood of $\bm{z}$ during differentiation. ∎

Lemma 4 (Sandwich-bound on leave- $j$ -out score).

Let $\bm{v}_{j}^{V}=\frac{\bm{V}_{j}}{\|\bm{V}_{j}\|_{2}}$ and $\bm{r}_{j\mid E}^{V}=\bm{P}_{\bm{V}_{E}}^{\perp}\bm{V}_{j}$ . For every $\bm{z}\in\mathbb{R}^{n+1}$ and $h>0$ ,

T_{j}^{V}(\bm{z}-\rho_{ij}^{V}h\bm{v}_{j}^{V})\leq T_{j}^{V}(\bm{z}+h\bm{e}_{i})\leq T_{j}^{V}(\bm{z}+\rho_{ij}^{V}h\bm{v}_{j}^{V}),

where, for $i\in[n+1]$ ,

\rho_{ij}^{V}=\sup_{{E\in\mathcal{E}_{-j}^{V},\bm{r}_{j\mid E}^{V}\neq 0}}\frac{|(\bm{r}_{j\mid E}^{V})_{i}|\|\bm{V}_{j}\|_{2}}{\|\bm{r}_{j\mid E}^{V}\|_{2}^{2}}.

Proof.

By Lemma 3, for almost every $\bm{z}$ , if $E\in\mathcal{E}_{-j}^{V}$ is the locally active leave- $j$ -out model at $\bm{z}$ , then

D_{\bm{v}_{j}^{V}}T_{j}^{V}(\bm{z})=\bigl(\bm{r}_{j\mid E}^{V}\bigr)^{\top}\bm{v}_{j}^{V}=\frac{\|\bm{r}_{j\mid E}^{V}\|_{2}^{2}}{\|\bm{V}_{j}\|_{2}}\geq 0.

When $\bm{r}_{j\mid E}^{V}\neq 0$ , the definition of $\rho_{ij}^{V}$ gives

	$\displaystyle\|\partial_{i}T_{j}^{V}(\bm{z})\|$	$\displaystyle=\|(\bm{r}_{j\mid E}^{V})_{i}\|$
		$\displaystyle\leq\rho_{ij}^{V}\frac{\\|\bm{r}_{j\mid E}^{V}\\|_{2}^{2}}{\\|\bm{V}_{j}\\|_{2}}$
		$\displaystyle=\rho_{ij}^{V}D_{\bm{v}_{j}^{V}}T_{j}^{V}(\bm{z}).$

When $\bm{r}_{j\mid E}^{V}=0$ , both sides vanish. Hence,

|\partial_{i}T_{j}^{V}|\leq\rho_{ij}^{V}D_{\bm{v}_{j}^{V}}T_{j}^{V}

Lebesgue-almost everywhere.

By linearity of directional derivatives in the direction of differentiation, it follows almost everywhere that

	$\displaystyle D_{\bm{e}_{i}+\rho_{ij}^{V}\bm{v}_{j}^{V}}T_{j}^{V}$	$\displaystyle=\partial_{i}T_{j}^{V}+\rho_{ij}^{V}D_{\bm{v}_{j}^{V}}T_{j}^{V}\geq 0,$
	$\displaystyle D_{\rho_{ij}^{V}\bm{v}_{j}^{V}-\bm{e}_{i}}T_{j}^{V}$	$\displaystyle=\rho_{ij}^{V}D_{\bm{v}_{j}^{V}}T_{j}^{V}-\partial_{i}T_{j}^{V}\geq 0.$

To apply Lemma 6, observe that

	$\displaystyle(\bm{z}+h\bm{e}_{i})-(\bm{z}-\rho_{ij}^{V}h\bm{v}_{j}^{V})$	$\displaystyle=h(\bm{e}_{i}+\rho_{ij}^{V}\bm{v}_{j}^{V}),$
	$\displaystyle(\bm{z}+\rho_{ij}^{V}h\bm{v}_{j}^{V})-(\bm{z}+h\bm{e}_{i})$	$\displaystyle=h(\rho_{ij}^{V}\bm{v}_{j}^{V}-\bm{e}_{i}).$

Thus, applying the lemma first from $\bm{z}-\rho_{ij}^{V}h\bm{v}_{j}^{V}$ in direction $\bm{e}_{i}+\rho_{ij}^{V}\bm{v}_{j}^{V}$ , and then from $\bm{z}+h\bm{e}_{i}$ in direction $\rho_{ij}^{V}\bm{v}_{j}^{V}-\bm{e}_{i}$ , gives

T_{j}^{V}(\bm{z}-\rho_{ij}^{V}h\bm{v}_{j}^{V})\leq T_{j}^{V}(\bm{z}+h\bm{e}_{i})\leq T_{j}^{V}(\bm{z}+\rho_{ij}^{V}h\bm{v}_{j}^{V}).

This proves the claim. ∎

A.2.2 Stability bound for feature–selection probabilities

Lemma 5 (Bounding probabilities under randomization).

Let $A\subseteq\mathbb{R}^{n+1}$ be measurable and define

p_{A}(\bm{z})=\mathbb{P}_{\bm{\omega}}(\bm{z}+\bm{\omega}\in A),\qquad\bm{\omega}\sim N(0,\tau^{2}\bm{I}_{n}).

Then, for every unit vector $\bm{v}$ , $D_{\bm{v}}p_{A}(\bm{z})=\frac{1}{\tau^{2}}\mathbb{E}_{\bm{\omega}}\left[(\bm{v}^{\top}\bm{\omega})\mathds{1}\{\bm{z}+\bm{\omega}\in A\}\right]$ , and $|D_{\bm{v}}p_{A}(\bm{z})|\leq\frac{1}{\tau\sqrt{2\pi}}$ .

Proof.

Differentiation under the Gaussian convolution gives us

D_{\bm{v}}p_{A}(\bm{z})=\frac{1}{\tau^{2}}\mathbb{E}_{\bm{\omega}}\left[(\bm{v}^{\top}\bm{\omega})\mathds{1}\{\bm{z}+\bm{\omega}\in A\}\right].

If $G=\bm{v}^{\top}\bm{\omega}$ , then $G\sim N(0,\tau^{2})$ and

-\mathbb{E}_{\bm{\omega}}[(-G)_{+}]\leq\mathbb{E}_{\bm{\omega}}\left[G\mathds{1}\{\bm{z}+\bm{\omega}\in A\}\right]\leq\mathbb{E}_{\bm{\omega}}[G_{+}],\text{ where }G_{+}=\max\{G,0\}.

By symmetry, $\mathbb{E}_{\bm{\omega}}[G_{+}]=\mathbb{E}_{\bm{\omega}}[(-G)_{+}]=\frac{\tau}{\sqrt{2\pi}}$ , and the result follows. ∎

Theorem 2 (Sensitivity bound for selection probabilities under randomized lasso).

Under Assumption 1, for every $\bm{z}\in\mathbb{R}^{n+1}$ and $i\in[n+1]$ ,

|\partial_{i}\pi_{j}^{\circ}(\bm{z})|\leq\sqrt{\frac{2}{\pi}}\frac{\rho_{ij}^{V}}{\tau}.

(15)

Consequently, we have

\sup_{\bm{z}\in\mathbb{R}^{n}}\|\nabla\pi_{j}^{\circ}(\bm{z})\|_{\infty}\leq\sqrt{\frac{2}{\pi}}\frac{a_{j}^{V}}{\tau},

(16)

where

a_{j}^{V}:=\max_{i\in[n+1]}\rho_{ij}^{V}=\sup_{{E\in\mathcal{E}_{-j}^{V},\bm{r}_{j\mid E}^{V}\neq 0}}\frac{\|\bm{r}_{j\mid E}^{V}\|_{\infty}\|\bm{V}_{j}\|_{2}}{\|\bm{r}_{j\mid E}^{V}\|_{2}^{2}}.

Proof.

By Lemma 2,

\pi_{j}^{\circ}(\bm{z})=p_{j,+}(\bm{z})+p_{j,-}(\bm{z}),

where

p_{j,+}(\bm{z})=\mathbb{P}_{\bm{\omega}}\{T_{j}^{V}(\bm{z}+\bm{\omega})>\lambda\},\qquad p_{j,-}(\bm{z})=\mathbb{P}_{\bm{\omega}}\{T_{j}^{V}(\bm{z}+\bm{\omega})<-\lambda\}.

Then, applying Lemma 4 yields

p_{j,+}(\bm{z}-\rho_{ij}^{V}h\bm{v}_{j}^{V})\leq p_{j,+}(\bm{z}+h\bm{e}_{i})\leq p_{j,+}(\bm{z}+\rho_{ij}^{V}h\bm{v}_{j}^{V}).

Letting $h\downarrow 0$ , we have

|\partial_{i}p_{j,+}(\bm{z})|\leq\rho_{ij}^{V}|D_{\bm{v}_{j}^{V}}p_{j,+}(\bm{z})|.

Analogously, applying Lemma 4 to $p_{j,-}(\bm{z})$ yields

|\partial_{i}p_{j,-}(\bm{z})|\leq\rho_{ij}^{V}|D_{\bm{v}_{j}^{V}}p_{j,-}(\bm{z})|

Finally, applying Lemma 5 to bound the directional derivatives on the right-hand side directional derivatives of the two above-stated inequalities, we have

|\partial_{i}p_{j,+}(\bm{z})|\leq\frac{\rho_{ij}^{V}}{\tau\sqrt{2\pi}},\qquad|\partial_{i}p_{j,-}(\bm{z})|\leq\frac{\rho_{ij}^{V}}{\tau\sqrt{2\pi}}.

Therefore,

|\partial_{i}\pi_{j}^{\circ}(\bm{z})|\leq\frac{2\rho_{ij}^{V}}{\tau\sqrt{2\pi}}=\sqrt{\frac{2}{\pi}}\frac{\rho_{ij}^{V}}{\tau}.

Taking the maximum over $i$ proves (16). ∎

A.3 Auxiliary results

For $E\in\mathcal{E}_{-j}^{V}$ , recall from Lemma 3 and Lemma 4 that $\bm{P}_{\bm{V}_{E}}$ denotes the orthogonal projection onto the column space of $\bm{V}_{E}$ , and that

\bm{r}_{j\mid E}^{V}=\bm{P}_{\bm{V}_{E}}^{\perp}\bm{V}_{j},\quad\text{where, }\quad\bm{P}_{\bm{V}_{E}}^{\perp}=\bm{I}_{n+1}-\bm{P}_{\bm{V}_{E}},

and that for $i\in[n+1]$ ,

\rho_{ij}^{V}=\sup_{\begin{subarray}{c}E\in\mathcal{E}_{-j}^{V}\\ \bm{r}_{j\mid E}^{V}\neq 0\end{subarray}}\frac{|(\bm{r}_{j\mid E}^{V})_{i}|\,\|\bm{V}_{j}\|_{2}}{\|\bm{r}_{j\mid E}^{V}\|_{2}^{2}}.

Proposition 1 (Sensitivity bound under randomized local distillation).

Suppose that $\bm{y}\longmapsto\hat{\phi}(\bm{x}^{*},\bm{y})$ belongs to $C^{1}(\mathbb{R}^{n})$ , and that Assumption 1 holds. Then, for every $\bm{y}\in\mathbb{R}^{n}$ , $i\in[n]$ , and $j\in[p]$ ,

\displaystyle|\partial_{i}\pi_{j}(\bm{y})|\leq\sqrt{\frac{2}{\pi}}\frac{1}{\tau}\left\{\sqrt{S_{i}}\,\rho_{ij}^{V}+\sqrt{S_{n+1}}\,|\partial_{i}\hat{\phi}(\bm{x}^{*},\bm{y})|\,\rho_{n+1,j}^{V}\right\}.

Proof.

By Lemma 1,

\pi_{j}(\bm{y})=\pi_{j}^{\circ}\bigl(\bm{z}(\bm{y})\bigr)

is continuously differentiable.

The coordinates of $\bm{z}(\bm{y})$ satisfy

z_{k}(\bm{y})=\sqrt{S_{k}}\,y_{k},\qquad k\in[n],\qquad z_{n+1}(\bm{y})=\sqrt{S_{n+1}}\,\hat{\phi}(\bm{x}^{*},\bm{y}),

where $S_{n+1}=\frac{\mu}{\sqrt{n_{\mathrm{eff}}}}$ . Consequently, for $i\in[n]$ and $k\in[n+1]$ ,

\frac{\partial z_{k}(\bm{y})}{\partial y_{i}}=\sqrt{S_{i}}\,\mathds{1}\{k=i\}+\sqrt{S_{n+1}}\,\partial_{i}\hat{\phi}(\bm{x}^{*},\bm{y})\,\mathds{1}\{k=n+1\}.

The chain rule therefore gives

	$\displaystyle\partial_{i}\pi_{j}(\bm{y})$	$\displaystyle=\sum_{k=1}^{n+1}\partial_{z_{k}}\pi_{j}^{\circ}\bigl(\bm{z}(\bm{y})\bigr)\frac{\partial z_{k}(\bm{y})}{\partial y_{i}}$
		$\displaystyle=\sqrt{S_{i}}\,\partial_{z_{i}}\pi_{j}^{\circ}\bigl(\bm{z}(\bm{y})\bigr)+\sqrt{S_{n+1}}\,\partial_{i}\hat{\phi}(\bm{x}^{*},\bm{y})\,\partial_{z_{n+1}}\pi_{j}^{\circ}\bigl(\bm{z}(\bm{y})\bigr).$		(17)

Applying Theorem 2 gives, for every $\bm{z}\in\mathbb{R}^{n+1}$ and $k\in[n+1]$ ,

\left|\partial_{z_{k}}\pi_{j}^{\circ}(\bm{z})\right|\leq\sqrt{\frac{2}{\pi}}\frac{\rho_{kj}^{V}}{\tau}.

Substituting these bounds into (17) and applying the triangle inequality proves the bound. ∎

Lemma 6 (Directional monotonicity).

Let $f:\mathbb{R}^{n}\to\mathbb{R}$ be Lipschitz, and let $\bm{d}\in\mathbb{R}^{n}$ . If $D_{\bm{d}}f(\bm{z})\geq 0$ for Lebesgue-almost every $\bm{z}$ , then $f(\bm{z}+t\bm{d})\geq f(\bm{z})$ for every $\bm{z}\in\mathbb{R}^{n}$ and $t\geq 0$ .

Proof.

The result is immediate if $\bm{d}=0$ . Otherwise, let $\widetilde{\bm{d}}=\bm{d}/\|\bm{d}\|_{2}$ and decompose $\mathbb{R}^{n}=\widetilde{\bm{d}}^{\perp}\oplus\operatorname{span}(\widetilde{\bm{d}})$ , and for $\bm{u}\in\widetilde{\bm{d}}^{\perp}$ , define $g_{\bm{u}}(t)=f(\bm{u}+t\widetilde{\bm{d}})$ . By Fubini’s theorem, for almost every $\bm{u}$ , the inequality

D_{\widetilde{\bm{d}}}f(\bm{u}+t\widetilde{\bm{d}})\geq 0

holds for almost every $t\in\mathbb{R}$ . Since $f$ is Lipschitz, $g_{\bm{u}}$ is absolutely continuous, and

g_{\bm{u}}^{\prime}(t)=D_{\widetilde{\bm{d}}}f(\bm{u}+t\widetilde{\bm{d}})\geq 0

for almost every $t$ . Hence $g_{\bm{u}}$ is nondecreasing.

Now fix an arbitrary $\bm{u}\in\widetilde{\bm{d}}^{\perp}$ . Choose a sequence $\bm{u}_{m}\to\bm{u}$ such that $g_{\bm{u}_{m}}$ is nondecreasing. For $t_{1}<t_{2}$ ,

f(\bm{u}_{m}+t_{1}\widetilde{\bm{d}})\leq f(\bm{u}_{m}+t_{2}\widetilde{\bm{d}}).

Passing to the limit and using continuity of $f$ yields $f(\bm{u}+t_{1}\widetilde{\bm{d}})\leq f(\bm{u}+t_{2}\widetilde{\bm{d}})$ . Thus $f$ is nondecreasing along every line parallel to $\bm{d}$ . ∎

Lemma 7 (Effective similarity neighborhood: size and weights).

Let $\mathcal{I}_{S}=\left\{i\in[n]:S_{i}\geq\frac{1}{2n_{\mathrm{eff}}}\right\}$ . Under Assumption 2, $|\mathcal{I}_{S}|\geq\frac{n_{\mathrm{eff}}}{2C_{S}^{2}}$ . Moreover, for every $i\in\mathcal{I}_{S}$ , $\frac{1}{2n_{\mathrm{eff}}}\leq S_{i}\leq\frac{C_{S}}{n_{\mathrm{eff}}}$ .

Proof.

For every $i\notin\mathcal{I}_{S}$ , the definition of $\mathcal{I}_{S}$ gives $S_{i}<\frac{1}{2n_{\mathrm{eff}}}$ , and therefore

\sum_{i\notin\mathcal{I}_{S}}S_{i}^{2}\leq\frac{1}{2n_{\mathrm{eff}}}\sum_{i\notin\mathcal{I}_{S}}S_{i}\leq\frac{1}{2n_{\mathrm{eff}}},

where the last inequality uses $\sum_{i=1}^{n}S_{i}=1$ . Since $\sum_{i=1}^{n}S_{i}^{2}=\frac{1}{n_{\mathrm{eff}}}$ , it follows that

\sum_{i\in\mathcal{I}_{S}}S_{i}^{2}\geq\frac{1}{2n_{\mathrm{eff}}}.

On the other hand, Assumption 2 implies $S_{i}\leq S_{\max}\leq\frac{C_{S}}{n_{\mathrm{eff}}}$ , for $i\in[n]$ . Therefore,

\sum_{i\in\mathcal{I}_{S}}S_{i}^{2}\leq|\mathcal{I}_{S}|\frac{C_{S}^{2}}{n_{\mathrm{eff}}^{2}}.

Combining the preceding two displays yields $|\mathcal{I}_{S}|\frac{C_{S}^{2}}{n_{\mathrm{eff}}^{2}}\geq\frac{1}{2n_{\mathrm{eff}}}$ , i.e.,

|\mathcal{I}_{S}|\geq\frac{n_{\mathrm{eff}}}{2C_{S}^{2}}.

Finally, for $i\in\mathcal{I}_{S}$ , the lower bound on $S_{i}$ follows from the definition of $\mathcal{I}_{S}$ , while the upper bound follows from Assumption 2. ∎

Lemma 8 (Bounds on residual and design in augmented regression).

Under Assumptions 2–4, define $C_{V}=B_{X}\sqrt{1+C_{\mu}}$ , and $c_{R}=\frac{\kappa_{X}}{2C_{S}}$ . Then, uniformly over $j\in[p]$ ,

\|\bm{V}_{j}\|_{2}\leq C_{V},

and, uniformly over $j\in[p]$ and $E\in\mathcal{E}_{-j}^{V}$ ,

\|\bm{r}_{j\mid E}^{V}\|_{2}\geq c_{R}.

In particular, $\bm{r}_{j\mid E}^{V}\neq 0$ for every $j\in[p]$ and every $E\in\mathcal{E}_{-j}^{V}$ .

Proof.

First, because the training weights are nonnegative and sum to one,

\sum_{i=1}^{n}S_{i}^{2}\leq\left(\sum_{i=1}^{n}S_{i}\right)^{2}=1.

Therefore, $n_{\mathrm{eff}}\geq 1$ . Using $S_{n+1}=\mu/\sqrt{n_{\mathrm{eff}}}$ and Assumption 3, we obtain

\displaystyle\|\bm{V}_{j}\|_{2}^{2}

\displaystyle=\sum_{i=1}^{n}S_{i}x_{ij}^{2}+S_{n+1}(x_{j}^{*})^{2}\leq B_{X}^{2}\sum_{i=1}^{n}S_{i}+\frac{\mu}{\sqrt{n_{\mathrm{eff}}}}B_{X}^{2}\leq B_{X}^{2}(1+C_{\mu})=C_{V}^{2}.

This proves the asserted upper bound on $\|\bm{V}_{j}\|_{2}$ .

Next, fix arbitrary $j\in[p]$ and $E\in\mathcal{E}_{-j}^{V}$ . By the definition of an orthogonal projection, $\|\bm{r}_{j\mid E}^{V}\|_{2}^{2}=\inf_{\bm{\gamma}\in\mathbb{R}^{|E|}}\|\bm{V}_{j}-\bm{V}_{E}\bm{\gamma}\|_{2}^{2}$ . For every $\bm{\gamma}\in\mathbb{R}^{|E|}$ ,

	$\displaystyle\\|\bm{V}_{j}-\bm{V}_{E}\bm{\gamma}\\|_{2}^{2}$	$\displaystyle=\sum_{i=1}^{n}S_{i}\bigl(x_{ij}-\bm{x}_{i,E}^{\top}\bm{\gamma}\bigr)^{2}+S_{n+1}\bigl(x_{j}^{}-\bm{x}_{E}^{\top}\bm{\gamma}\bigr)^{2}$
		$\displaystyle\geq\sum_{i\in\mathcal{I}_{S}}S_{i}\bigl(x_{ij}-\bm{x}_{i,E}^{\top}\bm{\gamma}\bigr)^{2}\geq\frac{1}{2n_{\mathrm{eff}}}\left\\|\bm{X}_{\mathcal{I}_{S},j}-\bm{X}_{\mathcal{I}_{S},E}\bm{\gamma}\right\\|_{2}^{2},$

where the final inequality follows from the definition of $\mathcal{I}_{S}$ .

After ordering the columns of $\bm{X}_{\mathcal{I}_{S},E\cup\{j\}}$ so that $j$ is last, we write

\bm{X}_{\mathcal{I}_{S},j}-\bm{X}_{\mathcal{I}_{S},E}\bm{\gamma}=\bm{X}_{\mathcal{I}_{S},E\cup\{j\}}\begin{pmatrix}-\bm{\gamma}\\ 1\end{pmatrix}.

Assumption 4 therefore implies

\displaystyle\left\|\bm{X}_{\mathcal{I}_{S},j}-\bm{X}_{\mathcal{I}_{S},E}\bm{\gamma}\right\|_{2}

\displaystyle\geq\kappa_{X}\sqrt{|\mathcal{I}_{S}|}\left\|\begin{pmatrix}-\bm{\gamma}\\ 1\end{pmatrix}\right\|_{2}\geq\kappa_{X}\sqrt{|\mathcal{I}_{S}|}.

It follows that

\|\bm{V}_{j}-\bm{V}_{E}\bm{\gamma}\|_{2}^{2}\geq\frac{\kappa_{X}^{2}|\mathcal{I}_{S}|}{2n_{\mathrm{eff}}}.

By Lemma 7, $|\mathcal{I}_{S}|\geq\frac{n_{\mathrm{eff}}}{2C_{S}^{2}}$ , and therefore,

\|\bm{V}_{j}-\bm{V}_{E}\bm{\gamma}\|_{2}^{2}\geq\frac{\kappa_{X}^{2}}{4C_{S}^{2}}=c_{R}^{2}.

Taking the infimum over $\bm{\gamma}$ proves $\|\bm{r}_{j\mid E}^{V}\|_{2}\geq c_{R}$ , which completes the proof. ∎

Lemma 9 (Bounds for weighted contributions from local distillation).

Under Assumptions 2–4, there exists a constant $C_{\rho}<\infty$ , such that, uniformly over $j\in[p]$ ,

\max_{i\in[n]}\sqrt{S_{i}}\,\rho_{ij}^{V}\leq\frac{C_{\rho}}{\sqrt{n_{\mathrm{eff}}}},\quad\sqrt{S_{n+1}}\,\rho_{n+1,j}^{V}\leq\frac{C_{\rho}}{n_{\mathrm{eff}}^{1/4}}.

Proof.

By Lemma 8, uniformly over $j\in[p]$ and $E\in\mathcal{E}_{-j}^{V}$ , we have

\|\bm{V}_{j}\|_{2}\leq C_{V},\qquad\|\bm{r}_{j\mid E}^{V}\|_{2}\geq c_{R}.

Since $|(\bm{r}_{j\mid E}^{V})_{k}|\leq\|\bm{r}_{j\mid E}^{V}\|_{2}$ , we obtain, for every $k\in[n+1]$ ,

\rho_{kj}^{V}\leq\frac{C_{V}}{c_{R}}=:\overline{C}_{\rho}.

For $i\in[n]$ , Assumption 2 gives $S_{i}\leq S_{\max}\leq\frac{C_{S}}{n_{\mathrm{eff}}}$ , and therefore

\sqrt{S_{i}}\,\rho_{ij}^{V}\leq\frac{\sqrt{C_{S}}\,\overline{C}_{\rho}}{\sqrt{n_{\mathrm{eff}}}}.

Finally,

\sqrt{S_{n+1}}\,\rho_{n+1,j}^{V}=\frac{\sqrt{\mu}}{n_{\mathrm{eff}}^{1/4}}\rho_{n+1,j}^{V}\leq\frac{\sqrt{C_{\mu}}\,\overline{C}_{\rho}}{n_{\mathrm{eff}}^{1/4}}.

Letting $C_{\rho}=\sqrt{C_{S}}\,\overline{C}_{\rho}\vee\sqrt{C_{\mu}}\,\overline{C}_{\rho}$ proves both claims. ∎

Appendix B Performance on UCI ML and OpenML datasets

Here we report the complete results across all 17 benchmark datasets and methods, including a second tabular foundation-model teacher, TabFM [6], and the cross-validated teacher-selection rule (Section 2.3) omitted from the main-text figure for legibility. Table 1 lists per-dataset sample sizes and feature counts; Figure 6 shows test $R^{2}$ .

Dataset	$n$	$p$	Source
Automobile	159	51	UCI
Servo	167	10	UCI
Liver Disorders	341	5	UCI
Auto MPG	392	8	UCI
Real Estate Valuation	414	6	UCI
Student Performance	649	39	UCI
Cars	804	17	OpenML
QSAR Fish Toxicity	907	6	OpenML
Concrete Compressive Strength	1005	8	OpenML
Infrared Thermography Temperature	1018	43	UCI
Socmob	1156	35	OpenML
Red Wine	1359	11	OpenML
Airfoil Self-Noise	1503	5	OpenML
Auction Verification	2043	7	OpenML
Space GA	3107	6	OpenML
White Wine	3961	11	OpenML
Abalone	4177	9	OpenML

Table 1: Benchmark datasets after complete-case filtering and de-duplication, sorted by sample size.

p

is the number of columns in the design matrix, i.e. after one-hot encoding of categorical variables.

Figure 6: Predictive performance (test

R^{2}

) across 17 UCI ML and OpenML regression datasets. Each teacher is shown with its two distilled local-linear students, lasso (L) and ridge (R). Local distillation approaches its teacher’s accuracy across datasets. Dist (best teacher) selects, per split, the teacher with the lowest cross-validated error. LLF and LOESS are teacher-unaware local methods and are more variable.

Figure 7: (Continued.) Remaining datasets, same methods and axes as above.

Appendix C Ablation study

Local distillation has two components: (1) the similarity weights and (2) the “prior” prediction from the teacher. Here, we find that both combined usually have the best predictive performance (Figure 8).

Figure 8: Similarity weights and prediction prior are both useful for prediction. Ablation of local distillation’s two teacher-derived components: the similarity weights and the prediction prior. There are four comparators: Neither (a global lasso), Weights only, Prior only, and Both (the full method); each panel shows one dataset. Points show median test

R^{2}

over 20 splits (

\pm 1

SE); the teacher is TabPFN, the student a lasso. Both components help, and using both is best or tied-best on nearly every dataset.

Appendix D Performance with simulated data

Local distillation is intended to summarize the feature–outcome relationship linearly at each prediction point. Here we study its performance when the data generating process is known, and test how the method responds to changes in the amount of noise, and the sizes of $n$ and $p$ . We simulate training data as follows:

$\displaystyle x_{i}$	$\displaystyle\sim N_{p}(0,\Sigma),\qquad\Sigma_{jk}=0.3^{\|j-k\|},$
$\displaystyle y_{i}$	$\displaystyle=f(x_{i})+\varepsilon_{i},\qquad\varepsilon_{i}\sim N(0,\sigma^{2}),\quad\text{where}$
$\displaystyle f(x)$	$\displaystyle=(1+x_{3})\,x_{1}+(1-x_{3})\,x_{2}+2x_{3}$
	$\displaystyle=x_{1}+x_{2}+(2+x_{1}-x_{2})x_{3}.$	(18)

The true local coefficients $\nabla f(x)=(1+x_{3},\ 1-x_{3},\ 2+x_{1}-x_{2},\ 0,\ldots,0)$ vary with respect to $x_{1},x_{2}$ and $x_{3}$ , but the global linear model cannot model this heterogeneity. Our testing data in all cases is $40$ query points drawn from the same distribution.

We vary $n$ , $p$ and $\sigma$ ; for each combination, we run 50 iterations of local distillation with a gradient boosting teacher and lasso student. Table 2 reports (1) correlations across query points between the fitted local coefficients $\hat{\beta}_{1},\hat{\beta}_{2}$ and $\hat{\beta}_{3}$ and the true local gradients; (2) support recovery described by feature selection true positive rate and the number of false positives (incorrectly selected features); (3) $\hat{\mu}$ and the fraction of runs in which $\hat{\mu}\leq 1$ and the method reverted to the global fit.

We find that local distillation is generally strong across $\sigma$ , $n$ and $p$ (including $p>n$ ), though it naturally degrades as the level of noise grows (and eventually reverts to the global linear model). In particular, the first two coefficients $(1+x_{3})$ and $(1-x_{3})$ are typically well fit, with correlation to the true values between $0.63$ and $0.79$ . However, the third coefficient, $2+x_{1}-x_{2}$ , is harder to recover (correlation $0.14$ – $0.34$ ). This is a result of our definition of “similarity”: the difference $x_{1}-x_{2}$ affects the response only through its interaction with $x_{3}$ , so observations with very different values of $x_{1}-x_{2}$ can have similar teacher predictions, and each local fit returns roughly the average coefficient over its neighborhood.

$n$	$p$	$\sigma$	$\hat{\mu}$	reverted	$\mathrm{corr}(\hat{\beta}_{1})$	$\mathrm{corr}(\hat{\beta}_{2})$	$\mathrm{corr}(\hat{\beta}_{3})$	TPR	FP
400	100	0.3	2.89 (0.49)	0/50	0.79 (0.06)	0.78 (0.06)	0.34 (0.13)	0.97 (0.02)	7.4 (3.7)
400	100	1	1.46 (0.15)	0/50	0.77 (0.07)	0.77 (0.07)	0.30 (0.13)	0.96 (0.02)	9.3 (5.6)
400	100	2	1.07 (0.06)	11/50	0.73 (0.08)	0.72 (0.08)	0.24 (0.16)	0.94 (0.04)	10.3 (5.9)
400	100	3	1.02 (0.02)	40/50	0.63 (0.09)	0.64 (0.06)	0.14 (0.14)	0.92 (0.05)	9.3 (5.3)
200	100	0.3	1.86 (0.37)	0/50	0.76 (0.07)	0.75 (0.07)	0.33 (0.17)	0.96 (0.02)	7.7 (4.6)
400	800	0.3	2.43 (0.33)	0/50	0.78 (0.06)	0.77 (0.07)	0.34 (0.15)	0.95 (0.02)	10.4 (7.6)
200	500	0.3	1.66 (0.36)	0/50	0.77 (0.06)	0.75 (0.08)	0.25 (0.17)	0.94 (0.03)	9.8 (7.8)

Table 2: Recovery of local structure using a gradient boosting teacher and lasso student. Means (standard deviations) reported over runs, with

50

replicates per row. “reverted” counts replicates with

\hat{\mu}\leq 1

, in which the method reverted to the global fit. The correlation columns show the correlations between the fitted and true coefficients, for runs when

\hat{\mu}>1

. “TPR” is the true positive rate of the selected coefficients, “FP” is the number of false positive selections.

We then repeated this experiment to test the effect of filtering by selection probabilities. For each configuration, we took the first $10$ replicates, selected $\hat{t}$ using the rule of Section 4.1, and ran $B=100$ randomized refits per query point to compute $\hat{\pi}_{j}$ for every selected feature. Table 3 reports the results. We find that filtering reduces the false positive rate for feature selection, and retains most of the true positives.

					share with $\hat{\pi}_{j}>0.9$		FP per query
$n$	$p$	$\sigma$	reverted	$\hat{t}$	active	spurious	unfiltered	filtered	TPR (filtered)
400	100	0.3	0/10	1.14 (0.98)	0.98 (0.02)	0.12 (0.14)	8.3 (5.1)	0.9 (1.0)	0.95 (0.02)
400	100	1	0/10	0.67 (0.76)	0.98 (0.03)	0.21 (0.23)	10.4 (6.5)	2.0 (2.2)	0.93 (0.05)
400	100	2	0/10	0.89 (0.92)	0.95 (0.06)	0.12 (0.12)	12.4 (6.7)	1.8 (2.2)	0.87 (0.09)
400	100	3	8/10	0.62 (0.53)	0.93 (0.03)	0.05 (0.07)	9.1 (0.6)	0.5 (0.7)	0.86 (0.01)
200	100	0.3	0/10	0.39 (0.37)	0.98 (0.02)	0.21 (0.27)	7.3 (4.9)	1.6 (2.0)	0.94 (0.03)
400	800	0.3	0/10	0.66 (0.76)	0.98 (0.02)	0.15 (0.22)	10.1 (9.1)	1.7 (2.3)	0.93 (0.02)
200	500	0.3	0/10	0.42 (0.75)	0.98 (0.02)	0.27 (0.28)	9.0 (5.5)	2.4 (2.9)	0.94 (0.01)

Table 3: Validation of the selection probabilities against ground truth. Means (standard deviations) across replicates with

\hat{\mu}>1

; maximum

10

per row, with

B=100

randomized refits per query point at the scale

\hat{t}

selected per replicate as described in Section 4.1. “Reverted” counts replicates with

\hat{\mu}\leq 1

, which are excluded from the other columns. “Active” and “spurious” refer to selections of the

3

truly active and the

p-3

null features, respectively; “filtered” retains only selections with

\hat{\pi}_{j}>0.9

. The selection probabilities identify the false selections: filtering at

\hat{\pi}_{j}>0.9

reduces them from roughly

8

–

12

per query point to fewer than

2.5

, and retains

86

–

95\%

of the truly active features.

	$\displaystyle\|\partial_{i}T_{j}^{V}(\bm{z})\|$	$\displaystyle=\|(\bm{r}_{j\mid E}^{V})_{i}\|$
		$\displaystyle\leq\rho_{ij}^{V}\frac{\\|\bm{r}_{j\mid E}^{V}\\|_{2}^{2}}{\\|\bm{V}_{j}\\|_{2}}$
		$\displaystyle=\rho_{ij}^{V}D_{\bm{v}_{j}^{V}}T_{j}^{V}(\bm{z}).$

Interpretable AI with Local Distillation

Interpretable AI with Local Distillation

Abstract

1 Introduction

2 Local distillation

2.1 Model fitting and prediction

2.1.1 The teacher as a Bayesian prior

2.2 A worked example of local distillation

2.3 Computability, generalizations and further localization

Computability.

Other forms of regularization, and nonlinearities in the student.

Generalization to other losses.

Further localization.

Teacher selection.

Cross-modal and cross-domain distillation.

3 Related work

Local linear modeling.

Prediction-powered inference.

4 Interpretability through stability

4.1 Randomized local distillation

Remark 1.

4.2 Interpreting a single local model

Remark 2.

Remark 3.

4.3 Interpreting the local models in aggregate

5 Theoretical analysis of stability

5.1 Notation and preliminaries

5.2 Stability of randomized local distillation

Lemma 1 (Smoothness of feature–selection probabilities).

Assumption 1 (Weight-adjusted design in general position).

Assumption 2 (Bounded weight concentration).

Assumption 3 (Bounded design and local distillation weight).

Assumption 4 (Non-degeneracy of restricted design).

Theorem 1 (Sensitivity bound for selection probabilities).

Corollary 1 (Stability guarantee).

Extension to stability over a regularization-grid.

6 A high-dimensional example: predicting gene expression

7 Benchmark comparisons

8 Discussion

Extensions.

Local regression more generally.

References

Appendix A Details of theoretical results

A.1 Proof of main results

Proof of Lemma 1.

Proof of Theorem 1.

A.2 Theoretical analysis of stability under randomized lasso

A.2.1 Leave-jj-out score and lasso geometry

Lemma 2 (KKT characterization).

Proof.

Lemma 3 (Almost-everywhere gradient of the score).

Proof.

Lemma 4 (Sandwich-bound on leave-jj-out score).

Proof.

A.2.2 Stability bound for feature–selection probabilities

Lemma 5 (Bounding probabilities under randomization).

Proof.

Theorem 2 (Sensitivity bound for selection probabilities under randomized lasso).

Proof.

A.3 Auxiliary results

Proposition 1 (Sensitivity bound under randomized local distillation).

Proof.

Lemma 6 (Directional monotonicity).

Proof.

Lemma 7 (Effective similarity neighborhood: size and weights).

Proof.

Lemma 8 (Bounds on residual and design in augmented regression).

Proof.

Lemma 9 (Bounds for weighted contributions from local distillation).

Proof.

Appendix B Performance on UCI ML and OpenML datasets

Appendix C Ablation study

Appendix D Performance with simulated data

Comments

A.2.1 Leave- $j$ -out score and lasso geometry

Lemma 4 (Sandwich-bound on leave- $j$ -out score).