Technical Summary: WARP – Weight-Space Analysis for Recovering Training Data Portfolios

Problem Statement

Foundation models are increasingly released to the public, yet the specific "data recipes"—particularly the domain mixture weights determining how diverse sources (e.g., web text, code, scientific literature) are sampled during training—remain proprietary and opaque. This creates a significant access asymmetry: researchers can study the resulting model weights but lack visibility into the training distribution that produced them.

Existing approaches to inferring training data, such as Membership Inference (MI), operate at the level of individual samples to detect whether a specific data point was in the training set. However, MI cannot characterize the global composition or domain proportions of the training corpus. Furthermore, prior methods for recovering training trajectories often require access to intermediate checkpoints, which are rarely available for released models (where only the base model and the final fine-tuned weights are typically public). The core challenge addressed by this work is the inverse problem: given only the released base model ($\theta_{base}$), the fine-tuned model ($\theta_{ref}$), and a small probing dataset, how can one recover the unknown domain mixture proportions ($\pi^*$) used to train the model, without access to the true training trajectory or intermediate checkpoints?

Methodology: The WARP Framework

The authors propose WARP (Weight-space Analysis for Recovering Training Data Portfolios), a framework that recovers domain mixtures directly from model weights by exploiting the geometric structure of the weight space. The methodology proceeds in three stages:

1. Simulating the Training Trajectory

Since the true path from $\theta_{base}$ to $\theta_{ref}$ is unobserved, WARP uses model merging techniques to interpolate between the two endpoints. By applying merging operators (such as Linear Interpolation/LERP, Spherical Interpolation/SLERP, or TIES) with a schedule of mixing coefficients $\alpha_t \in (0,1)$, the framework generates a sequence of pseudo-checkpoints ($\tilde{\theta}_t$). These pseudo-checkpoints serve as proxies for the unobserved parameter evolution, effectively reconstructing a "simulated" training trajectory.

2. Distilling Geometric Footprints

From these simulated checkpoints, WARP extracts a domain-level signal using a metric adapted from the Mimic Score. For each pseudo-checkpoint $\tilde{\theta}_t$, the method defines a direction vector $v_t = \theta_{ref} - \tilde{\theta}_t$ pointing toward the target model.

• Sample Alignment: For a small probing dataset $D_{probe}$ (sampled from the same source domains as the training data), the method computes the negative gradient of each sample at $\tilde{\theta}_t$.
• Projection: The Mimic Score ($m_{i,t}$) is calculated as the projection of this negative gradient onto the unit vector of $v_t$. This score measures how strongly a sample aligns with the direction the model actually took during fine-tuning.
• Pooling: These scores are normalized and averaged within each domain $k$ at each step $t$, resulting in a feature matrix $M \in \mathbb{R}^{K \times T}$. This matrix represents the "geometric footprint" of the training data, tracing how strongly each domain aligns with the fine-tuning direction as the model evolves.

3. Mapping Geometry to Mixture

The final step translates the geometric footprint $M$ into a predicted domain mixture $\hat{\pi}$. The authors propose two variants:

• Unsupervised Readout: A parameter-free approach that averages the pooled scores across time steps and applies a softmax function to estimate the mixture.
• Supervised Projection: An MLP projector trained on synthetic data. The authors generate synthetic pairs by fine-tuning base models on random mixtures to create ground-truth labels, then training the MLP to map the geometric footprint $M$ to the known mixture $\pi$. This allows the model to learn how to weight different stages of the simulated trajectory adaptively.

Key Contributions

1. Closing the Access Asymmetry: WARP formalizes the recovery of a fine-tuned model's domain mixture using only artifacts available for most released models (base and final weights), bypassing the need for proprietary data or intermediate checkpoints.
2. Simulated Training Trajectories: The framework introduces a novel use of model merging not for capability fusion, but for simulating unobserved training paths to expose usable geometric structure in weight space.
3. Strong Empirical Performance: In controlled experiments, WARP outperforms sample-level membership inference baselines and even an "oracle" variant that has access to the true intermediate training checkpoints.
4. Robustness: The method remains accurate across different training regimes, including early-stopped, converged, and overtrained checkpoints.

Experimental Results

The framework was evaluated on BERT and GPT-2 models fine-tuned on four text datasets (SNLI, AGNews, Yelp, Yahoo) with known ground-truth mixtures.

• Accuracy: WARP achieved a Mean Absolute Error (MAE) as low as 0.046 on BERT and 0.104 on GPT-2 (using the supervised MLP variant with TIES merging).
• Comparison: These results significantly outperformed:
  • Random Guess and Centroid Guess baselines.
  • Sample-level Membership Inference (MI), which yielded higher errors (e.g., 0.063 MAE on BERT vs. 0.046 for WARP).
  • Oracle Variants: Surprisingly, WARP using simulated trajectories (LERP/TIES) often outperformed the variant using real intermediate checkpoints. The authors attribute this to the "smoothness gap": real training paths contain noise (stochastic mini-batches, learning rate warmup), whereas the interpolated paths provide a cleaner, monotone trajectory that makes the geometric footprint easier to decode.
• Unsupervised vs. Supervised: The parameter-free unsupervised softmax readout was competitive, outperforming all baselines and closing most of the gap to the supervised variant, demonstrating that the geometric signal is strong even without learned projectors.

Significance and Claims

The paper claims that WARP demonstrates the feasibility of reverse-engineering a model's training distribution solely from its weight-space geometry. By showing that the "path" between a base and fine-tuned model encodes a distributional footprint of the training data, the work offers a new tool for:

• Auditing: Detecting data contamination or verifying the provenance of model training.
• Interpretability: Understanding why models exhibit specific behaviors based on their training mixtures.
• Transparency: Addressing the opacity of frontier model development where weights are shared but data recipes are not.

The authors emphasize that this approach breaks the traditional forward-only paradigm of domain mixture optimization, providing a practical method to infer the "ingredients" of a model when only the "final dish" (weights) is available.