DRIVE: Fine-tuning via Data Contribution- and Diversity-aware Weighting with Prior Regularization

What it claims

DRIVE, a NeurIPS 2026 paper (arXiv, October 2026), weights each fine-tuning sample by solving a small concave problem at every training step. The objective adds three terms: a contribution score mixing anti-forgetting and task adaptation, η·log det of the weighted feature covariance for diversity, and −τ·KL(w‖p) toward a prior p. The authors describe the log-det term as the first differentiable diversity term for sample weighting. Code is at https://github.com/luciabcd/DRIVE.

How it works

Both contribution parts are approximated by losses, because the setting rules out pre-training data and held-out validation sets. The anti-forgetting score aᵢ is the sample's loss under the frozen pre-trained model. The adaptation score tᵢ is its loss under the current model. Contribution is cᵢ = λaᵢ + (1−λ)tᵢ.

The default schedule decays λ linearly from 1 to 0 over training. Early steps favour retention and late steps favour the new task, so no per-task tuning is needed. A task-specific constant λ is the stated alternative.

The diversity gradient for sample i is its statistical leverage, x̄ᵢᵀΣ_w⁻¹x̄ᵢ. Samples whose features the rest of the set explains poorly gain weight. The weights have a closed form, wᵢ ∝ pᵢ·exp((cᵢ + η·leverageᵢ)/τ). Because Σ_w depends on w, the update runs as a fixed-point iteration from uniform weights, which the authors report converges in a few steps.

The objective is strictly concave for τ > 0. Convergence is linear under a standard contraction condition.

flowchart LR
  A[Loss under pre-trained model] --> C[Contribution c_i]
  B[Loss under current model] --> C
  L[λ schedule 1→0] --> C
  F[Sketched features, buffer] --> S[Covariance Σ_w]
  C --> U[Weight update]
  S --> U
  P[Prior p, uniform default] --> U
  U -->|not converged| S
  U -->|converged| W[Weighted batch loss]

Setting η = 0 removes diversity. With η = 0, λ = 1 reduces to FLOW-style anti-forgetting weighting and λ = 0 to loss-based task weighting.

How to use it

Each step needs one extra forward pass through a frozen copy of the pre-trained model. Features come from the penultimate layer for vision models and the final block for LLMs. They are projected to k = 128 dimensions with one fixed Gaussian matrix, which cuts inversion cost from O(d³) to O(k³).

A single batch may leave Σ_w singular, so the batch is padded with a FIFO buffer of the 128 most recent samples. The buffer also supplies the centring mean. A ridge εI is added before inversion. Weights are then restricted to the batch and renormalised to form the weighted loss.

The main experiments use a uniform prior. The ablation (Sec. 5.3) reports that an entropy-informed prior helps.

Evidence

The loss-based scores correlate with validation-based Shapley contributions at Spearman ρ = 0.83. They match gradient-based estimates within 0.3% accuracy and run about 2× faster (Appendices A.2–A.3). Ensembling several random projections gives negligible gains (Appendix F.3).

The authors also report that the determinant term beats the clustering method DEFT-UCS (Sec. 5.2). The fetched text ends before the results section, so that comparison carries no visible numbers. The headline claims of better accuracy, faster convergence and robustness across vision models and LLMs also lack benchmarks, model names and baselines in the available text.

Limits

The guarantees are proved for linear models and for a local nonlinear regime under local smoothness and linearisation, not for general deep fine-tuning. The setting is single-task fine-tuning without task sequences, so continual learning is out of scope. The loss proxies assume cross-entropy objectives. The method adds hyperparameters η, τ, ε, k and the λ schedule. Sensitivity to k is in Appendix F.5, which is not in the fetched text.