PaperlayerView exact revision on arXiv

Source-linked AI summary

WARP: Weight-Space Analysis for Recovering Training Data Portfolios

Tzu-Heng Huang, Aditya Goyal, John Cooper, Frederic Sala

arXiv:2607.01686v1cs.LG

TL;DR

Training data recipes for released foundation models are rarely disclosed, limiting visibility into their domain mixtures. WARP infers those mixtures from model weights by simulating the missing training trajectory, achieving MAE as low as 0.046 on BERT and 0.104 on GPT-2 while outperforming membership-inference baselines and an oracle with the true trajectory.

  • Problem

    Training datasets and mixing recipes behind publicly released models are often undisclosed, limiting visibility into their training-data portfolios.

  • Method

    WARP uses model merging to simulate missing training trajectories, extracts weight-space geometric footprints, and maps them to domain proportions.

  • Results

    MAE reaches 0.046 on BERT and 0.104 on GPT-2, outperforming membership-inference baselines and an oracle with access to the true trajectory.

  • Takeaways & Limitations

    WARP demonstrates the feasibility of recovering training distributions from weight-space geometry, supporting greater transparency when model weights are shared but data recipes are not.

  • Takeaways & Limitations

    The parameter-free softmax readout cannot exploit how domain influence varies across early and late learning because it collapses the temporal dimension.

Abstract

from arXiv · show

Foundation models are routinely released to the public, yet the data recipes used to train them -- such as domain mixture weights that determine how different sources are sampled -- are rarely disclosed. This creates an access asymmetry: researchers study the resulting models but lack visibility into the training distribution that produces them. Prior works for inferring training data, such as membership inference, detect at the level of individual samples and thus cannot characterize the global composition of the training corpus. We introduce WARP, a framework that recovers a fine-tuned model's training mixtures directly from its released weights. WARP interpolates between the base and fine-tuned models using model merging, generating pseudo-checkpoints that approximate the missing training trajectory and expose a geometric footprint of the training data in the weight space. From these simulated footprints, WARP extracts geometric features and maps them to domain proportions using either a parameter-free softmax readout or an MLP projector trained on synthetic mixtures. In controlled experiments with BERT and GPT-2, WARP recovers domain mixtures with an average MAE as low as 0.046 and 0.104 respectively, outperforming membership inference and a variant with access to the true training trajectory.

1 Introduction

WARP addresses the hidden-data problem behind released foundation models by inferring domain mixtures from model weights rather than individual training examples. It reconstructs a usable training footprint between base and fine-tuned endpoints and maps that geometry to mixture proportions.

  • The access asymmetry: Released models expose final and base weights while their proprietary data mixtures and training paths remain hidden.This opacity can cause unintended capability drift during continued fine-tuning and complicate data-contamination auditing.
  • The access asymmetry: Membership inference detects whether particular samples were used but cannot characterize the corpus-wide proportions of training domains.Existing mixture-optimization methods generally work forward from data to model, making them poorly suited to already-released weights.
  • WARP: WARP uses model merging to interpolate between base and fine-tuned models, creating pseudo-checkpoints that approximate the missing training trajectory.Geometric features from these simulated paths encode a distributional footprint that can be read out as domain proportions.
  • WARP: WARP offers either a parameter-free softmax readout for coarse estimates or an MLP projector trained on synthetic mixtures for stronger prediction.Both approaches operate with a small probing dataset while the corpus and true trajectory remain withheld.
  • Results: 0.046 MAE on BERT and 0.104 on GPT-2 are WARP’s lowest reported recovery errors in controlled experiments.The method outperformed sample-level membership-inference baselines and a variant with access to the true training trajectory across forty trials.
  • Results: WARP remains accurate on early-stop, converged, and overtrained checkpoints, indicating robustness across training recipes.This matches post-training practice in which released models are often pushed beyond compute-optimal budgets.

2 Related Work

