Interpretable AI with Local Distillation

Interpretable AI with Local Distillation

arXiv:2608.23538v1 [stat.ME] 24 Aug 2026

Erin Craig∗ Department of Biostatistics, University of Michigan and Yiling Huang Department of Statistics, University of Michigan and Snigdha Panigrahi Department of Statistics, University of Michigan August 25, 2026

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 ∗

Corresponding author. Email: ercr@umich.edu.

1

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.

2

local coefficients

horsepower

4

cylinders

2

displacement

0

acceleration

−2

origin: European

−4

origin: Japanese model year

coefficient

global coefficients weight

Test car, ordered by increasing predicted mpg

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 3

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 X ∈ Rn×p (standardized), a continuous response y ∈ Rn , and a test point x∗ . We also have a teacher ϕ̂ : Rp → R, a fitted model that predicts y from x. Our goal is to predict y at 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.

4

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 " n # n X
2
2 1 X ⊤ ⊤ β̂ = arg min yj − x j β + µ ϕ̂(xj ) − xj β , 2n j=1 β j=1

(1)

where hyperparameter µ determines the influence of the teacher. When E[y | 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 x∗ , replacing the global distillation sum with a term that pulls the student’s prediction x∗ ⊤ β 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 {Ŝj (x∗ ) : j ∈ {1, . . . , n}} that upweight training points whose teacher prediction is close to ϕ̂(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 µ̂ from the data as the cross-validated student-to-teacher loss ratio, √ rather than tuning it; see (2). We scale it further by 1/ n̂eff , where n̂eff is the effective sample size defined in (3). The teacher’s influence therefore grows both when the teacher outperforms the student globally (larger µ̂) and when the local neighborhood is sparse (smaller n̂eff ), where the local fit most needs anchoring. When µ̂ ≤ 1, the teacher offers no improvement and we instead return the global linear fit for all test observations.1 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).

5

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 ϕ̂(−j) (xj ) for the teacher’s prediction at training point j computed without access to (xj , yj ): 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), ϕ̂(−j) is the out-of-fold (OOF) prediction; if the teacher makes no use of the training data, ϕ̂(−j) = ϕ̂. 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 (X, y) with n observations; test observation x∗ ; teacher ϕ̂; elastic-net parameter α ∈ [0, 1]. Output: Prediction ŷ ∗ and local coefficients β̂ ∗ at x∗ . 1. Estimate distillation strength as the cross-validated student-to-teacher loss ratio: µ̂ =

L̂(fβ̂ ) L̂(ϕ̂)

n

,

2 1X L̂(g) = yj − g (−j) (xj ) , n j=1

(2)

where fβ̂ is the global elastic-net fit with shrinkage parameter λ̂ chosen by crossvalidation, and g (−j) denotes prediction at xj without access to (xj , yj ). If µ̂ ≤ 1, return ŷ ∗ = fβ̂ (x∗ ) and β̂ ∗ = β̂, and stop. 2. Compute similarity weights and effective sample size:
(−j) ∗ 2 ϕ̂ (x ) − ϕ̂(x ) exp(−d ) j j Ŝj = Ŝj (x∗ ) = Pn , dj = , σ̂ϕ2 k=1 exp(−dk )

n̂eff =

X n

Ŝj2

−1 ,

j=1

(3) where σ̂ϕ2 is the empirical variance of {ϕ̂(−j) (xj )}nj=1 . 3. Fit the local model and predict: ŷ ∗ = x∗⊤ β̂ ∗ , where β̂ ∗ = arg min β

n
2
2   1X µ̂ (1 − α) Ŝj yj −x⊤ ϕ̂(x∗ )−x∗⊤ β + λ̂ α∥β∥1 + ∥β∥22 . j β + √ 2 j=1 2 2 n̂eff

(4) Note: For a set of test observations, run step 1 once and steps 2–3 independently (in parallel) for each x∗ .

6

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 x∗⊤ β, centered at ϕ̂(x∗ ) with precision µ̂ √ . This connects to data augmentation priors [25, 26], in which prior beliefs are encoded n̂eff via pseudo-observations rather than parameter distributions. These priors express beliefs about observable quantities y | 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 √ from the data, in n̂eff 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; µ̂ 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 12.81 teacher 6.31; therefore, µ̂ = = 2.03. Since µ̂ > 1, the teacher improves on the 6.31 global lasso and we proceed; had µ̂ ≤ 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 n̂eff range from 39 to 172 (median 141) out of ntrain = 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

7

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 x∗ , Equation (4) is a weighted elasticnet problem with fixed λ̂ and α: append x∗ to X and ϕ̂(x∗ ) to y, and use weights √ {Ŝ1 , . . . , Ŝn , µ̂/ n̂eff }. Any solver that supports observation weights (e.g. glmnet [28] or adelie [29]) can fit it directly. Fitting many models scales naturally: λ̂ is estimated once on the training set, and the local fits are then independent single-λ 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 x with a basis expansion Φ(x) (e.g. pairwise interactions) yields a student that is linear in Φ(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 ℓ(y, η) convex in η = x⊤ β: the local objective of Algorithm 1 becomes ∗

β̂ = arg min β

n X

µ̂ Ŝj ℓ(yj , x⊤ ℓ(ϕ̂(x∗ ), x∗⊤ β) + λ̂[α∥β∥1 + (1−α) ∥β∥22 ], j β) + √ 2 n̂eff j=1

(5)

where µ̂ is the cross-validated student-to-teacher loss ratio with ℓ in place of squared error. Taking ℓ(y, η) = 21 (y − η)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 dj are computed on the logit scale.

8

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 (ϕ̂(x∗ ) − x∗⊤ β)2 on the log-relative-hazard scale. Further localization. For each test observation, one could select λ by local cross-validation and estimate µ from locally weighted losses. We instead reuse the global λ̂: the local loss in (4) is a weighted mean on the same scale as the cross-validation loss used to select λ̂, so the same value is a sensible default, and refitting locally is expensive and in our experiments rarely improved prediction. Local estimates of µ 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 L̂(ϕ̂) 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 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, µ̂ guards against domain shift: if the teacher does not transfer well, µ̂ ≤ 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. 9

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 pseudoobservation 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 α = 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 populationlevel 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 10

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 µ̂ and reversion to the global fit. The power-tuning parameter λPPI ∈ [0, 1] of PPI++ (chosen from data to minimize variance) mirrors the role of µ̂, though µ̂ 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 x∗ and features 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 w = (w1 , . . . , wn+1 ) denote a vector of (n + 1) independent and identically distributed randomization variables, with wj ∼ N (0, τ 2 ) for j = 1, . . . , n + 1. Rather than relying on a single local fit, we generate repeated randomized local fits by solving n
2
2   1X µ̂ (1 − α) w ∗ ∗⊤ 2 Ŝj yj − x⊤ β + ϕ̂(x ) − x β + λ̂ α∥β∥ + ∥β∥ β̂ = arg min 1 j 2 2 j=1 2(n̂eff )1/2 2 β !⊤ n X µ̂1/2 1/2 Ŝj wj xj + 1/4 wn+1 x∗ − β, neff j=1 (6) where β̂ w denotes the coefficient vector obtained at x∗ for a given randomization draw 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 11

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 wj in the linear perturbation term is scaled by the square root of its corresponding weight, Ŝj . Consequently, if Ŝj is zero or close to zero, wj has no or little effect on the optimization. The amount of randomization in each randomized local fit is controlled by τ 2 , the variance σ̂ 2 of the normal randomization variables in w. In practice, we set τ 2 = t ∗ , where n̂eff = n̂eff 1 is the effective sample size (Equation (3)), σ̂ 2 is the variance of the teacher’s out-ofŜ 2 +...+Ŝ 2 n

1

fold residuals, and t is a constant chosen by the analyst. Choosing the randomization variance σ̂ 2 as a fixed fraction of ensures that the scale of the randomization is comparable to the scale n̂eff 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 τ 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 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 : β̂j ̸= 0} (for α > 0), providing an interpretation of its prediction at 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 x∗ . To address this instability, we compute for each feature its empirical selection frequency

