Interpretable AI with Local Distillation

Erin Craig; Yiling Huang; Snigdha Panigrahi

Source-linked AI summary

Interpretable AI with Local Distillation

Erin Craig, Yiling Huang, Snigdha Panigrahi

arXiv:2608.23538v1Aug 24, 2026stat.MEcs.LGstat.ML

TL;DR

Accuracy alone is insufficient for acting on AI predictions, motivating interpretable models. Local distillation uses a black-box teacher to guide local linear models, nearly matching the teacher’s predictive accuracy while supporting stable feature interpretation.

Problem
Acting on AI predictions requires more than accuracy alone, creating a need for interpretable modeling.
Method
Local distillation uses a teacher-guided local linear model and randomized refits to assess feature-selection stability.
Results
Local distillation nearly matches its teacher’s predictive accuracy, while randomization has theoretical stability guarantees under small response perturbations.
Takeaways & Limitations
Well-constructed local linear models occupy a “sweet spot” between interpretable modeling and predictive accuracy.
Takeaways & Limitations
The method performed well experimentally, but the authors make no claims about optimality; estimating locality and deciding when to revert to a simpler model remain challenges.

Abstract

from arXiv · show

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.

1 Introduction

Local distillation constructs sparse linear models at individual query points, using a black-box teacher to preserve predictive accuracy while making predictions interpretable as produced. Randomized refits further assess the stability of local feature selections and reveal heterogeneous feature–outcome relationships across observations.

Motivation and method: Local distillation uses a high-performing black-box teacher to guide a local linear student fit at each observation or query point.The approach is motivated by local linear approximation of smooth regression functions.
Empirical illustration: 48%: local distillation improves the global lasso’s prediction squared error on Auto MPG, with PSE 5.59 versus TabPFN 5.25 and global lasso 10.81.Each test-set car receives its own sparse linear model.
Empirical illustration: Local coefficients expose heterogeneity that a global linear model averages away: cylinder count predicts fuel economy among least-efficient cars, whereas displacement is more useful among the most efficient.The global model assigns both coefficients zero.
Motivation and method: The teacher defines locality through similarity of predicted responses, reducing the feature space to a single interpretable axis and avoiding input-space locality’s curse of dimensionality.Its influence also pulls the student toward the teacher’s prediction, with strength estimated from the data; if the teacher does not outperform the student, the method reverts to a global linear fit.
Stability and interpretation: Randomized local distillation uses refits to identify stable features at query points and stable subgroups across observations without recomputing teacher predictions.The framework is supported by stability guarantees for feature selection under small perturbations and can reveal relationships a global linear model cannot express.
Stability and interpretation: Across two case studies, the framework reveals heterogeneous feature–outcome relationships that global linear models cannot express, including relationships relevant to personalized medicine.The paper also evaluates the method across 17 datasets.

2 Local distillation

Local distillation fits a separate sparse linear student at each query point, using teacher-based locality and a teacher prediction anchor. The method estimates teacher influence from data and nearly matches teacher accuracy in the Auto MPG example.

Local distillation fits a separate regularized linear student for each test observation, rather than one global student model.
A pseudo-observation anchors the local student at the teacher’s prediction, with influence scaled by estimated teacher strength and effective sample size.
Teacher-prediction similarity weights define each query point’s local neighborhood by upweighting training observations with similar predicted outcomes.
When the teacher offers no global improvement, the method conservatively returns the global linear fit because localized loss ratios were too variable to be reliable.
Auto MPG test performance was PSE 5.59 for local distillation versus 5.25 for TabPFN and 10.81 for the global lasso.
The pseudo-observation construction supports alternative regularizers and feature expansions while retaining sparse, interpretable local models.

3 Related work

Local distillation differs from existing local explanation and prediction methods in how the teacher defines locality and enters the fit. It combines pointwise sparse modeling with stability guarantees for feature selection.

Unlike kernel-based local methods, local distillation defines neighborhoods through similarity of predicted responses and accepts any accurate regressor as teacher.
Teacher influence is estimated from a cross-validated loss ratio rather than tuned, and the method reverts to the global fit when the teacher adds no improvement.
Each query point receives an elastic-net fit on a weighted, augmented dataset with one active set, rather than a blend of separate models.
Unlike LIME, local distillation fits observed responses, while the teacher defines locality and anchors the fit rather than serving as the regression target.
With lasso regularization, local distillation provides stability guarantees for feature-selection probabilities, unlike the related local methods discussed here.

4 Interpretability through stability

Randomized local distillation assesses whether local feature selections and patterns across query points remain stable under objective perturbations. Selection frequencies support pointwise interpretation, while clustering reveals heterogeneous subgroups.