WARP connects domain-mixture optimization, training-data inference, and model merging, but changes the direction of analysis from data-to-model prediction to weight-to-portfolio recovery. Its distinctive use of merging is analytical: it exposes training-history structure instead of merely combining capabilities.

  • Domain-mixture optimization: Domain-mixture optimization studies how static mixtures or dynamic reweighting affect foundation-model training and downstream performance.This literature traditionally seeks an ideal portfolio before or during training rather than recovering one from released weights.
  • Training-data inference: Training-data inference is motivated by privacy concerns and by extracting data insights for future curation, with membership inference as a common formulation.Membership inference focuses on whether a particular sample appeared in training, not on global domain proportions.
  • Model merging: Model merging combines models through weight-space operations such as interpolation, stitching, or averaging, usually to preserve or combine capabilities.WARP departs from that objective by using merging to simulate an unobserved training trajectory and infer its data mixture.

3 Framework

WARP estimates hidden domain mixtures from a released fine-tuned model and its base model, without the original dataset or intermediate checkpoints. It simulates the missing trajectory, extracts domain-aligned weight-space geometry, and maps that footprint to mixture proportions.

  • Setup: WARP assumes access to the base and fine-tuned weights plus a small labeled probing set sampled from the relevant source domains.The probing set need not match the hidden training mixture; it only needs coverage of the domains of interest.
  • Simulating the Training Trajectory: Model merging interpolates the base and fine-tuned models to create pseudo-checkpoints that proxy intermediate training states.The mixing schedule α_t acts as normalized training progress, from states near θbase to states near θref.
  • Distilling Geometric Footprints: At each pseudo-checkpoint, WARP computes Mimic Scores by projecting probing samples’ negative gradients onto the direction toward the fine-tuned model.The scores estimate how strongly samples from each domain would pull the model along the observed fine-tuning direction, followed by within-step min-max normalization.
  • Distilling Geometric Footprints: Averaging normalized sample scores within domains and stacking them across checkpoints produces a K × T geometric footprint matrix.Each row tracks one domain’s alignment as the simulated model evolves from the base toward the fine-tuned endpoint.
  • Mapping from Geometry to Mixture: WARP maps the footprint to a predicted mixture using either a parameter-free softmax readout or a supervised projector trained on synthetic mixture-footprint pairs.Synthetic pairs come from short fine-tuning runs on datasets sampled under randomly drawn domain mixtures.
  • Mapping from Geometry to Mixture: The unsupervised readout is fast and label-free but loses temporal information by averaging domain scores across pseudo-checkpoints.Its temperature τ controls the sharpness of the predicted distribution.

4 Experiments

WARP is evaluated as a weight-space method for recovering known domain mixtures across architectures, datasets, baselines, and training stages.

  • Readout comparison: The parameter-free softmax readout already outperforms all baselines without requiring a learned projector or synthetic training set.It closes most of the gap to the supervised variant.
  • Baselines: WARP’s best supervised variant outperforms random guessing, centroid guessing, and sample-level membership inference.The baselines include uninformed mixture predictors and aggregated per-domain membership rates.
  • Trajectory simulation: On BERT, supervised LERP and TIES reach 0.048 avg. MAE versus 0.080 for the oracle on real checkpoints.On GPT-2-Small, LERP reaches 0.117 versus 0.138 for the oracle, suggesting smoother simulated paths can help recovery.
  • Robustness: WARP remains accurate on early-stop, converged, and overtrained GPT-2-Small checkpoints, with supervised MLP consistently outperforming baselines.Table 2 uses AG News and TIES merging to compare these training stages.

5 Conclusion

