PaperlayerView exact revision on arXiv

Source-linked AI summary

Datamodels: Predicting Predictions from Training Data

Andrew Ilyas, Sung Min Park, Logan Engstrom, Guillaume Leclerc, Aleksander Madry

arXiv:2202.00622v1stat.MLcs.CVcs.LG

TL;DR

The paper asks how training data and learning algorithms combine to produce predictions. It introduces datamodeling, a framework that uses simple functions of training-example presence to predict model outputs for fixed targets, and shows accurate prediction alongside applications to counterfactuals, similarity, leakage, and embeddings.

  • Problem

    It is difficult to understand how learning algorithms and training data combine to produce model predictions.

  • Method

    Datamodeling fits a surrogate function, including a linear model, that predicts a target-specific model output from the presence of training examples in a subset.

  • Results

    Datamodels accurately predict model outputs and support dataset-counterfactual, brittleness, similarity, leakage, and embedding applications.

  • Takeaways & Limitations

    Datamodeling provides a framework for analyzing a model class in terms of its training data.

  • Takeaways & Limitations

    Prediction quality is limited by estimation error and misspecification, so perfect correspondence across all training-set subsets is not expected.

Abstract

from arXiv · show

We present a conceptual framework, datamodeling, for analyzing the behavior of a model class in terms of the training data. For any fixed "target" example $x$, training set $S$, and learning algorithm, a datamodel is a parameterized function $2^S \to \mathbb{R}$ that for any subset of $S' \subset S$ -- using only information about which examples of $S$ are contained in $S'$ -- predicts the outcome of training a model on $S'$ and evaluating on $x$. Despite the potential complexity of the underlying process being approximated (e.g., end-to-end training and evaluation of deep neural networks), we show that even simple linear datamodels can successfully predict model outputs. We then demonstrate that datamodels give rise to a variety of applications, such as: accurately predicting the effect of dataset counterfactuals; identifying brittle predictions; finding semantically similar examples; quantifying train-test leakage; and embedding data into a well-behaved and feature-rich representation space. Data for this paper (including pre-computed datamodels as well as raw predictions from four million trained deep neural networks) is available at https://github.com/MadryLab/datamodels-data .

1 Introduction

Datamodeling models how training data and a learning algorithm combine to produce predictions for a fixed target example. Simple datamodels can predict outputs and support analyses of model classes, counterfactual sensitivity, similarity, leakage, and embeddings.

  • Motivation: Datamodeling replaces the complex map from training data to predictions with an easier surrogate function for a fixed target example.The surrogate is trained to approximate fA(x; S′), the outcome of training on a subset and evaluating on x.
  • Framework: A datamodel is fit from subset-output pairs generated by training models on sampled subsets and evaluating them on the target.The framework represents subsets by characteristic vectors indicating which original training examples are present.
  • Framework: Datamodeling studies the distribution of models induced by a learning algorithm rather than the behavior of one specific trained model.This distinction matters because neural-network behavior can vary substantially with training randomness.
  • Framework: Datamodels are target-specific and can be computed for collections of targets when the training set and learning algorithm are fixed.Targets may be test examples, synthetic examples, or examples from the training set itself.
  • Applications: The framework predicts model outputs with Pearson correlation r > 0.99 and supports counterfactual, similarity, leakage, and embedding applications.For CIFAR-10, 50% of test predictions can be flipped by removing fewer than 200 target-specific training points, or by relabeling 35 points.

2 Constructing (linear) datamodels

