TRAK: Attributing Model Behavior at Scale

Motivation

• Training data influences model errors, biases, and capabilities, motivating attribution methods for explaining predictions, debugging behavior, valuing data, detecting poisoned or mislabeled examples, and curating datasets.Improving training-data quality generally improves resulting model performance [HAE16] [LIN+22].Source: §1 · p. 1
• Existing attribution methods face a computational-efficacy tradeoff: tractable influence and gradient-agreement methods can be unreliable for non-convex models, whereas sampling-based methods require thousands or tens of thousands of model trainings.The paper frames this tradeoff as the motivation for seeking methods that are both scalable and effective in large-scale non-convex settings.Source: §1 · p. 1

Evaluation

• The paper evaluates attribution methods by their ability to make accurate counterfactual predictions, measuring correlation between true model outputs and attribution-derived predictions.The benchmark asks what would happen if the model were trained on a specified subset S′ of the training set and is inspired by datamodeling [IPE+22].Source: §1 · p. 2

Method

• TRAK (Tracing with the Randomly-projected After Kernel) approximates parametric differentiable models with a kernel machine and derives attribution scores in the resulting kernel domain.The approximation can use the empirical neural tangent kernel.Source: §1 · p. 2

Results

• TRAK matches datamodel efficacy on CIFAR-10 and QNLI while running 100-1000x faster, and outpredicts TracIn and influence-function variants at comparable speed.Figure 1 benchmarks efficacy and runtime on CIFAR-10-trained ResNet-9 and QNLI-trained BERT-BASE models; error bars show 95% bootstrap confidence intervals.Source: §1 · p. 2, Figure 1 · p. 2

Applications

• TRAK supports large-scale studies of ImageNet classifiers, CLIP image-text embeddings trained on MS COCO [RKH+21] [LMB+14], and fact-tracing in a 300M-parameter mT5-small model.The mT5 model is finetuned on FTRACE-TREX.Source: §1 · p. 2

2 Motivation and Setup

• Data attribution assigns each training input a real-valued importance score for a specified model output, with models trained by minimizing empirical training loss.The setup uses an ordered training set S, a model output function f, and parameters θ⋆(S); the output function may be the model loss or another setting-dependent quantity.Source: §2 · p. 3

2.1 The linear datamodeling score (LDS)

• LDS evaluates attribution methods by how accurately their scores predict model outputs after training on randomly sampled subsets of the training data.Attribution scores are summed over each subset, and LDS is the Spearman rank correlation between predicted and actual outputs, averaged across test examples; it is quantitative, simple to compute, and task- and modality-independent.Source: §2.1 · pp. 4–5

2.2 An oracle for data attribution

• Sampling-based datamodel attributions accurately predict counterfactual model outputs but may require tens or hundreds of thousands of separately trained models, making them costly at scale.The approach samples training subsets, retrains a model on each, evaluates outputs, and fits a linear predictor; Ilyas et al. [IPE+22] use this oracle-style procedure, which also underlies other sampling-based methods such as empirical influences [FZ20] [CIJ+22] and Shapley values [GZ19; JDW+19].Source: §2.2 · p. 5

2.3 Data attribution methods beyond sampling

• Influence-function and heuristic attribution methods use few trained models but can fail to make meaningful counterfactual predictions in large-scale non-convex settings.Influence methods can accurately approximate linear and other simple models yet perform poorly for deep neural networks [BPF21] [IPE+22] [BNL+22], while representation-similarity [HYH+21] and gradient-agreement [HL22a] methods may look qualitatively compelling without predicting model outputs reliably.Source: §2.3 · pp. 5–6

3 TRAK: Tracing with the Randomly-Projected After Kernel

• TRAK is a data attribution method designed to be effective and scalable for large-scale differentiable models.It derives τTRAK by using a canonical logistic-regression attribution method as a primitive and reducing complex models to that setting.Source: §3 · p. 6

3.1 Warmup: Data attribution for logistic regression

• The logistic-regression warm-up establishes the one-step Newton approximation as TRAK’s canonical, effective attribution primitive.TRAK then uses this convex-setting method to extend attribution beyond logistic regression and other convex settings.Source: §3.1 · p. 7

3.2 TRAK for binary (non-linear) classifiers

• For binary non-linear classifiers, TRAK linearizes model outputs, projects gradients into lower dimensions, estimates influences, ensembles models trained on random subsets, and applies soft-thresholding.The linearization is centered at the final parameters; random projections preserve gradient inner products, ensembling smooths seed variability, and the threshold is selected by cross-validation.Source: §3.2 · pp. 7–8