12

across the randomized local fits B o 1 X n wb π̂j = I β̂j ̸= 0 , B b=1

where B is the number of refits of (6) and wb denotes the randomization used in the b-th fit. That is, π̂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 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

OOF median squared error

flags these as uncertain.

3.2

Feature

Coef.

Sel. prob.

3.0

weight

−5.11

100%

2.8

model year

2.28

100%

2.6

origin: European

0.40

99%

2.4

acceleration

0.52

99%

origin: Japanese

0.41

98%

cylinders

0.24

75%

horsepower

−0.17

70%

2.2 2.0

+5% 10−2

10−1

100

Noise multiple

101

𝘵

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 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 13

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 n̂eff is small, i.e., when the weights Ŝ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 ntest 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 y rather than artifacts of a single fit. We therefore aggregate the clustering obtained across the randomized refits 14

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 ntest fitted local models, with coefficient vectors in Rp . 2. Choose the number of clusters: cluster the ntest unperturbed coefficient vectors and select the appropriate k. 3. Cluster the randomized refits: draw B independent randomization vectors w1 , . . . , wB , as described in Section 4.1. For each draw wb , run randomized local distillation (solve (6)) to obtain ntest 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 ntest = 157 observations and (2) select the number of clusters 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 cooccurrence 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.

15

global coefficients

C1 (mean 16.4)

local coefficients

C2 (mean 26.9)

C3 (mean 35.7)

weight horsepower cylinders

2

displacement

0

coefficient

4

acceleration

−2

origin: European

−4

origin: Japanese model year