The paper constructs linear datamodels by sampling training subsets, recording model outputs on a target, and fitting a regularized surrogate to subset characteristic vectors. It uses random α-fraction subsets and continuous classifier outputs rather than binary correctness.

  • Construction pipeline: The construction pipeline samples subsets, trains and evaluates models, collects subset-output pairs, splits them, and fits the surrogate with validation-based model selection.The collected data are divided into training, validation, and test sets before parameter estimation.
  • Surrogate choice: The paper uses a simple linear surrogate mapping each subset’s characteristic vector to the predicted model output.A bias term may be fitted, although the paper often omits it from notation.
  • Subset distribution: Training subsets are sampled as random α-fraction subsets, reducing the sampling-distribution choice to the subsampling fraction α.The most useful α can vary by setting, so the authors estimate datamodels for several values.
  • Output choice: For classifiers, correctness is suboptimal because it discards confidence information and is poorly matched to least-squares fitting.The authors therefore consider continuous outputs such as cross-entropy loss, correct-label confidence, or a correct-class margin.
  • Loss and regularization: An ℓ1 regularization term incorporates a sparsity prior because a target prediction may not depend on every training example.The regularization parameter λ can be selected separately for each target example using cross-validation.

3 Accurately predicting outputs with datamodels

The paper constructs datamodels for CIFAR-10 and FMoW that predict model outputs for specific examples from training subsets. Linear datamodels show near-perfect correspondence with neural-network margins on unseen subsets, while regularization is important for generalization.

  • Evaluation: The evaluation tests datamodel predictions on fresh subsets sampled from the same distribution as the training subsets.Held-out subset-output pairs are used to measure on-distribution prediction quality.
  • Linear datamodel performance: Linear datamodels show near-perfect correspondence with expected correct-class margins from neural networks trained on random subsets.Ground-truth margins are estimated by training 100 models on each subset and averaging their outputs.
  • Sample complexity: The datamodel training-set size m is evaluated using average mean-squared error against an optimal error level determined by inherent training noise.Figure 6 plots average MSE as a function of m, with OPT representing the irreducible variance component.
  • Regularization: Setting λ = 0 produces overfit datamodels that perform well on estimation subsets but poorly on new subsets.With m = 300,000 models, unregularized datamodels have higher MSE than models trained with optimal λ using m = 10,000.

4 Leveraging datamodels

Datamodels extend from predicting model outputs on random training subsets to estimating counterfactual effects, identifying brittle predictions, finding similar training examples, and detecting train-test leakage.

  • Counterfactual prediction: Datamodels predict outputs on arbitrary, including out-of-distribution, training subsets, enabling counterfactual estimates and brittleness analysis.They were originally trained on random α-fraction subsets but can estimate effects for selected removals.
  • Measuring brittleness: Over 90% of 300 CIFAR-10 examples received certified upper bounds on their true data support.Certification used retraining after excluding the estimated support and checking average misclassification.
  • Measuring brittleness: 20% of CIFAR-10 test images had estimated data support below 40 training images, or 0.08% of the training set.Around half had estimated support of 250 images or less, corresponding to 0.4% of the training set.
  • Measuring brittleness: Datamodel-guided estimates required 40 removals to misclassify 20% of test examples, versus over 600 for the best baseline.Every tested baseline produced looser data-support estimates.
  • Predicting data counterfactuals: Counterfactual predictions correlated strongly with true effects, with Pearson r = 0.96 for CIFAR-10 and r = 0.90 for FMoW.Across factors of variation, Spearman correlations were ρ = 0.98 and ρ = 0.94 for CIFAR-10 and FMoW respectively.
  • Train-test similarity: Highest-magnitude datamodel coordinates identify visually similar training examples and support train-test leakage detection across CIFAR-10 and FMoW.Datamodel-derived leakage candidates consistently outperform nearest neighbors from a pretrained representation space.

5 Discussion: The role of the subsampling fraction α