3.3 Extending to multi-class classification

• For multi-class classification, TRAK uses the correct-class softmax probability to reduce the c-class problem to a single logistic regression instead of solving c^2 binary problems.After linearizing this output around the optimum, it applies the binary procedure’s projection, estimation, ensembling, and sparsification steps.Source: §3.3 · pp. 10–11

3.4 Implementing TRAK

• Algorithm 1 implements multi-class TRAK as an attribution matrix and uses a highly optimized random projector to support very large models.The implemented algorithm returns soft-thresholded attribution scores.Source: §3.4 · p. 11

4.1 Experimental setup

• The evaluation covers ResNet classifiers on CIFAR and ImageNet and pretrained BERT models finetuned on QNLI, comparing TRAK with influence, TracIn, GAS, representation-similarity, empirical-influence, Arnoldi, and datamodel baselines.The study samples 100 half-sized training subsets, trains five models per subset, computes LDS using Spearman correlation, and averages scores over 2,000 validation examples with 95% bootstrap intervals.Source: §4.1 · p. 12
• Computational cost is measured by both single-A100 wall time and the number of trained models, which produce similar conclusions and emphasize training cost for large models.The trained-model metric is hardware- and implementation-agnostic, while per-example gradient computation is generally parallelizable and bounded by model-training time.Source: §4.1 · p. 12

4.2 Results

• Across all evaluated models and datasets, TRAK achieves significantly better efficacy–efficiency tradeoffs than competing methods, with datamodel-level efficacy at an average computational footprint over 100x smaller.Efficacy is measured by LDS, and computational cost by trained-model count; the result is reported across CIFAR-2, CIFAR-10, ImageNet, and QNLI.Source: Figure 2 · p. 13, §4.2 · p. 13
• TRAK identifies training examples that are semantically similar to targets, with helpful examples usually sharing the target class and detracting examples usually belonging to a different class.This qualitative pattern is shown for QNLI, CIFAR-10, and ImageNet examples.Source: Figure 3 · p. 14, §4.2 · p. 14
• TRAK scores correlate with datamodel scores, whereas scores from the other evaluated attribution methods do not.The correlation is defined as the averaged Spearman rank correlation between attribution-score vectors over multiple examples; Table 4 evaluates this on CIFAR-10.Source: Table 4 · p. 15, §4.2 · p. 15
• Using incompletely trained models or multiple checkpoints from one training run can dramatically reduce TRAK’s computational cost without significantly degrading performance.The study varies linearization, projection dimension, Newton-step application, and aggregation, with the two lower-cost optimizations reported in Tables 5 and 6.Source: Table 5 · p. 15, Table 6 · p. 15, §4.2 · p. 15

5.1 Attributing CLIP models

• The CLIP application targets local attribution by identifying training examples that cause a model to learn a specific image-caption association, complementing aggregate representation metrics.Visual comparisons show TRAK can emphasize caption semantics differently from CLIP similarity distance, including helpful examples whose images do not share the target object.Source: §5.1 · p. 16, §5.1.2 · p. 18

5.1.1 Computing TRAK for CLIP

• For CLIP, TRAK converts the symmetric image-text contrastive objective into two classification problems and applies the same logistic-regression influence approximation used for classification.The model output sums the canonical output functions for image-to-caption and caption-to-image matching.Source: §5.1.1 · p. 17

5.1.2 Results

• Removing just 400 TRAK-identified training examples decreases average CLIP similarity distance by 0.36, versus 0.11 for the same number of nearest neighbors, while TracIn has no significant effect.These examples comprise less than 0.5% of the training set; Figure 8 measures the effect after retraining a CLIP model with influential examples removed.Source: §5.1.2 · p. 19, Figure 8 · p. 19

5.2 Fact tracing for large language models (mT5)

• For mT5 fact tracing, TRAK ranks abstracts using a masked-language-model output function formed by summing canonical classification outputs over masked tokens.FTRACE-TREX evaluates whether ranked abstracts retrieve the ground-truth proponents of novel facts using mean reciprocal rank.Source: §5.2 · p. 20

5.2.1 Results and discussion