Randomized local distillation perturbs the optimization objective while holding the query point and feature design fixed, then uses repeated refits to assess stability.
The randomization scale balances stability against predictive accuracy: larger randomization improves the stability bound but can reduce accuracy, so t is chosen within a 5% error tolerance.
Feature-selection frequencies quantify how often each feature is retained across randomized local fits; low frequencies flag uncertain contributions.
In the Auto MPG example, five of seven selected features appeared in over 90% of refits, while cylinders and horsepower appeared in 75% and 70%.
In the gene-expression example, a median local fit selected 94 genes, but typically only 15 were retained in over 90% of refits.
Clustering randomized local fits identifies subgroups with distinct local predictors, including cylinders for least-efficient cars and displacement for most-efficient cars.

5 Theoretical analysis of stability

The paper establishes theoretical stability guarantees for randomized local-distillation feature-selection probabilities under the lasso. The guarantees rely on smooth teacher predictions and regularity conditions on the weighted design and similarity weights.

Theorem 1 bounds how feature-selection probabilities change when training responses change, yielding a uniform stability guarantee under small response perturbations.
The analysis proves smoothness of selection probabilities by convolving randomized lasso selection indicators with Gaussian noise.
The guarantee assumes a smooth teacher prediction with bounded gradient, general-position weighted design, bounded weight concentration, bounded covariates and distillation weight, and non-degenerate restricted design.
The proof combines randomized-lasso stability for perturbed responses with a chain-rule analysis of the teacher prediction’s contribution.
The bound improves when teacher sensitivity decreases with effective sample size, and simplifies when the teacher is fully pretrained and independent of training samples.
The stability bound improves with randomization standard deviation τ, although excessive randomization can reduce predictive accuracy.

6 A high-dimensional example: predicting gene ex-

In a high-dimensional BRCA1 gene-expression study, local distillation matched the teacher’s predictive performance while producing sparse, stable local models. Randomized local fits also revealed patient-specific gene-selection patterns that global linear modeling missed.

Data and setup: The study predicts BRCA1 expression from 17,322 gene predictors using tumor samples from 536 patients.The data were split into 321 training and 215 test patients.
Teacher and student models: The screened TabPFN teacher performed far better than the full-gene alternative and was selected for local distillation.The teacher used the 500 genes most correlated with BRCA1, screened within each training fold, while the student used all 17,322 genes.
Predictive performance: PSE 0.148 (R2 = 0.670) for local distillation versus PSE 0.189 (R2 = 0.579) for the global lasso, a 22% reduction relative to the global lasso.Local distillation slightly exceeded the teacher’s PSE of 0.150 (R2 = 0.666) while retaining transparency.
Sparsity and stability: The median local fit selected 94 genes versus 123 for the global lasso, and stability screening reduced this to 15 stably selected genes.The stability screen used 100 randomized refits and retained a median of 15 genes, with IQR 13–17.
Patient heterogeneity: Clustering 100 randomized local fits revealed patient heterogeneity: FAM107A and KLF14 were selected almost exclusively in different clusters, unlike the global lasso.Both genes had negative local coefficients and zero global-lasso coefficients; this heterogeneity was difficult to surface in a black-box model.
Robustness: Across 100 random 60/40 splits, local distillation improved on the global lasso in 96% of runs and had a sparser median fit in 97 splits.The median PSE reduction was 22%, and the median reduction in selected genes was 15%.

7 Benchmark comparisons

Across 17 regression datasets, the benchmark compares local distillation with global linear, black-box, and other local linear models. Local distillation closely matches its teacher’s predictive performance, especially when the teacher outperforms the global linear student.

Compared methods: The comparison includes global lasso or ridge students, TabPFN or XGBoost teachers, and LOESS or local linear forests.Local-distillation results are labeled by teacher and student regularization choice.
Benchmark design: The benchmark evaluates 17 regression datasets spanning sample sizes n ∈[159, 4177] and feature counts p ∈[5, 51].Datasets come from the UCI Machine Learning Repository and the OpenML-CTR23 regression benchmark.
Benchmark design: Each dataset uses 20 random 80/20 train/test splits with complete-case data, one-hot encoding, and training-set feature normalization.Sample sizes and feature counts are reported in Appendix Table 1.
Results: Local distillation closely matches the predictive performance of its teacher across a wide range of datasets using test R2.Representative results are shown in Figure 5, with median performance and one standard error across splits.
Results: 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

The discussion presents local distillation as a modular framework that combines black-box predictive strength with sparse, interpretable local linear models. It emphasizes teacher-defined locality, randomized stability assessment, theoretical guarantees, broad predictive fidelity, and open choices around distillation strength and model extensions.