The subsampling fraction α controls which training-data relationships a datamodel reveals: smaller α surfaces broader groups, while larger α isolates smaller groups.

  • Granularity of relationships: α determines the granularity of datamodel relationships, with large α identifying local relationships and small α revealing broader ones.Figure 15 illustrates this trade-off using a nearest-neighbor classifier in two dimensions.
  • Why α matters: Decreasing α lets datamodels observe when entire groups are absent and distribute the group effect across their constituent examples.This helps expose relationships driven by larger groups of training examples.
  • Why α matters: For atypical examples whose predictions rely on a few training images, larger α can better isolate the relevant training points.Small α may spread the observed variation across a large group of non-present images.
  • Practical implication: Larger α is better suited to predictions driven by smaller groups, whereas smaller α is better suited to predictions driven by larger groups.The appendix’s underdetermined linear-regression analysis reports the same pattern for overlapping subpopulations.

6 Related work

Datamodeling relates to influence estimation, surrogate modeling, robustness, data selection, memorization, and representation learning, while serving a distinct goal: predicting model behavior from training-data subsets.

  • Influence estimation: Influence estimation can require training C · d models, motivating approximation schemes that reuse models to improve sample efficiency.Using random subsets of half the training set maximizes model reuse, while subset size also trades off accuracy against efficiency.
  • Influence estimation: Datamodels construct an explicit statistical model of predictions from training data, whereas influence functions measure the counterfactual effect of individual training points.Where quantitative comparisons are available, datamodels predict model behavior better.
  • Surrogate models: Unlike pixel-space surrogates that explain predictions locally around a fixed model input, datamodels model the mapping from training-data subsets to predictions.Datamodels also provide a clear notion of missingness because training examples can be removed directly.
  • Data evaluation: Datamodels differ from active-learning and coreset surrogates, which use simpler models as proxies for evaluating data points.Their broader framing concerns predictions across training-data subsets rather than only data evaluation.
  • Learning and memorization: Datamodels may help study how memorization, generalization, differential privacy, and pruning affect model behavior across subgroups.Prior work reports poorer or disparate generalization in some privacy and pruning settings.
  • Robustness and data dependence: Datamodeling connects to robustness research because it studies how conclusions change when data points or subsets are omitted.Related work includes methods for estimating influential subsets in econometric analyses.

7 Future work

The paper identifies several directions for improving datamodel estimation and using datamodels to study generalization, interpretation, and data exploration. Key open questions concern correlated outputs, uncertainty, subset sampling, priors, linearity, the role of α, and extending applications.

  • Future work: The current instantiation provides good behavior predictors and applications but remains basic, leaving substantial room for improvement.The authors frame these limitations as opportunities for further methodological and conceptual development.
  • Improving datamodel estimation: Correlated outputs across test inputs reduce estimation independence, motivating correlation-aware methods to improve calibration and sample efficiency.The same models estimate datamodels for an entire test set, creating correlated parameter-estimation errors.
  • Improving datamodel estimation: Confidence intervals could replace point estimates with uncertainty ranges and help determine whether training inputs significantly predict particular test inputs.The paper notes that post-selection inference is also needed because data-driven coefficient selection can miscalibrate classical confidence intervals.
  • Improving datamodel estimation: Better subset distributions and domain-specific priors could improve datamodels beyond random α-subsets and least-squares regression with ℓ1 regularization.The authors specifically suggest structured sparsity as one way to incorporate additional data structure.
  • Studying generalization: Simple linear datamodels accurately predict outputs, including out-of-distribution counterfactual datasets, but why this works for deep networks remains an open question.The authors connect this question to training dynamics and implicit biases in overparameterized training.
  • Studying generalization: Datamodel sparsity may quantify how memorized a test input is, while α changes the granularity of correlations detected and lacks a crisp theoretical characterization.The paper also proposes extending datamodels for counterfactual interpretation and user-friendly data exploration tools.

8 Conclusion

The paper presents datamodeling as a framework that models training outputs through the presence of individual training data points. A simple linear instantiation predicts model outputs accurately and supports multiple applications.

  • Conclusion: Datamodeling represents model-training outputs as a simple function of which training data points are present.The framework is applied to fixed target examples and training processes.
  • Conclusion: A simple linear datamodel accurately predicts model outputs and facilitates a variety of applications.The conclusion summarizes the framework's predictive and application-oriented value.