• TRAK achieves 0.42 MRR on FTRACE-TREX, significantly outperforming TracIn at 0.09 but underperforming BM25 at 0.77.The reported values are specific to this model class and are not directly comparable to prior exact numbers because of a model-class mismatch.Source: §5.2.1 · p. 20
• A 34% accuracy decrease follows removing TRAK-identified abstracts, exceeding BM25’s 10% and ground-truth proponents’ 12%, showing TRAK traces behaviorally important data.This counterfactual result suggests TRAK’s weaker benchmark retrieval reflects FTRACE-TREX’s focus on fact proponents rather than a failure to identify causes of model behavior.Source: §5.2.1 · p. 21, Figure 9 · p. 21

5.3.1 Estimating prediction brittleness

• TRAK scores from 100 models estimate CIFAR-10 prediction brittleness about as effectively as datamodel scores from 50,000 models.This makes TRAK an orders-of-magnitude-faster alternative for the brittleness procedure of Ilyas et al. [IPE+22].Source: §5.3.1 · p. 22, Figure 10 · p. 22

5.3.2 Learning algorithm comparisons

• MODELDIFF using TRAK recovers similar distinguishing features to the datamodel-based analysis while reducing total computational cost by a factor of 100.The case study compares image classifiers trained with and without data augmentation, following Shah et al. [SPI+22].Source: §5.3.2 · p. 23

Training data attribution

• TRAK applies the linearized model directly to the original model and at scale, unlike Achille et al. [AGR+21], whose method modifies model components and targets smaller-scale finetuning.Related estimators use random projections, Fisher matrices [KKG+19], or linearized models [AGR+21], with the latter experiments focusing on smaller-scale finetuning.Source: §6 · p. 23
• TRAK addresses scalability–efficacy tradeoffs differently from Jia et al., who evaluate leave-one-out and Shapley methods using a k-nearest-neighbors approximation over pretrained embeddings.Their estimator offers an alternative to working with the linearized model.Source: §6 · p. 23
• Unlike Hessian-based influence functions [KL17] [KAT+19] [BPF21], TRAK uses positive-semidefinite curvature from the linearized model, avoiding additional regularization beyond random-projection regularization.Related generalized Gauss–Newton [BNL+22] and Fisher [TBG+21] matrices are likewise guaranteed positive semidefinite.Source: §6 · pp. 23–24

Neural tangent kernel

• TRAK is motivated by finite-width empirical NTK work [MLL20] [AGR+21] [Lon21] [ABP22] [MWY+22] [ABS+23] [MGF22], which finds kernel regression can approximate the original model.NTK research also studies generalization [ADH+19], optimization, and robustness [GCL+19], while eNTK applications include active learning [HZK+22] and model selection [DAR+21].Source: §6 · p. 24

Kernel methods and random projections

• TRAK’s random projections build on widely used kernel-approximation techniques, which improve computational efficiency and can explain phenomena such as subspace-based finetuning [MWY+22].Gradient-as-feature ideas also connect to the Fisher kernel and related methods.Source: §6 · p. 24

• TRAK formalizes data attribution and provides an effective, efficiently scalable method.Source: §7 · p. 24
• TRAK is demonstrated on image classifiers trained on CIFAR and ImageNet, language models including BERT and mT5, and the image-text model CLIP.Source: §7 · p. 24
• TRAK requires differentiable models and depends on the suitability of its linear approximation, while experiments and prior NTK success on language modeling suggest it may extend to larger models.This creates an opportunity to apply data attribution in previously untenable domains, such as large generative models.Source: §7 · p. 24

A.1 Datasets and models

• The evaluation covers CIFAR-2/CIFAR-10, ImageNet, QNLI, CLIP on MS COCO, mT5 fact tracing, and LIVING17 MODELDIFF under specified architectures and training setups.CIFAR uses ResNet-9; ImageNet and LIVING17 use ResNet variants; QNLI fine-tunes bert-base-cased; CLIP combines ResNet-50 and a Transformer; fact tracing uses mt5-small.Source: §A.1 · pp. 34–35
• The appendices define task-specific implementation choices, including normalized cosine similarity for CLIP counterfactual evaluation and augmentation/no-augmentation training for LIVING17 MODELDIFF.The CLIP metric rescales raw image-caption correlation using dataset-wide similarity percentiles, while MODELDIFF distinguishes learning algorithms through resulting feature priors.Source: §A.1 · pp. 34–35

A.2 TRAK hyperparameters

• Soft-thresholding is an optional hyperparameter that helps non-binary CIFAR-10 and ImageNet classification, but not CIFAR-2 or QNLI, possibly because attribution vectors are sparser with more classes.For CIFAR-10, k is cross-validated using the same M models when M ≥20; for ImageNet, k = 1000 is chosen because classes contain about 1,300 training examples on average.Source: §A.2 · p. 35