Framework: The framework combines transparent, computationally simple linear modeling with predictive performance that can rival black-box models.The discussion characterizes local linear modeling as a modular framework whose weights, penalty, and pseudo-observations can be adapted to the problem.
Empirical performance: Across 17 benchmark datasets and multiple teacher models, local distillation consistently matches or approaches teacher predictive accuracy.The evaluated teachers include TabPFN, TabFM, and XGBoost.
Interpretation and stability: Randomized refits identify stable feature selections at individual test points and stable subgroups across the test dataset.Selection frequencies provide pointwise interpretation, while clustering aggregates randomized fits across observations.
Interpretation and stability: Under lasso penalization, the randomization has theoretical stability guarantees for feature-selection probabilities under small perturbations of training responses.The stability theory is presented as extending beyond this specific local-model construction.
Open choices and extensions: Cross-modal distillation can incorporate external datasets or different feature modalities, with the gene expression example providing one instance.The authors describe distillation across genuinely different modalities as a promising avenue.
Open choices and extensions: The estimated distillation strength performed well empirically, but its optimality is not established and alternative rules for reverting to simpler models remain possible.The authors also note that richer student classes may narrow accuracy gaps when a linear student fails to recover teacher performance.
Framework: Local distillation uses a black-box teacher to define locality, anchor each local linear fit, and support randomized stability assessment.Teacher predictions determine which observations inform each query-point fit and provide anchoring pseudo-observations.

A.1 Proof of main results

The appendix proves smoothness and stability properties for randomized lasso feature-selection probabilities, using KKT and leave-j-out characterizations of the active set.

Selection probabilities are defined as the probability that feature j belongs to the randomized lasso active set.
The map from responses y to the augmented response z is k times differentiable when the teacher prediction is k times differentiable.
Theorem 1 establishes differentiability of the feature-selection probabilities after Gaussian convolution.
Theorem 2 gives a sensitivity bound for feature-selection probabilities under randomized lasso perturbations.
The lasso active set is characterized through leave-j-out scores and KKT conditions under the uniqueness assumption.

A.3 Auxiliary results

Auxiliary results control directional behavior, effective similarity neighborhoods, and the sensitivity of randomized local distillation through augmented-regression bounds.

Proposition 1 provides a uniform featurewise sensitivity bound for randomized local distillation under the stated assumptions.
The augmented response coordinates and resulting prediction map are continuously differentiable, enabling application of the randomized-lasso stability theorem.
A Lipschitz function with nonnegative directional derivative almost everywhere is nondecreasing along that direction.
The effective similarity neighborhood contains observations whose weights are bounded between inverse-effective-sample-size scales.
The auxiliary lemmas bound residual and design quantities in the augmented regression used by local distillation.

B Performance on UCI ML and OpenML datasets

The benchmark evaluation reports test R2 comparisons across 17 datasets and multiple methods, with results summarized over repeated train/test splits.

17 benchmark datasets are evaluated across local linear, ridge, lasso, and teacher-based methods.
The evaluation includes a tabular foundation-model teacher, TabFM, alongside the other teacher-based comparisons.
Test R2 is compared across methods, with medians and 1 SE error bars over 20 train/test splits.
The remaining datasets are presented in continuation figures using the same methods and axes.

C Ablation study

The ablation study separates local distillation’s similarity weights from its teacher prediction prior and finds that their combination is usually strongest.

Local distillation combines similarity weights with a teacher prediction prior as its two teacher-derived components.
The combined components usually provide the best predictive performance.
The ablation compares neither component, weights only, prior only, and both components using a lasso student.
Both components help, and using both is best or tied-best on nearly every dataset.

D Performance with simulated data

The simulated-data study evaluates local distillation with a gradient-boosting teacher and lasso student across noise and dimensionality settings. It generally recovers local structure, while higher noise can cause reversion to the global linear model; selection-probability filtering reduces spurious features while retaining active ones.

Experimental design: The experiment varies n, p, and σ, using 50 local-distillation runs per configuration and 40 query points drawn from the same distribution.The evaluation uses a gradient-boosting teacher and lasso student, measuring coefficient correlations, support recovery, and reversion to the global fit.
Overall performance: Local distillation is generally strong across σ, n, and p, including the high-dimensional regime p > n.Performance is evaluated when the data-generating process is known.
Overall performance: As noise grows, performance degrades and the method eventually reverts to the global linear model.Reverted runs are those with μ̂ ≤ 1.
Local-structure recovery: Correlations with the true local coefficients are 0.63–0.79 for the first two coefficients but 0.14–0.34 for the third.The third coefficient is harder to recover because the similarity definition can group observations with different x1 − x2 values but similar teacher predictions, producing neighborhood-average coefficients.
Selection-probability filtering: Filtering at π̂j > 0.9 reduces false selections from roughly 8–12 to fewer than 2.5 per query point while retaining 86–95% of active features.The filtered set retains only selections whose estimated selection probability exceeds 0.9.