WARP recovers a fine-tuned model’s training domain mixture from released weights by simulating the missing trajectory and reading its geometric footprint. Controlled BERT and GPT-2 experiments show accurate recovery beyond membership inference, supporting greater transparency when data recipes remain undisclosed.

  • Method: WARP uses model merging to simulate the missing training trajectory between base and fine-tuned models.The simulated path yields a geometric footprint that is mapped to domain proportions.
  • Results: 0.046 MAE on BERT and 0.104 on GPT-2 are WARP’s lowest reported recovery errors.These results outperform membership inference baselines and an oracle with access to the true trajectory.
  • Implication: WARP demonstrates the feasibility of recovering training distributions from weight-space geometry when model weights are public but data recipes are not.The evaluation is controlled and uses known mixtures.

Appendix Roadmap

The appendix roadmap organizes implementation details around the full WARP pipeline, its experimental setup, and future work.

  • Appendix A: Appendix A summarizes the complete WARP framework in pseudocode.The algorithm covers trajectory simulation, footprint distillation, and mixture readout.
  • Appendix B–C: Appendix B documents experimental setup and training configurations, while Appendix C outlines future-work directions.The roadmap does not provide the future-work details themselves.
  • Framework: WARP has three stages: model merging, Mimic Score footprint extraction, and mixture prediction.Prediction uses either an unsupervised softmax readout or a supervised projector trained on synthetic pairs.
  • Algorithm inputs: Algorithm 1 requires base and fine-tuned models, a K-domain data source, a merging operator, and a schedule.It returns an estimated domain mixture ˆπ in the simplex ∆K−1.
  • Pipeline: The procedure initializes a labeled probing dataset, creates pseudo-checkpoints, computes Mimic Scores, and maps the resulting footprint to a mixture.The supervised branch samples synthetic mixtures and fine-tunes reference models to train its projector.

B Experimental Details

The experiments construct controlled reference models with known mixtures, use balanced probing data, train a supervised projector, and report mixture-recovery errors.

  • Reference models: Each reference model is fine-tuned for 9 epochs on 5,000 examples sampled according to a known mixture π⋆.BERT and GPT-2-SMALL start from publicly released pretrained weights used as θbase.
  • Table 3: Table 3 compares MSE recovery error across BERT and GPT-2-Small datasets, averaged over 40 trials, with lower values preferred.Best and second-best results are marked within each model block.
  • Probing data: WARP’s probing dataset contains 2,500 examples, with 2,500/K examples drawn uniformly from each domain.The same probing set is reused across all 40 mixtures for each dataset and architecture pair.
  • Projector: The supervised projector is a 2-layer ReLU MLP with a softmax output trained for 200 epochs at learning rate 1e−4.Five-fold cross-validation uses 32 mixture–reference-model pairs for training and 8 for evaluation in each fold.
  • Metrics: The evaluation averages MAE over 40 ground-truth mixtures and additionally reports MSE.Table 3 uses MSE between predicted ˆπ and ground-truth π⋆ across BERT and GPT-2-Small.

C Future Work

WARP’s next steps focus on making its learned signals more interpretable, improving label-free aggregation, and testing whether its geometric assumptions survive modern large-scale training pipelines.

  • Interpretability of the Learned Projector: The supervised MLP projector implicitly maps Mimic Scores across pseudo-checkpoints to domain proportions, but this mapping remains difficult to interpret.Linear projectors, more synthetic checkpoints, and regularization across t could expose when each domain contributes during the simulated evolution.
  • Advanced Unsupervised Readouts: The unsupervised variant assumes uniform averaging across pseudo-checkpoints and treats resulting per-domain quantities as softmax-compatible logits.Graphical-model inference and weak supervision could learn stage reliabilities and down-weight uninformative steps without labels.
  • Scaling to Large-Scale LLMs: Scaling WARP to large-scale LLMs is an open direction because current experiments isolate recovery using BERT and GPT-2 under known mixtures.Modern LLM adaptation combines continued mid-training, supervised fine-tuning, and preference optimization across many domains.
  • Scaling to Large-Scale LLMs: Extending WARP requires testing whether linear interpolation remains useful when training trajectories cross qualitatively different objectives.The geometric footprint may also become harder to interpret as domain counts and complexity increase.
Loading 2607.01686v1…