A.3 Baselines

• In the evaluation, approximation-based methods were averaged over multiple independent checkpoints, while datamodel estimators used up to 60,000 CIFAR-2 models and 300,000 CIFAR-10 models trained on different random 50% subsets.Source: §A.3 · pp. 35, 37

A.4 Hardware and wall-time measurements

• All experiments use NVIDIA A100 GPUs with 40GB memory and 12 CPU cores, evaluating attribution cost by total wall-time and total trained-model count.Runtime components include TRAIN_TIME for one model and GRAD_TIME for dataset-wide gradients and random projections; the fast implementation supports dimensions up to 80,000 with little runtime increase.Source: §A.4 · p. 37
• TRAK’s estimated compute is M × (TRAIN_TIME + GRAD_TIME), whereas datamodels and empirical influences cost approximately M × TRAIN_TIME because post-training linear solves are negligible.TRAK computes gradients for every trained model, while datamodel and empirical-influence estimation adds comparatively little cost after training.Source: §A.4 · p. 37
• LISSA influence functions require thousands of Hessian-vector-product iterations per target, so runtime is measured on 50–100 validation examples and extrapolated to the full test set.The Arnoldi influence-function runtime is measured using the authors’ code on translated CIFAR models, while TracIn and GAS use projected gradients across checkpoints.Source: §A.4 · p. 37

B TRAK implementation

• TRAK’s library computes scores through model training or checkpoint acquisition, per-example gradient computation, random projection, and final matrix-based aggregation.Steps (ii)–(iv) are automated using functorch, GPU matrix multiplication or a custom CUDA kernel, and simple matrix operations.Source: B TRAK implementation · p. 38

B.1 Fast random projections on GPU

• The GPU projection kernel avoids storing the projection matrix by generating pseudo-random coefficients as needed, eliminating projection-matrix memory bandwidth and preserving exact reproducibility.Storing a 300-million-weight by 1024-dimensional matrix would require more than 1TB, making block generation memory-throughput limited.Source: §B.1 · p. 38
• Over 200x speed-ups are achieved over block-by-block random projection through an optimized GPU implementation using parallelization, shared memory, and Tensor Cores.The implementation targets NVIDIA GPUs with compute capability 7.0 or newer and supports batched inputs and Gaussian or Rademacher coefficients.Source: §B.1 · p. 38

C.1 The one-step Newton approximation for leave-one-out influence

• TRAK’s leave-one-out influence estimator uses a single Newton step from the global optimum to approximate parameters obtained after removing one training example, then maps parameter change to prediction change.The derivation uses the Sherman–Morrison formula and extends the logistic-regression result to example-dependent bias terms independent of model parameters.Source: §C.1 · pp. 39–40
• The one-step approximation assumes updates converge in one step, but TRAK’s effectiveness across settings suggests this approximation is accurate in practical regimes.The derivation can generalize beyond logistic regression when the training loss is convex in the model output, although output-function choice limits linear-datamodeling performance when outputs are not sufficiently linear.Source: §C.1 · pp. 39–40

C.2 Random projections preserve gradient flow

• Under small-learning-rate gradient flow, model-output evolution depends on gradient vectors only through their inner products, so preserving those products justifies replacing gradients with random projections.The analysis assumes gradients remain unchanged over time and follows prior gradient-flow analyses.Source: §C.2 · p. 40

C.3 Subsampling the training set

• Ensembling attribution scores over multiple models significantly improves TRAK’s performance, motivating training each model on a different random subset.This design choice is investigated further in Appendix E.2.Source: §C.3 · p. 40
• Subsampling different random training subsets lets TRAK average single-subset estimates to approximate empirical influence, which is a difference-in-means estimator [FZ20].Datamodels are described as regularized empirical influences, following Ilyas et al. [IPE+22].Source: §C.3 · p. 41
• Although each subsampled model formally applies only to examples included in its subset, updating TRAK scores for the entire training set produced marginally better estimates.This is an empirical implementation choice that differs from the formal restriction.Source: §C.3 · p. 41

D.1 Correlation distribution

• TRAK scores computed with α = 0.5 remain highly predictive when evaluated at other α values, although correlation degrades somewhat.Figure D.1 compares α = 0.5 scores across α = 0.5 and α = 0.75 evaluations; broader counterfactuals indicate predictivity near α = 1.Source: §C.3 · p. 41, §D.1 · p. 43, Figure D.1 · p. 43