Appendices

The appendices describe the datamodeling procedure, counterfactual evaluation, and a simulated analysis of how the subsampling fraction α affects detected correlations. They also specify the linear-regression setup and feature-frequency evaluation.

  • Framework and algorithms: The framework samples subsets, trains models on them, evaluates a target example, and records subset indicators with model outputs as datamodel training pairs.Each subset is represented by a binary characteristic vector indicating which training examples it contains.
  • Counterfactual evaluation: Counterfactual evaluation compares the full-training-set margin with margins after removing the top-k datamodel-weighted training examples.The procedure averages full-set outputs over 100 trials and removal outputs over 20 trials for k values from 20 to 320.
  • Subsampling fraction α: Higher α is intended to detect local effects from small example groups, whereas lower α is intended to detect global effects from larger groups.The appendix uses a simulated linear model to test this intuition across feature frequencies.
  • Simulation setup: The simulation uses an underdetermined linear regression with binary features, Bernoulli feature frequencies, noisy labels, and an equally sized held-out evaluation set.The reported setup uses d = 150 features and n = 125 training samples.
  • Simulation setup: Datamodels are fit from one million subsampled training subsets and their evaluations using ordinary least squares.The OLS fit is separate from the linear regression used as the original model class.
  • Feature-frequency analysis: The analysis tests whether α-dependent datamodels predict the effects of subpopulation removal across feature frequencies.Actual effects compare predictions with and without a feature-defined subpopulation, while predicted effects sum datamodel weights over that subpopulation.
  • Feature-frequency analysis: Feature-correlation evaluation computes Pearson correlation between actual and datamodel-predicted counterfactual effects across features and test examples.The procedure evaluates features at a specified frequency p and repeats the analysis for different α values.
  • Feature-frequency analysis: Figure B.1 examines how α affects feature detection by correlating actual and predicted effects across feature frequencies.The figure uses a simulated linear model with binary vectors and overlapping feature-defined subpopulations.

C Selecting Output Function to Model

This appendix selects an output function for datamodel fitting by considering output variation and residual normality. Correct-class margins are chosen because their residual tests are consistent with normally distributed outputs.

  • Selection heuristic: The selection heuristic addresses output saturation and seeks output distributions more compatible with least-squares datamodel fitting.The authors state that the heuristic is neither sufficient nor necessary for least-squares regression to work.
  • Output saturation: Over 30% of the CIFAR-10 test set is always correct or always incorrect across trained models, making correctness-based datamodel estimation impossible there.Insufficient output variation is the most pronounced for model correctness, but the issue is not unique to correctness.
  • Residual normality: In the idealized least-squares model, outputs are modeled as normally distributed around a subset-dependent linear mean with unknown variance.The paper uses residual normality as an empirical signal for comparing candidate output functions.
  • Residual normality: The evaluation retrains 200 models on one fixed subset and tests output normality for each CIFAR-10 test image.This repeated-retraining setup provides multiple response values for the same subset and target image.
  • Output choice: Correct-class margin is selected because its residual-test p-values are nearly uniformly distributed, unlike the other candidate outputs.The tested alternatives are correctness, correct-class confidence, cross-entropy loss, and correct-class margin.
  • Figure C.1: Figure C.1 compares residual-normality test p-value distributions for four output functions.The key comparison is whether p-values are broadly uniform or concentrated near rejection.
  • Output choice: Correct-class margin is the correct-class logit minus the highest incorrect-class logit, and its sign indicates classification correctness.Unlike bounded correctness or confidence measures, the margin is unbounded by definition.

D.1 Datasets