Test car, ordered by increasing predicted mpg within cluster Figure 3: Vehicle subgroups and their local miles per gallon models. Local coefficients for the Auto MPG data, with test cars grouped into three subgroups (C1–C3) by the stable clustering of Section 4.3. Subgroups are ordered left to right by mean MPG, and cars by predicted MPG within subgroup.

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 Ŝj and the regularization parameters µ̂ and λ̂ as fixed. To avoid confusion, we drop the hats and write these quantities as Sj , µ, and λ, respectively; similarly we write neff for the effective sample size. We slightly modify our notation for the teacher’s prediction at the query point, denoting it by ϕ̂(x∗ ; 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., α = 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 16

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     y=  

y1 .. .





x⊤ 1



 .   .   .  X=  ∈ R(n+1)×p , ⊤  x   n  (x∗ )⊤

    ∈ Rn+1 , yn   ϕ̂(x∗ ; y)

denote the augmented response vector and design matrix, obtained by appending the teacher’s prediction ϕ̂(x∗ ; y) and the query point x∗ to the response vector and design matrix, respectively. Let Ω = diag(S1 , . . . , Sn , Sn+1 ) ∈ R(n+1)×(n+1) , denote the diagonal matrix of weights for the n + 1 observations in the augmented dataset, where Sn+1 = √nµeff is the weight assigned to the pseudo-response at the query point. We then define z = z(y1 , . . . , yn ) = Ω1/2 y ∈ Rn+1 ,

V = Ω1/2 X ∈ R(n+1)×p ,

(7)

to be the weight-adjusted response and design, respectively, i.e., coordinatewise, Zk = √ √ √ µ Sk yk (k ≤ n), and Zn+1 = Sn+1 ϕ̂(x∗ ; y) = 1/4 ϕ̂(x∗ ; y). neff

Under this notation, when α = 1 and λ denotes the tuning parameter for the lasso penalty, the randomized local distillation problem (6) is equivalent to solving βb ω = argmin 21 ∥z + ω − V β∥22 + λ∥β∥1 ,

(8)

β∈Rp

where ω ∼ N (0, τ 2 In+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. 17

  bV (z, ω) = supp βbω (z) denote the active set obtained by solving the randomized Let E lasso in (8). The quantity of interest is the feature-selection probability for predictor j, evaluated as a function of y and denoted by h i bV (z, ω) , πj (y) = Pω j ∈ E

for j ∈ [p],

(9)

where the probability is over the randomization ω, with y and hence z = z(y1 , . . . , yn ) treated as fixed, and [p] = {1, 2, . . . , p}. The subscript ω 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 w, they are smooth provided that the teacher’s prediction is a smooth function of the input 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 ∈ N ∪ {∞} and suppose that y 7→ ϕ̂(x∗ ; y) belongs to C k (Rn ). Then for every j ∈ [p], the selection probability πj (y), defined in (9), also belongs to C k (Rn ). We now state Theorem 1, which bounds the sensitivity of the feature–selection probabilities to perturbations of the training response y, leveraging the smoothness of these probabilities established in Lemma 1. We derive this result under the following assumptions on the design matrix X, the similarity weights Sj and the weight-adjusted design matrix V . Assumption 1 (Weight-adjusted design in general position). The columns of V are in general position [43]. Assumption 2 (Bounded weight concentration). We assume that there exists a constant CS < ∞ such that neff Smax ≤ CS . Assumption 3 (Bounded design and local distillation weight). For some constants BX , Cµ >  0, we assume that max maxi∈[n] ∥xi ∥∞ , ∥x∗ ∥∞ ≤ BX , and 0 < µ ≤ Cµ . Assumption 4 (Non-degeneracy of restricted design). For some constant κX > 0, we assume V V that, uniformly over j ∈ [p] and E ∈ E−j , where E−j denotes the collection of essential active

18

sets from lasso regression on V−j (i.e., active sets from the leave-j-covariate-out lasso fit whose corresponding selection regions have positive Lebesgue measure), ! 1 σmin p XIS ,E∪{j} ≥ κX , |IS | o n where IS = i ∈ [n] : Si ≥ 2n1eff 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 IS , 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 y 7→ ϕ̂(x∗ ; y) belongs to C 1 (Rn ) with Lϕ̂ (x∗ ) = supy∈Rn ∥∇ϕ̂(x∗ ; y)∥∞ < ∞, and that Assumptions 1, 2, 3 and 4 hold. Then there exists a constant C < ∞ such that, for every j ∈ [p], ! r Lϕ̂ (x∗ ) 2C 1 sup ∥∇πj (y)∥∞ ≤ . + √ 1/4 πτ neff y∈Rn neff 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 −1/2

first is a direct contribution, excluding the effect of distillation, which is of order neff , and the second contribution arises through the teacher’s prediction, which is of order −1/4

neff Lϕ̂ (x∗ ). (ii) Although we make no assumption on Lϕ̂ (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. 19

(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 r 2 C −1/2 sup ∥∇πj (y)∥∞ ≤ n , π τ eff y∈Rn since Lϕ̂ (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 Sj = 1/n, the same bound holds with neff = 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 < ∞ such that, for every j ∈ [p], i ∈ [n], and for any y, y ′ ∈ Rn satisfying yk = yk′ for all k ̸= i (equivalently, differing only in the i-th coordinate), ! r ∗ L (x ) C 1 2 ϕ̂ + πj (y) − πj (y ′ ) ≤ |yi − yi′ |. √ 1/4 πτ neff neff

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 τ . 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 τ 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 Λ of lasso regularization parameters λ, rather than on a single fixed value of λ. As proposed in the stability selection framework of [41], a natural quantity on which to base decisions is max πj (y; λ), where πj (y; λ) denotes the λ∈Λ

selection probability at λ (with the dependence on λ 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. 20

Instead, one can consider a smooth approximation to max πj (y; λ), namely ΠΛ (y) = λ∈Λ  P  1 1 + log |Λ| λ∈Λ exp (ηπj (y; λ)) , for fixed η ∈ R . Like the selection probability for any η fixed λ, this quantity also takes values in [0, 1]. In practice, using its empirical counterpart b Λ (y), one can then determine the set of selected features as {j : Π b Λ (y) ≥ pthr }, for a chosen Π threshold pthr . Furthermore, the gradient of ΠΛ (y) can be written as X exp(ηπj (y; λ)) ∇ΠΛ (y) = , wλ (y) ∇πj (y; λ), with wλ (y) = P ′ ′ ∈Λ exp(ηπj (y; λ )) λ λ∈Λ 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 Λ would yield a corresponding stability guarantee for ΠΛ (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 ntrain = 321 and ntest = 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 (R2 = 0.579), TabPFN 0.150 (R2 = 0.666), and local distillation 0.148 (R2 = 0.670): a 22% reduction relative to the global lasso, 21

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 Applying the stability screen of Section 4 (π̂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. global

local

coefficients

coefficients C1 (mean -2.37)

C2 (mean -1.84)

C3 (mean -1.46)

C4 (mean -1.07)

C5 (mean -0.43)

FGFRL1 FAM107A KLF14

0.10

0.05

FBXL13 TOP2A

0.00

CENPK

coefficient

CTXN3

−0.05

RPL27

−0.10

LOC201229 POC1A DTL NBR2

Clustered test patient, ordered by predicted BRCA1 expression within cluster

Figure 4: Local BRCA1-model coefficients across patient clusters. Columns are test patients (n = 215); rows are the genes whose local coefficients vary most across clusters, together with two of the largest-magnitude genes for reference. Patients are grouped into five clusters using 100 randomized fits (Section 4.3), and ordered within cluster by teacherpredicted BRCA1 expression ϕ̂(x∗ ). The leftmost column is the global lasso fit on the same training data, for comparison. Color scale is clipped at the 99th percentile of |β̂|. 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.

22

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 ∈ [159, 4177] and feature counts p ∈ [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 R2 , 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.

23

R2 comparison across methods Datasets ordered by TabPFN's median R² improvement over the lasso. Error bars represent 1 SE across 20 train/test splits. airfoil self noise

servo

automobile

socmob

cars

student performance

TabPFN LD (TabPFN, L) LD (TabPFN, R) XGBoost LD (XGB, L) LD (XGB, R) Lasso Ridge LLF LOESS

TabPFN LD (TabPFN, L) LD (TabPFN, R) XGBoost LD (XGB, L) LD (XGB, R) Lasso Ridge LLF LOESS 0.50

1.00

0.50

1.00

0.50

1.00

Test R2

Figure 5: Predictive performance (test R2 , median ± 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: 24

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 µ̂ 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 µ̂ ≤ 1, the estimated teacher error is worse than that of the student). We considered a continuous alternative, weighting the teacher by its excess performance (µ̂ − 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 µ̂ and determining when to revert to a simpler model may nevertheless be worth exploring.

25

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: claudeopus-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).

26

References [1] John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Žı́dek, Anna Potapenko, et al. Highly accurate protein structure prediction with AlphaFold. Nature, 596(7873):583–589, 2021. [2] Remi Lam, Alvaro Sanchez-Gonzalez, Matthew Willson, Peter Wirnsberger, Meire Fortunato, Ferran Alet, Suman Ravuri, Timo Ewalds, Zach Eaton-Rosen, Weihua Hu, et al. Learning skillful medium-range global weather forecasting. Science, 382(6677): 1416–1421, 2023. [3] Noah Hollmann, Samuel Müller, Lennart Purucker, Arjun Krishnakumar, Max Körfer, Shi Bin Hoo, Robin Tibor Schirrmeister, and Frank Hutter. Accurate predictions on small data with a tabular foundation model. Nature, 637(8045):319–326, 2025. [4] Christopher Kolberg, Jules Kreuer, Jonas Huurdeman, Sofiane Ouaari, Katharina Eggensperger, and Nico Pfeifer. TabPFN-wide: Continued pre-training for extreme feature counts. arXiv preprint arXiv:2510.06162, 2025. [5] Junwei Ma, Valentin Thomas, Rasa Hosseinzadeh, Alex Labach, Jesse Cresswell, Keyvan Golestan, Guangwei Yu, Anthony L Caterini, and Maks Volkovs. TabDPT: Scaling tabular foundation models on real data. Advances in Neural Information Processing Systems, 38, 2026. [6] Weihao Kong and Abhimanyu Das. Introducing TabFM: A zero-shot foundation model for tabular data. Google Research Blog, June 2026. URL https://research.google/ blog/introducing-tabfm-a-zero-shot-foundation-model-for-tabular-data/. Accessed July 7, 2026. [7] Jingang Qu, David Holzmüller, Gaël Varoquaux, and Marine Le Morvan. TabICL: A tabular foundation model for in-context learning on large data. arXiv preprint arXiv:2502.05564, 2025. [8] Marco Tulio Ribeiro, Sameer Singh, and Carlos Guestrin. “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, pages 1135–1144, 2016. 27

[9] Scott M Lundberg and Su-In Lee. A unified approach to interpreting model predictions. Advances in Neural Information Processing Systems, 30, 2017. [10] Damien Garreau and Ulrike von Luxburg. Explaining the explainer: A first theoretical analysis of LIME. In International Conference on Artificial Intelligence and Statistics, pages 1287–1296. PMLR, 2020. [11] Dylan Slack, Sophie Hilgard, Emily Jia, Sameer Singh, and Himabindu Lakkaraju. Fooling LIME and SHAP: Adversarial attacks on post hoc explanation methods. In Proceedings of the AAAI/ACM Conference on AI, Ethics, and Society, pages 180–186, 2020. [12] I Elizabeth Kumar, Suresh Venkatasubramanian, Carlos Scheidegger, and Sorelle Friedler. Problems with Shapley-value-based explanations as feature importance measures. In International Conference on Machine Learning, pages 5491–5500. PMLR, 2020. [13] Kjersti Aas, Martin Jullum, and Anders Løland. Explaining individual predictions when features are dependent: More accurate approximations to Shapley values. Artificial Intelligence, 298:103502, 2021. [14] Cynthia Rudin. Stop explaining black box machine learning models for high stakes decisions and use interpretable models instead. Nature Machine Intelligence, 1(5):206– 215, 2019. [15] R. Quinlan. Auto mpg. UCI Machine Learning Repository, 1993. https://doi.org/ 10.24432/C5859H. Dataset from the 1983 American Statistical Association Data Exposition. [16] William S Cleveland and Susan J Devlin. Locally weighted regression: an approach to regression analysis by local fitting. Journal of the American Statistical Association, 83 (403):596–610, 1988. [17] Rui Qiu, Zhou Yu, and Ruoqing Zhu. Random forest weighted local Fréchet regression with random objects. Journal of Machine Learning Research, 25(107):1–69, 2024. [18] Adam Bloniarz, Ameet Talwalkar, Bin Yu, and Christopher Wu. Supervised neighborhoods for distributed nonparametric regression. In Artificial Intelligence and Statistics, pages 1450–1459. PMLR, 2016. 28

[19] Rina Friedberg, Julie Tibshirani, Susan Athey, and Stefan Wager. Local linear forests. Journal of Computational and Graphical Statistics, 30(2):503–517, 2020. [20] Susan Athey, Julie Tibshirani, and Stefan Wager. Generalized random forests. The Annals of Statistics, 47(2):1148–1178, April 2019. [21] Bin Yu. Stability. Bernoulli, 19(4):1484–1500, 2013. [22] Bin Yu and Karl Kumbier. Veridical data science. Proceedings of the National Academy of Sciences, 117(8):3920–3929, 2020. [23] Zachary T Rewolinski and Bin Yu. PCS workflow for veridical data science in the age of AI. arXiv preprint arXiv:2508.00835, 2025. [24] Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge in a neural network. arXiv preprint arXiv:1503.02531, 2015. [25] Joseph B Kadane, James M Dickey, Robert L Winkler, Wayne S Smith, and Stephen C Peters. Interactive elicitation of opinion for a normal linear model. Journal of the American Statistical Association, 75(372):845–854, 1980. [26] Edward J Bedrick, Ronald Christensen, and Wesley Johnson. A new perspective on priors for generalized linear models. Journal of the American Statistical Association, 91 (436):1450–1460, 1996. [27] Joseph G Ibrahim and Ming-Hui Chen. Power prior distributions for regression models. Statistical Science, 15(1):46–60, 2000. [28] Jerome Friedman, Trevor Hastie, and Robert Tibshirani. Regularization paths for generalized linear models via coordinate descent. Journal of Statistical Software, 33(1): 1–22, 2010. [29] James Yang and Trevor Hastie. A fast and scalable pathwise-solver for group lasso and elastic net penalized regression via block-coordinate descent.

arXiv preprint

arXiv:2405.08631, 2024. [30] Michael Lim and Trevor Hastie. Learning interactions via hierarchical group-lasso regularization. Journal of Computational and Graphical Statistics, 24(3):627–654, 2015. [31] Guo Yu, Jacob Bien, and Ryan Tibshirani. Reluctant interaction modeling. arXiv preprint arXiv:1907.08414, 2019. 29

[32] Erin Craig and Robert Tibshirani. Supervised learning pays attention. arXiv preprint arXiv:2512.09912, 2025. [33] Gregory Plumb, Denali Molitor, and Ameet S Talwalkar. Model agnostic supervised local explanations. Advances in Neural Information Processing Systems, 31, 2018. [34] Anastasios N. Angelopoulos, Stephen Bates, Clara Fannjiang, Michael I. Jordan, and Tijana Zrnic. Prediction-powered inference. Science, 382(6671), 2023. doi: 10.1126/ science.adi6000. [35] Anastasios N Angelopoulos, John C Duchi, and Tijana Zrnic.

PPI++: Efficient

prediction-powered inference. arXiv preprint arXiv:2311.01453, 2023. [36] Xiaoying Tian and Jonathan Taylor. Selective inference with a randomized response. The Annals of Statistics, 46(2):679–710, 2018. [37] Snigdha Panigrahi, Junjie Zhu, and Chiara Sabatti. Selection-adjusted inference: an application to confidence intervals for cis-eqtl effect sizes. Biostatistics, 22(1):181–197, 2021. [38] Soham Bakshi, Walter Dempsey, and Snigdha Panigrahi. Selective inference for timevarying effect moderation. arXiv preprint arXiv:2411.15908, 2024. [39] Yiling Huang, Sarah Pirenne, Snigdha Panigrahi, and Gerda Claeskens. Selective inference using randomized group lasso estimators for general models. Electronic Journal of Statistics, 19(2):3489–3531, 2025. [40] Ronan Perry, Snigdha Panigrahi, and Daniela Witten. Post-selection inference for penalized m-estimators via score thinning. arXiv preprint arXiv:2601.13514, 2026. [41] Nicolai Meinshausen and Peter Bühlmann. Stability selection. Journal of the Royal Statistical Society: Series B (Statistical Methodology), 72(4):417–473, 2010. [42] Ana LN Fred and Anil K Jain. Combining multiple clusterings using evidence accumulation. IEEE Transactions on Pattern Analysis and Machine Intelligence, 27(6):835–850, 2005. [43] Ryan J. Tibshirani. The lasso problem and uniqueness. Electronic Journal of Statistics, 7:1456–1490, 2013.

30

[44] Olivier Bousquet and André Elisseeff. Stability and generalization. Journal of Machine Learning Research, 2(Mar):499–526, 2002. [45] The Cancer Genome Atlas Network. Comprehensive molecular portraits of human breast tumours. Nature, 490(7418):61–70, 2012. doi: 10.1038/nature11412. [46] Patrick Breheny and Jian Huang. hdrm: High-dimensional regression modeling, 2025. URL https://github.com/pbreheny/hdrm. R package version 0.17.1. [47] Dehua Ou, Zhiqin Zhang, Zesong Wu, Peilin Shen, Yichuan Huang, Sile She, Sifan She, and Ming-en Lin. Identification of the putative tumor suppressor characteristics of FAM107A via pan-cancer analysis. Frontiers in Oncology, 12:861281, 2022. [48] Jian Chu, Xing-Chi Hu, Chang-Chun Li, Tang-Ya Li, Hui-Wen Fan, and Guo-Qin Jiang. KLF14 alleviated breast cancer invasion and m2 macrophages polarization through modulating SOCS3/RhoA/Rock/STAT3 signaling. Cellular Signalling, 92:110242, 2022. [49] Tianqi Chen and Carlos Guestrin. XGBoost: A scalable tree boosting system. In Proceedings of the 22nd ACM SIGKDD International Conference on Knowledge Discovery and Data Mining, pages 785–794, 2016. [50] Markelle Kelly, Rachel Longjohn, and Kolby Nottingham. The uci machine learning repository. https://archive.ics.uci.edu, 2025. [51] Bernd Bischl, Giuseppe Casalicchio, Matthias Feurer, Pieter Gijsbers, Frank Hutter, Michel Lang, Rafael G Mantovani, Jan N van Rijn, and Joaquin Vanschoren. OpenML benchmarking suites. arXiv preprint arXiv:1708.03731, 2017. [52] Sebastian Felix Fischer, Matthias Feurer, and Bernd Bischl. OpenML-CTR23–a curated tabular regression benchmarking suite. In AutoML Conference 2023 (Workshop), 2023.

A

Details of theoretical results

A.1

Proof of main results

bV (z, ω)] denote the selection probability of preProof of Lemma 1. Let πj◦ (z) = Pω [ j ∈ E dictor j under lasso regression of z + ω on V , evaluated as a function of the weight-adjusted

31

response z. Then, by definition,
πj (y) = πj◦ z(y) = πj◦ ◦ z(y). The map y 7→ z(y) has entries

√

Sk yk for k ≤ n and

√ −1/4 µ neff ϕ̂(x∗ ; y) for k = n + 1; the

first n entries are linear in y, while the final entry is k times differentiable by assumption, and hence, z : Rn → Rn+1 is k times differentiable. Observe that we can write πj◦ (z) := Pω



 bV (z, ω) = j∈E

Z χj (z + w) φτ (w) dw = (χj ∗ φτ )(z), Rn+1

b }, and φτ (u) is the N (0, τ 2 In+1 ) density function at u. where χj (t) = 1{ j ∈ E(t) Since χj is bounded and φτ has integrable derivatives of all orders, differentiation under the integral sign yields πj◦ ∈ C ∞ (Rn+1 ), with ∂ α πj◦ = χj ∗ ∂ α φτ for every multi-index α. Applying the chain rule to πj = πj◦ ◦ z shows that, for each r ≤ k, every derivative of πj of order r can be expressed in terms of derivatives of πj◦ and derivatives of z of order at most r. Since all such derivatives exist, πj is k times differentiable on Rn . Proof of Theorem 1. Combining Proposition 1 with Lemma 9 (both stated in Appendix A.3) yields, for any y ∈ Rn and i ∈ [n], that r o p 2 1 np V ∗ V |∂i πj (y)| ≤ Si ρij + Sn+1 |∂i ϕ̂(x ; y)| ρn+1,j π τ ( ) r 2 C 1 |∂i ϕ̂(x∗ ; y)| + . ≤ √ 1/4 π τ neff neff Taking supremum over y and i ∈ [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 h i ◦ b πj (z) = Pω j ∈ EV (z, ω) . with respect to perturbations in 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.

32

A.2.1

Leave-j-out score and lasso geometry

For a feature j ∈ [p], define the leave-j-out lasso   1 (−j) 2 βbV (z) = arg min ∥z − V−j b∥2 + λ∥b∥1 . 2 b∈Rp−1

(10)

(−j) V (z) = z − V−j βbV (z) be the leave-j-out residual, and define Furthermore, let R−j V TjV (z) = Vj⊤ R−j (z).

(11)

Lemma 2 (KKT characterization). Under Assumption 1, it holds that bV (z, ω) j∈ /E

|TjV (z + ω)| ≤ λ,

⇐⇒

where TjV (·) is as defined in (11). bV (z, ω) is the support of the full Lasso solution with Proof. Set t = z + ω. By definition, E response t. The KKT conditions for the full Lasso problem state that a vector β ∈ Rp is optimal if and only if there exists γ ∈ Rp such that V ⊤ (t − V β) = λγ,

where, for every k ∈ [p], γk =

 sign(βk ),

(12)

βk ̸= 0,

u ∈ [−1, 1], β = 0. k k (−j)

Let βbV

(−j)

V (t) = t−V−j βbV (t) be the solution of the leave-j-out Lasso, and define R−j

(t).

The KKT conditions for the restricted problem imply that there exists γ−j ∈ Rp−1 such that ⊤ V V−j R−j (t) = λγ−j ,

with

(13)

 sign βb(−j) (t)
, βb(−j) (t) ̸= 0, V,k V,k (γ−j )k = (−j) u ∈ [−1, 1], βbV,k (t) = 0. k

Now construct the p-dimensional candidate βe with βej = 0,

(−j) βe−j = βbV (t).

V Its residual in the full Lasso problem is exactly t − V βe = R−j (t). Equation (13) verifies the full KKT conditions for every coordinate other than j. Since βej = 0, the remaining

33

V V KKT condition is Vj⊤ R−j (t) ∈ λ[−1, 1]. By definition, TjV (t) = Vj⊤ R−j (t), which gives

|TjV (t)| ≤ λ. Therefore, |TjV (t)| ≤ λ

⇐⇒

βe is a solution of the full Lasso problem.

(14)

If |TjV (t)| ≤ λ, then (14) shows that βe is a full Lasso solution. By uniqueness of the b e Since βej = 0, it follows that full Lasso solution, guaranteed by Assumption 1, β(t) = β. bV (z, ω). j∈ /E bV (z, ω), so that βbj (t) = 0. The vector βb−j (t) must then Conversely, suppose that j ∈ /E solve the leave-j-out problem. Indeed, if some b ∈ Rp−1 had a strictly smaller restricted objective, then the full vector (b, 0) would have a strictly smaller full Lasso objective than b β(t), contradicting its optimality. By uniqueness of the restricted solution, implied by (−j)

Assumption 1, βb−j (t) = βbV

(t). The full KKT condition at the zero coefficient βbj (t) = 0

therefore gives Vj⊤



 (−j) b t − V−j βV (t) ≤ λ.

Equivalently, |TjV (t)| ≤ λ. We have thus proved bV (z, ω) j∈ /E

⇐⇒

|TjV (t)| ≤ λ.

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 ⊆ [p] \ {j} and sign vector s ∈ {−1, 1}|E| , define the polyhedral selection region o n

n b(−j) (z) = E, sign βb(−j) (z) = s , β RV,−j = z ∈ R : supp V V,E E,s V and let E−j =

n

o
E : Lebn RV,−j > 0 for some s be the set of essential active sets from E,s

lasso regression on V−j , i.e., whose corresponding selection regions have positive Lebesgue measure. Lemma 3 (Almost-everywhere gradient of the score). For every z ∈ ∪E,s int(RV,−j E,s ), let
(−j) E = supp βb (z) . Then V

∇TjV (z) = PV⊥E Vj , where PVE = VE (VE⊤ VE )−1 VE⊤ , and PV⊥E = I − PVE .

34

Proof. Fix z ∈ ∪E,s int(RV,−j E,s ), and let (E, s) denote the support and sign of the leave-j-out lasso solution. Assumption 1 implies that VE has full column rank, and therefore (−j) βbV,E (z) = (VE⊤ VE )−1 (VE⊤ z − λs)

throughout this neighborhood. It follows that (−j) V R−j (z) = z − VE βbV,E (z) = PV⊥E z + λVE (VE⊤ VE )−1 s,

and therefore,
V ∇TjV (z) = ∇ Vj⊤ R−j (z) = PV⊥E Vj . V,−j Here, z ∈ ∪E,s int(RV,−j E,s ) guarantees that z is in the interior of RE,s , and therefore (E, s)

stay constant in a local neighborhood of z during differentiation. V

V Lemma 4 (Sandwich-bound on leave-j-out score). Let vjV = ∥Vjj∥2 and rj|E = PV⊥E Vj . For

every z ∈ Rn+1 and h > 0, TjV (z − ρVij hvjV ) ≤ TjV (z + hei ) ≤ TjV (z + ρVij hvjV ), where, for i ∈ [n + 1], ρVij =

V )i |∥Vj ∥2 |(rj|E

sup

V ∥22 ∥rj|E

V ,r V ̸=0 E∈E−j j|E

.

V Proof. By Lemma 3, for almost every z, if E ∈ E−j is the locally active leave-j-out model

at z, then DvjV TjV (z) =

⊤ V V rj|E vj

=

V ∥rj|E ∥22

∥Vj ∥2

V ̸= 0, the definition of ρVij gives When rj|E V |∂i TjV (z)| = |(rj|E )i |

≤ ρVij

V ∥rj|E ∥22

∥Vj ∥2

= ρVij DvjV TjV (z). V When rj|E = 0, both sides vanish. Hence,

|∂i TjV | ≤ ρVij DvjV TjV Lebesgue-almost everywhere. 35

≥ 0.

By linearity of directional derivatives in the direction of differentiation, it follows almost everywhere that Dei +ρVij vjV TjV = ∂i TjV + ρVij DvjV TjV ≥ 0, DρVij vjV −ei TjV = ρVij DvjV TjV − ∂i TjV ≥ 0. To apply Lemma 6, observe that (z + hei ) − (z − ρVij hvjV ) = h(ei + ρVij vjV ), (z + ρVij hvjV ) − (z + hei ) = h(ρVij vjV − ei ). Thus, applying the lemma first from z −ρVij hvjV in direction ei +ρVij vjV , and then from z +hei in direction ρVij vjV − ei , gives TjV (z − ρVij hvjV ) ≤ TjV (z + hei ) ≤ TjV (z + ρVij hvjV ). This proves the claim. A.2.2

Stability bound for feature–selection probabilities

Lemma 5 (Bounding probabilities under randomization). Let A ⊆ Rn+1 be measurable and define ω ∼ N (0, τ 2 In ).

pA (z) = Pω (z + ω ∈ A),

  Then, for every unit vector v, Dv pA (z) = τ12 Eω (v ⊤ ω)1{z + ω ∈ A} , and |Dv pA (z)| ≤ τ

√1 . 2π

Proof. Differentiation under the Gaussian convolution gives us Dv pA (z) =

  1 Eω (v ⊤ ω)1{z + ω ∈ A} . 2 τ

If G = v ⊤ ω, then G ∼ N (0, τ 2 ) and −Eω [(−G)+ ] ≤ Eω [G1{z + ω ∈ A}] ≤ Eω [G+ ], where G+ = max{G, 0}. By symmetry, Eω [G+ ] = Eω [(−G)+ ] = √τ2π , and the result follows. Theorem 2 (Sensitivity bound for selection probabilities under randomized lasso). Under Assumption 1, for every z ∈ Rn+1 and i ∈ [n + 1], r |∂i πj◦ (z)| ≤

36

2 ρVij . π τ

(15)

Consequently, we have r sup ∥∇πj◦ (z)∥∞ ≤ z∈Rn where aVj

:= max i∈[n+1]

ρVij =

2 aVj , π τ

V ∥rj|E ∥∞ ∥Vj ∥2

sup V ,r V ̸=0 E∈E−j j|E

V ∥22 ∥rj|E

(16)

.

Proof. By Lemma 2, πj◦ (z) = pj,+ (z) + pj,− (z), where pj,+ (z) = Pω {TjV (z + ω) > λ},

pj,− (z) = Pω {TjV (z + ω) < −λ}.

Then, applying Lemma 4 yields pj,+ (z − ρVij hvjV ) ≤ pj,+ (z + hei ) ≤ pj,+ (z + ρVij hvjV ). Letting h ↓ 0, we have |∂i pj,+ (z)| ≤ ρVij |DvjV pj,+ (z)|. Analogously, applying Lemma 4 to pj,− (z) yields |∂i pj,− (z)| ≤ ρVij |DvjV pj,− (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 ρVij |∂i pj,+ (z)| ≤ √ , τ 2π

ρVij |∂i pj,− (z)| ≤ √ . τ 2π

Therefore, |∂i πj◦ (z)| ≤

2ρVij √ = τ 2π

r

2 ρVij . π τ

Taking the maximum over i proves (16).

A.3

Auxiliary results

V For E ∈ E−j , recall from Lemma 3 and Lemma 4 that PVE denotes the orthogonal projection

onto the column space of VE , and that V rj|E = PV⊥E Vj ,

where,

PV⊥E = In+1 − PVE ,

and that for i ∈ [n + 1], ρVij =

sup

V |(rj|E )i | ∥Vj ∥2 V ∥rj|E ∥22

V E∈E−j V rj|E ̸=0

37

.

Proposition 1 (Sensitivity bound under randomized local distillation). Suppose that y 7−→ ϕ̂(x∗ , y) belongs to C 1 (Rn ), and that Assumption 1 holds. Then, for every y ∈ Rn , i ∈ [n], and j ∈ [p], r |∂i πj (y)| ≤

o p 2 1 np Si ρVij + Sn+1 |∂i ϕ̂(x∗ , y)| ρVn+1,j . πτ

Proof. By Lemma 1, πj (y) = πj◦ z(y)

is continuously differentiable. The coordinates of z(y) satisfy p zk (y) = Sk yk , k ∈ [n],

zn+1 (y) =

p Sn+1 ϕ̂(x∗ , y),

where Sn+1 = √nµeff . Consequently, for i ∈ [n] and k ∈ [n + 1], p ∂zk (y) p = Si 1{k = i} + Sn+1 ∂i ϕ̂(x∗ , y) 1{k = n + 1}. ∂yi The chain rule therefore gives n+1 X

∂zk (y) ∂yi k=1 p
p
= Si ∂zi πj◦ z(y) + Sn+1 ∂i ϕ̂(x∗ , y) ∂zn+1 πj◦ z(y) .

∂i πj (y) =

∂zk πj◦ z(y)

(17)

Applying Theorem 2 gives, for every z ∈ Rn+1 and k ∈ [n + 1], r V 2 ρkj ◦ . ∂zk πj (z) ≤ π τ Substituting these bounds into (17) and applying the triangle inequality proves the bound.

Lemma 6 (Directional monotonicity). Let f : Rn → R be Lipschitz, and let d ∈ Rn . If Dd f (z) ≥ 0 for Lebesgue-almost every z, then f (z + td) ≥ f (z) for every z ∈ Rn and t ≥ 0. Proof. The result is immediate if d = 0. Otherwise, let de = d/∥d∥2 and decompose Rn = e and for u ∈ de⊥ , define gu (t) = f (u + td). e By Fubini’s theorem, for almost de⊥ ⊕ span(d), every u, the inequality e ≥0 Ddef (u + td) holds for almost every t ∈ R. Since f is Lipschitz, gu is absolutely continuous, and e ≥0 gu′ (t) = Ddef (u + td) 38

for almost every t. Hence gu is nondecreasing. Now fix an arbitrary u ∈ de⊥ . Choose a sequence um → u such that gum is nondecreasing. For t1 < t2 , e ≤ f (um + t2 d). e f (um + t1 d) e ≤ f (u + t2 d). e Thus f is Passing to the limit and using continuity of f yields f (u + t1 d) nondecreasing along every line parallel to d. o n Lemma 7 (Effective similarity neighborhood: size and weights). Let IS = i ∈ [n] : Si ≥ 2n1eff . neff 1 Under Assumption 2, |IS | ≥ 2C ≤ Si ≤ nCeffS . 2 . Moreover, for every i ∈ IS , 2n eff S

Proof. For every i ∈ / IS , the definition of IS gives Si < 2n1eff , and therefore X

Si2 ≤

i∈I / S

where the last inequality uses

1 X 1 Si ≤ , 2neff 2neff i∈I / S

Pn

i=1 Si = 1. Since

X

Si2 ≥

i∈IS

Pn

1 2 i=1 Si = neff , it follows that

1 . 2neff

On the other hand, Assumption 2 implies Si ≤ Smax ≤ nCeffS , for i ∈ [n]. Therefore, C2 Si2 ≤ |IS | 2S . neff i∈I X S

C2

Combining the preceding two displays yields |IS | n2S ≥ 2n1eff , i.e., eff

|IS | ≥

neff . 2CS2

Finally, for i ∈ IS , the lower bound on Si follows from the definition of IS , while the upper bound follows from Assumption 2. Lemma 8 (Bounds on residual and design in augmented regression). Under Assumptions 2– p κX 4, define CV = BX 1 + Cµ , and cR = 2C . Then, uniformly over j ∈ [p], S ∥Vj ∥2 ≤ CV , V and, uniformly over j ∈ [p] and E ∈ E−j , V ∥rj|E ∥2 ≥ cR . V V In particular, rj|E ̸= 0 for every j ∈ [p] and every E ∈ E−j .

39

Proof. First, because the training weights are nonnegative and sum to one, !2 n n X X Si2 ≤ Si = 1. i=1

i=1

√ Therefore, neff ≥ 1. Using Sn+1 = µ/ neff and Assumption 3, we obtain ∥Vj ∥22 =

n X

2 Si x2ij + Sn+1 (x∗j )2 ≤ BX

n X

i=1

µ 2 2 Si + √ BX ≤ BX (1 + Cµ ) = CV2 . neff i=1

This proves the asserted upper bound on ∥Vj ∥2 . V Next, fix arbitrary j ∈ [p] and E ∈ E−j . By the definition of an orthogonal projection, V ∥rj|E ∥22 = inf γ∈R|E| ∥Vj − VE γ∥22 . For every γ ∈ R|E| ,

∥Vj − VE γ∥22 =

n X

Si xij − x⊤ i,E γ

2

+ Sn+1 x∗j − x∗⊤ E γ

2

≥

2

i=1

≥

X

Si xij − x⊤ i,E γ

i∈IS

1 ∥XIS ,j − XIS ,E γ∥22 , 2neff

where the final inequality follows from the definition of IS . After ordering the columns of XIS ,E∪{j} so that j is last, we write ! −γ XIS ,j − XIS ,E γ = XIS ,E∪{j} . 1 Assumption 4 therefore implies p ∥XIS ,j − XIS ,E γ∥2 ≥ κX |IS | It follows that

−γ

!

1

≥ κX

p |IS |.

2

κ2X |IS | 2 ∥Vj − VE γ∥2 ≥ . 2neff

neff By Lemma 7, |IS | ≥ 2C 2 , and therefore, S

∥Vj − VE γ∥22 ≥

κ2X = c2R . 2 4CS

V Taking the infimum over γ proves ∥rj|E ∥2 ≥ cR , which completes the proof.

Lemma 9 (Bounds for weighted contributions from local distillation). Under Assumptions 2–4, there exists a constant Cρ < ∞, such that, uniformly over j ∈ [p], p Cρ Si ρVij ≤ √ , i∈[n] neff

p Cρ Sn+1 ρVn+1,j ≤ 1/4 . neff

max

40

V Proof. By Lemma 8, uniformly over j ∈ [p] and E ∈ E−j , we have V ∥2 ≥ cR . ∥rj|E

∥Vj ∥2 ≤ CV ,

V V ∥2 , we obtain, for every k ∈ [n + 1], )k | ≤ ∥rj|E Since |(rj|E

ρVkj ≤

CV =: C ρ . cR

For i ∈ [n], Assumption 2 gives Si ≤ Smax ≤ nCeffS , and therefore √ p CS C ρ V Si ρij ≤ √ . neff Finally,

√ p µ V Sn+1 ρn+1,j = 1/4 ρVn+1,j ≤ neff p √ Letting Cρ = CS C ρ ∨ Cµ C ρ proves both claims.

B

p

Cµ C ρ 1/4

.

neff

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 teacherselection 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 R2 .

41

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. R2 comparison across methods Points are medians; error bars are 1 SE across 20 train/test splits. airfoil self noise

servo

auction verification

concrete comp. strength

automobile

space ga

socmob

real estate valuation

white wine

TabPFN LD (TabPFN, L) LD (TabPFN, R) XGBoost LD (XGB, L) LD (XGB, R) TabFM LD (TabFM, L) LD (TabFM, R) Dist (best teacher) Lasso Ridge LLF LOESS

TabPFN LD (TabPFN, L) LD (TabPFN, R) XGBoost LD (XGB, L) LD (XGB, R) TabFM LD (TabFM, L) LD (TabFM, R) Dist (best teacher) Lasso Ridge LLF LOESS

TabPFN LD (TabPFN, L) LD (TabPFN, R) XGBoost LD (XGB, L) LD (XGB, R) TabFM

42

R2 comparison across methods Points are medians; error bars are 1 SE across 20 train/test splits. infrared therm. temp.

auto mpg

qsar fish toxicity

abalone

red wine

liver disorders

cars

student performance

TabPFN LD (TabPFN, L) LD (TabPFN, R) XGBoost LD (XGB, L) LD (XGB, R) TabFM LD (TabFM, L) LD (TabFM, R) Dist (best teacher) Lasso Ridge LLF LOESS

TabPFN LD (TabPFN, L) LD (TabPFN, R) XGBoost LD (XGB, L) LD (XGB, R) TabFM LD (TabFM, L) LD (TabFM, R) Dist (best teacher) Lasso Ridge LLF LOESS 0.00

0.50

1.00

TabPFN LD (TabPFN, L) LD (TabPFN, R) XGBoost LD (XGB, L) LD (XGB, R) TabFM LD (TabFM, L) LD (TabFM, R) Dist (best teacher) Lasso Ridge LLF LOESS 0.00

0.50

1.00

0.00

0.50

1.00

Test R2

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

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 7).

43

Ablation: both components are useful for prediction Similarity weights and the teacher prior, on/off. Points are medians; error bars 1 SE across 20 splits. airfoil self noise

servo

auction verification

concrete comp. strength

automobile

Neither Weights only Prior only Both (full) 0.60

0.80

1.00

0.60

space ga

0.80

socmob

0.70 0.80 0.90 1.00 0.60

0.70

real estate valuation

0.80

0.90

white wine

0.60

0.70

0.80

0.90

infrared therm. temp.

Neither Weights only Prior only Both (full) 0.60

0.70

0.80

auto mpg

0.90

0.60

0.65

qsar fish toxicity

0.70

0.75

0.30

abalone

0.35

0.40

red wine

0.55

0.60

0.65

liver disorders

Neither Weights only Prior only Both (full) 0.80

0.85 cars

0.90

0.55

0.60

0.52 0.54 0.56 0.58

0.36

0.38

0.40

0.15

0.20

student performance

Neither Weights only Prior only Both (full) 0.92

0.94

0.960.20 0.22 0.24 0.26 0.28

Test R2

Figure 7: 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 R2 over 20 splits (±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.

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: xi ∼ Np (0, Σ),

Σjk = 0.3|j−k| ,

yi = f (xi ) + εi ,

εi ∼ N (0, σ 2 ),

where

f (x) = (1 + x3 ) x1 + (1 − x3 ) x2 + 2x3 = x1 + x2 + (2 + x1 − x2 )x3 .

(18)

The true local coefficients ∇f (x) = (1 + x3 , 1 − x3 , 2 + x1 − x2 , 0, . . . , 0) vary with respect to x1 , x2 and x3 , but the global linear model cannot model this heterogeneity. Our testing 44

n

p

σ

µ̂

reverted

corr(β̂1 )

corr(β̂2 )

corr(β̂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 µ̂ ≤ 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 µ̂ > 1. “TPR” is the true positive rate of the selected coefficients, “FP” is the number of false positive selections. data in all cases is 40 query points drawn from the same distribution. We vary n, p and σ; 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 β̂1 , β̂2 and β̂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) µ̂ and the fraction of runs in which µ̂ ≤ 1 and the method reverted to the global fit. We find that local distillation is generally strong across σ, 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 + x3 ) and (1 − x3 ) are typically well fit, with correlation to the true values between 0.63 and 0.79. However, the third coefficient, 2 + x1 − x2 , is harder to recover (correlation 0.14–0.34). This is a result of our definition of “similarity”: the difference x1 − x2 affects the response only through its interaction with x3 , so observations with very different values of x1 − x2 can have similar teacher predictions, and each local fit returns roughly the average coefficient over its neighborhood. We then repeated this experiment to test the effect of filtering by selection probabilities. For each configuration, we took the first 10 replicates, selected t̂ using the rule of Section 4.1, and ran B = 100 randomized refits per query point to compute π̂j for every selected feature. Table 3 reports the results. We find that filtering reduces the false positive rate for feature

45

share with π̂j > 0.9 n

p

σ

reverted

t̂

active

spurious

400

100

400

100

0.3

0/10

1.14 (0.98)

0.98 (0.02)

0.12 (0.14)

1

0/10

0.67 (0.76)

0.98 (0.03)

0.21 (0.23)

400

100

2

0/10

0.89 (0.92)

0.95 (0.06)

0.12 (0.12)

400

100

3

8/10

0.62 (0.53)

0.93 (0.03)

200

100

0.3

0/10

0.39 (0.37)

400

800

0.3

0/10

200

500

0.3

0/10

FP per query unfiltered

filtered

TPR (filtered)

8.3 (5.1)

0.9 (1.0)

0.95 (0.02)

10.4 (6.5)

2.0 (2.2)

0.93 (0.05)

12.4 (6.7)

1.8 (2.2)

0.87 (0.09)

0.05 (0.07)

9.1 (0.6)

0.5 (0.7)

0.86 (0.01)

0.98 (0.02)

0.21 (0.27)

7.3 (4.9)

1.6 (2.0)

0.94 (0.03)

0.66 (0.76)

0.98 (0.02)

0.15 (0.22)

10.1 (9.1)

1.7 (2.3)

0.93 (0.02)

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 µ̂ > 1; maximum 10 per row, with B = 100 randomized refits per query point at the scale t̂ selected per replicate as described in Section 4.1. “Reverted” counts replicates with µ̂ ≤ 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 π̂j > 0.9. The selection probabilities identify the false selections: filtering at π̂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. selection, and retains most of the true positives.

46

Related documents

• Mapping the Impact of Generative AI on Disinformation: Insights from a Scoping Review#481882
• Smoothing the Semantic Landscape: Generalizable AI-Generated Image Detection via Text-Induced Flatness#1116218
• Neuron-Aware Data Selection for Annotation-Free LLM Self-Distillation#332513
• Contrastive On-Policy Distillation#1143849
• Critique-Guided Distillation for Robust Reasoning via Refinement#1050754