C.4 Linearity and model output function

• Linearity is a latent assumption behind datamodels, Shapley values, influence-function methods, empirical influences, and TracIn [GZ19] [JDW+19] [KL17] [KAT+19] [FZ20] [PLS+20].Datamodels and Shapley values satisfy additivity by construction, while several other methods use first-order or summed influence approximations.Source: §C.4 · p. 41
• TRAK uses task-specific output functions: a logit-equivalent margin for classification and language modeling, and a CLIP-loss-based function for CLIP.The multiclass logit choice is motivated by theoretical and empirical [IPE+22] observations and is well approximated by linear datamodels.Source: §C.4 · p. 42
• Choosing f balances linearizing the model output against preserving nonlinearity: f(z) = z leaves the approximation to capture more nonlinearity, whereas f = L linearizes too much.The former resembles the Hessian-based influence-function approximation, while the latter performs poorly because L is generally highly nonlinear in f.Source: §C.4 · p. 42

D.1 Correlation distribution

• TRAK and datamodel LDS correlations are correlated across examples, with TRAK generally performing better where datamodels perform better.This comparison uses the two estimators shown in Figure D.1.Source: §D.1 · p. 43, Figure D.1 · p. 43

D.2 Table for LDS evaluation

• Sampling-based datamodels and empirical influences can outperform TRAK when given more computation, but at substantially higher computational cost.Table D.2 evaluates predictiveness with linear datamodeling score and efficiency by wall-time or number of trained models, with 95% bootstrap confidence intervals.Source: Table D.2 · p. 44
• Table D.2 compares data-attribution methods jointly by linear datamodeling score and computational efficiency.Efficiency is measured using total computation time on one A100 GPU or the number of trained models required.Source: Table D.2 · p. 44

D.3 TRAK examples

• TRAK examples identify the most helpful and detracting training examples for random ImageNet and QNLI test cases.The ImageNet examples use highest- and lowest-scoring training examples, while the QNLI examples show top attributions for entailment and non-entailment labels.Source: Figure D.3 · p. 45, §D.3 · p. 45, Figure D.4 · p. 46
• For CLIP models trained on MS COCO, TRAK attributions are displayed alongside CLIP similarity distance and TracIn.The examples show corresponding most helpful and most detracting training examples for random test examples.Source: §D.3 · p. 45, Figure D.5 · p. 47

D.4 MODELDIFF with TRAK

• TRAK dramatically accelerates the MODELDIFF algorithm for comparing learning algorithms.The application compares LIVING17 models trained with and without data augmentation and uses TRAK scores in place of datamodel scores for the representation stage.Source: §D.4 · p. 48, Figure D.6 · p. 48
• 100x lower computational cost recovers Shah et al.’s [SPI+22] finding that augmentation-trained models use spider webs as a spurious cue for the spider class.MODELDIFF extracts distinguishing subpopulations before counterfactually testing the inferred feature.Source: Figure D.6 · p. 48

E.1 Dimension of the random projection

• LDS initially rises with projection dimension but eventually falls, while the optimal dimension increases with the number of ensembled models.The tradeoff reflects improved inner-product preservation versus projection cost and a possible regularizing effect at lower dimensions.Source: §E.1 · p. 49, Figure E.1 · p. 50

E.2 Number of models used in the ensemble

• Ensembling more models substantially improves attribution, although a single model still achieves LDS 0.096 on CIFAR-2, outperforming all evaluated prior gradient-based methods.The experiment varies ensemble sizes from 1 to 100.Source: §E.2 · p. 49, Figure E.1 · p. 50

E.3 Proxies for model ensembles in compute-constrained settings

• Multiple checkpoints from the same trajectories achieve comparable LDS to independently trained models, and early-training checkpoints retain efficacy before convergence.These proxies reduce the time required to train several independent models.Source: §E.3 · p. 50, Table E.3 · p. 51

E.4 Role of different terms.

• The reweighting matrix and loss-gradient term substantially improve TRAK, whereas the diagonal R term and leverage score are omitted because they lower or barely affect LDS.Averaging the estimator and loss-gradient term separately, then rescaling, also yields higher LDS.Source: §E.4 · p. 51, Table E.4 · p. 52

E.5 Choice of the kernel