The experiments use CIFAR-10 and FMoW with specified ResNet model classes, dataset restrictions, and training settings. Datamodel estimation relies on large-scale subset-output matrices and regularized linear regression, supported by GPU and data-loading optimizations.

  • Datasets: The study analyzes deep image classifiers on CIFAR-10 and FMoW.FMoW experiments use only in-distribution train/test splits and examples from 2012.
  • Model classes: CIFAR-10 uses a fast-training ResNet-9 variant, while FMoW uses standard ResNet-18 without data augmentation or ImageNet pretraining.The FMoW setup accepts lower accuracy to isolate the dataset’s role.
  • Datamodel estimation: Datamodel estimation uses training-subset characteristic vectors and model outputs rather than storing model checkpoints.These quantities are stored as matrices with rows for model instances and columns for examples.
  • Datamodel estimation: The estimation problem is large: at α = 50%, LASSO uses a 50,000 × 300,000 covariate matrix, while α = 10% requires five times more models.The α = 50% matrix corresponds to about 60GB of data.
  • Computation: The implementation uses SAGA-based GPU regression, simultaneous outputs across test examples, custom CUDA optimizations, and fast data loading.FFCV enables throughput above 5,000 CIFAR-10 models per day per GPU.
  • Experimental protocol: The regularization path contains 100 logarithmically spaced λ values, with separate training and validation samples for each regression.The maximum λ is dataset- and split-specific.

E.2 Omitted results

The omitted results describe regularization, sparsity, counterfactual evaluation, baseline construction, and transfer across model architectures. Together they specify how datamodel quality and comparison procedures were assessed beyond the main text.

  • Regularization: Figure E.2 compares in-sample and out-of-sample MSE for CIFAR-10 datamodels across regularization values λ.The datamodels are estimated by optimizing a regularized least-squares objective.
  • Sparsity: Higher α produces less sparse datamodels, consistent with higher α capturing relationships driven by smaller image groups.Figure E.1 reports sparsity distributions over CIFAR-10 test examples.
  • Additional results: The omitted figures reproduce main-text results for FMoW and for CIFAR-10 across all α values.Figures E.3 and E.4 extend the corresponding comparisons.
  • Counterfactual protocol: Counterfactual evaluations sample 300 class-balanced CIFAR-10 test images and 100 FMoW test images, using specified removal sizes and repeated trials.CIFAR-10 baselines use a 100-image subset, and T = 20 trials are used for the main counterfactuals.
  • Baselines: Baselines select examples by representation distance, influence estimates, or random same-class ordering.Representation distance uses penultimate-layer ℓ2 distance; influence baselines include influence functions and TracIn.
  • Baselines: Averaging representation distances over up to 1,000 models produces no discernible change in counterfactual effect size.The representation baseline therefore gains no additional signal from this averaging procedure.

F.3 Data support estimation

The data-support procedure uses datamodel-selected top-k removals, counterfactual retraining, and monotonicity-based estimation, followed by verification and comparison with alternative selection methods. Additional experiments examine mislabeling, architecture transfer, zero-weight groups, and leakage-related examples.

  • Data support estimation: For each target, the procedure selects top-k training examples by datamodel weight and evaluates models trained after removing them.The candidate sizes are k ∈ {10, 20, 40, 80, 160, 320, 640, 1280}.
  • Verification: Counterfactual evaluations train 20 models per removed subset to reduce variance, but direct verification found only 67% of images misclassified at the initial estimate.The procedure therefore evaluates additional removals to establish sufficient coverage.
  • Estimation methodology: The estimate assumes the expected margin decreases monotonically with the number of removed examples and locates its zero crossing.The resulting zero-crossing estimate is an upper bound on SUPPORT(x).
  • Heuristic estimation: A cheaper heuristic estimates support from the smallest k whose largest datamodel weights exceed the target’s average margin.The heuristic is compared with counterfactual-based estimates.
  • Mislabeling brittleness: Mislabeling selected training examples provides a stronger perturbation than removing them, requiring fewer points to alter predictions.The study estimates this brittleness by relabeling top-weighted examples toward a high-logit incorrect class.
  • Comparisons and transfer: Datamodels identify larger counterfactual effects than the baselines, while TracIn performs best among the baselines.Original datamodels remain accurate but degrade somewhat when evaluated on ResNet-18 models.
  • Zero-weight groups: Removing randomly chosen zero-weight examples has negligible impact on actual margins, matching datamodel predictions.This result is reported for α = 0.5 datamodels.