• Using penultimate-layer representations improves over existing gradient-approximation baselines but significantly underperforms full-gradient TRAK, with the gap larger on CIFAR-10 than binary tasks.Last-layer TRAK is nearly equivalent to influence-function approximation, which achieves LDS 0.19 on QNLI.Source: §E.5 · p. 52, Table E.5 · p. 52

E.6 Ensembling vs. Averaging the eNTK

• Averaging eNTKs before applying the estimator yields LDS 0.120 on CIFAR-2, versus 0.499 when averaging per-model estimators, favoring an ensemble of kernel predictors.TRAK follows the latter approach by averaging predictions derived from each kernel.Source: §E.6 · pp. 52–53

E.7 Summary

• TRAK performs best with sufficient model averaging, later-epoch gradients, moderately high projection dimension, reweighting, and full-model gradient features.Computationally cheaper ensemble proxies can also work.Source: §E.7 · p. 53

F Fact Tracing; F.1 The FTRACE-TREX Dataset; F.2 Fine-tuning details; F.3 Computing TRAK for masked language modeling

• FTRACE-TREX evaluates whether attribution ranks proponents among 300 distractors for novel facts, using MRR after masked-language-model fine-tuning and tokenwise TRAK output summation.The study uses a 300M-parameter mt5-small model and treats masked language modeling as separate V-way classification problems.Source: §F.1 · p. 54, §F.2 · p. 54, §F.3 · p. 55

F.4 Counterfactual experiment setup

• The experiment’s collections included TRAK-ranked examples, another set of 500 most-important examples per novel fact, and FTRACE-TREX proponents.The TRAK construction removed 17,914 total training examples, leaving 1,542,539 remaining.Source: §F.4 · p. 55
• Models were finetuned from mt5-small on each counterfactual set, with the pretrained model at 0% accuracy and full-data finetuning at 100% on the selected facts.The selected subset contained 50 novel facts, defined by incorrect pretrained handling and correct finetuned handling.Source: §F.4 · p. 55
• Removing FTRACE-TREX proponents was expected to yield very low accuracy because the corresponding set contained no direct evidence for the selected novel facts.The proponents were defined by the FTRACE-TREX benchmark.Source: §F.4 · p. 55
• Removing TRAK’s most important training examples changed model behavior significantly more than removing BM25-ranked examples or FTRACE-TREX proponents.This counterfactual effect was measured on 50 selected novel facts using average model performance.Source: §F.4 · p. 55

G.1 Further applications of TRAK

• TRAK could support downstream attribution applications including explaining predictions [KL17], dataset cleaning [JDW+19], poison removal [LZL+22], and uncertainty quantification [AV20].Its computational efficiency may expand the settings where these applications are feasible.Source: §G.1 · p. 56
• Future work proposes fact tracing for generative models, using TRAK to trace outputs to sources and improve understanding or trustworthiness in GPT-3 [BMR+20], ChatGPT, and stable diffusion [HJA20] [RBL+22].The authors also identify attribution for generative models as a broader future direction.Source: §G.1 · p. 56
• Aggregated TRAK scores could optimize datasets by selecting coresets, choosing active-learning data, and filtering large datasets while retaining scale benefits without equivalent computational cost.TRAK scores quantify each training example’s impact on a target prediction.Source: §G.1 · p. 56
• Because TRAK is differentiable with respect to inputs, it may enable dataset distillation, potentially leveraging its eNTK in light of NTK results [NNX+21].These input gradients differ from the model-parameter gradients used to compute TRAK.Source: §G.1 · p. 56

Empirical NTK

• Higher-order empirical-NTK approximations beyond TRAK’s linearization could improve estimator accuracy and efficacy [HY20] [BL20].TRAK uses the empirical NTK to approximate the original model, motivating improved analysis of when that approximation is accurate.Source: §G.2 · p. 56

Training dynamics and optimization

• Modeling the feature-learning stage that TRAK omits, alongside mini-batches, momentum, and weight decay, may reduce the gap between true outputs and TRAK predictions [LM20] [LBD+20].The final eNTK represents the later approximately invariant-feature regime, whereas early training may explain remaining limitations.Source: §G.2 · pp. 56–57

Ensembles

• Ensembling significantly improves TRAK efficacy, suggesting independently trained models’ eNTKs capture non-overlapping information and motivating checkpoint-based approximations to full ensembles.Understanding when checkpoints from one training trajectory can approximate a full ensemble could improve efficiency when fine-tuning several models is expensive.Source: §G.2 · p. 57