I.2 Omitted results

The omitted results extend the paper’s clustering analyses across α values and additional classes. The figures specifically include horse-class clusters and clusters for eight other classes.

  • Spectral clustering: Figure I.1 compares horse-class top clusters across α = 10%, 20%, 50%, and 75%.The caption lists the corresponding panel placements for these α values.
  • Spectral clustering: Figure I.2 presents omitted spectral-clustering results for eight classes not shown in the main paper.The main-paper comparison is identified as Figure 13.

J.1 CIFAR

The CIFAR-10 experiments use PCA to analyze datamodel embeddings, test whether principal directions reflect model behavior, and compare them with representation embeddings. Removing training examples selected by datamodel principal components produces targeted, predictable effects on aligned test examples while having little aggregate impact.

  • Effective dimensionality: Datamodel embeddings require up to 500 principal components to capture half their variance, whereas network-layer embeddings capture 95% with the first 10.The comparison uses a ResNet-18 representation embedding and datamodel embeddings.
  • Analyzing model-faithfulness: Removing the most positive principal-component coordinates decreases margins for highly positive test projections and increases margins for highly negative projections.The evaluation groups held-out examples by their projections onto each principal component.
  • Analyzing model-faithfulness: Removing the most negative coordinates produces the opposite margin changes for positive- and negative-projection test examples.The sign of the removed training-example coordinates determines the direction of the effect.
  • Analyzing model-faithfulness: Larger removed sets produce effects of greater magnitude, while removing principal components has negligible impact on the aggregate test set.The latter pattern indicates that the impacts of different principal components are roughly orthogonal.
  • Analyzing model-faithfulness: Datamodels accurately predict the counterfactual effects of removing the selected principal-component groups.The ground-truth effects are obtained by retraining models and are compared with datamodel predictions.
  • Representation baseline: Principal components from representation embeddings appear less coherent than those obtained from datamodel embeddings.The representation baseline visualizes images with the highest and lowest projections along top components.

K.1 Proof of Lemma 1

This section characterizes empirical influences as rescaled linear datamodel estimates and evaluates how estimation choices affect their ability to predict model outputs. The analysis emphasizes that explicit datamodel training, larger sample sizes, and margin outputs improve prediction and generalization.

  • Proof of Lemma 1: Lemma 1 reinterprets empirical influences as rescaled weights of a linear datamodel.The proof constructs a binary mask matrix for sampled training subsets and analyzes its limiting covariance and count structure.
  • Evaluating influence estimates as datamodels: Explicit datamodels differ from empirical influences through their estimation algorithm, scale, and output type.The comparison concerns squared-error fitting with sparsity and bias, the number of sampled models, and prediction targets such as margins versus correctness.
  • Evaluating influence estimates as datamodels: Explicit datamodels minimize squared error and use sparsity and bias terms that may help generalization.These design choices are contrasted with empirical influence estimation rather than established as independently causal effects.
  • Evaluating influence estimates as datamodels: The framework evaluates predicted versus true outputs using Spearman correlation and, where appropriate, mean squared error and AUC.AUC treats the influence-based score as a continuous predictor of binary correctness.
  • Testing Lemma 1 empirically: Empirical influences can predict correctness on the subsets used for estimation but show significant generalization error on held-out subsets.Train conditional prediction distributions separate correct and incorrect outcomes better than held-out distributions.
  • View of empirical influences as a Taylor approximation: Empirical influences provide a first-order Taylor approximation of the multilinear extension of the target set function.The multilinear extension interprets fractional inputs as inclusion probabilities, and its derivatives correspond approximately to subsampled influences.
Loading 2202.00622v1…