PMP-based Data Selection Framework

• PMP-based Data Selection (PDS) is a framework that formulates data selection as an optimal control problem using Pontryagin’s Maximum Principle.
• It integrates gradient-based optimization and proxy evaluations to assign quality scores to data, ensuring efficient pre-training of language models.
• Empirical results show PDS can double training speed and reduce token requirements, thereby enhancing model performance on downstream tasks.

PMP-based Data Selection (PDS) is a principled framework for selecting high-quality data from massive corpora for the pre-training and fine-tuning of LLMs (LMs), grounded in the theory of optimal control and leveraging Pontryagin’s Maximum Principle (PMP). In this approach, data selection is not treated as a heuristic or shallow filtering operation, but rather as a mathematically structured optimization over the entire training dynamics of an LM. The methodology synthesizes techniques from optimal control, machine learning, and gradient-based optimization to maximize downstream model performance while increasing data utilization efficiency.

1. Formulation: Data Selection as an Optimal Control Problem

PDS translates the data selection challenge into an optimal control framework, where the evolution of the LM’s parameters during training is the dynamic system and the data quality scores, $\gamma_n$, assigned to each data point $x_n$, are the control variables. This leads to a general weighted training loss:

$$L(\theta, \gamma) = \sum_{n=1}^{|D|} \gamma_n \cdot l(x_n, \theta)$$

where $l(x_n, \theta)$ is the model loss for instance $x_n$ under parameters $\theta$, and $\gamma$ is constrained to the simplex ($\sum_n \gamma_n = 1$, $\gamma_n \geq 0$ for all $n$).

The model parameters $\theta$ are updated by standard gradient descent:

$$\theta_{t+1} = \theta_t - \eta \nabla L(\theta_t, \gamma)$$

The training trajectory and final model performance, as measured by a downstream objective $J(\cdot)$, are fully determined by the choice of $\gamma$. The PDS problem is then:

$\min_{\gamma \in U} \sum_{t=0}^{T-1} J(\theta_t)$

where $U$ is the set of admissible data weightings.

2. Pontryagin’s Maximum Principle: Necessary Conditions for Optimal Data Selection

Within this formalization, PMP provides a set of necessary optimality conditions for the data score policy. A Hamiltonian is constructed that merges the running loss and the gradient-descent update:

$$H(\theta, \lambda, \gamma) = J(\theta) + \lambda^{\top} f(\theta, \gamma)$$

where $f(\theta, \gamma)$ encodes the parameter update dynamics, and $\lambda$ is a "co-state" or target vector backward-propagating the effects of future losses.

The PMP system leads to three recursions:

1. State Evolution – Model parameters update via gradient descent:

$$\theta^*_{t+1} = -\eta \nabla L(\theta^*_t, \gamma)$$

1. Co-State Recursion – The co-state vector aggregates future downstream effects:

$$\lambda_t^* = \lambda_{t+1}^* + \nabla J(\theta_{t+1}^*) + (\text{terms involving }\nabla^2L)$$

1. Optimal Data Weight Selection – The quality scores are computed by solving:

$$\gamma^* = \arg\max_{\gamma \in U} \sum_{t=0}^{T-1} (t+1)^{\top} \nabla l(x_n, \theta)$$

These conditions extend data selection from pointwise or short-term importance metrics to a global, trajectory-aware, and theoretically justified policy that accounts for the entire training process.

3. PDS Algorithmic Implementation

Because full-scale PMP computation is intractable for real-world LMs and corpora, the PDS framework introduces a practical multi-stage pipeline:

• Proxy Evaluation: Apply the PMP-based optimization on a reduced-size proxy dataset with a lightweight proxy model. Use forward and reverse training passes to solve the PMP conditions and label proxy points with data quality scores.
• Data Scorer Model Training: Train a small data scorer (e.g., a 125M parameter LM) with the proxy data and their quality labels to map new data to predicted scores based solely on input features.
• Full-Corpus Selection: Deploy the data scorer on the full corpus to label all candidates. Select the highest-scoring subset using diversity-promoting algorithms such as Gumbel-Top-K sampling.

This layered approach transforms the theoretical PMP solutions into practical, scalable data selection procedures compatible with modern large-scale LM pipelines (2410.07064).

4. Empirical Benefits and Scaling Law Connections

Empirical investigation reveals that PDS-selected corpora accelerate LM pre-training and consistently enhance downstream task performance—even for models with hundreds of billions of parameters trained on trillions of tokens. Notable findings include:

• Up to 2× speed-up in FLOPs to a fixed loss compared to uniform data usage.
• Substantially improved generalization on benchmarks such as MMLU and OLMo.
• When pre-training data is limited, PDS reduces the token requirement by 1.8× to achieve comparable outcomes.

This efficiency is mathematically tied to improvements in neural scaling laws:

$L(t) = \frac{C}{t^c} + L^{(\text{irreducible})}$

For the area under the reducible loss curve (AUC):

$$\mathrm{AUC}(\gamma) \approx \frac{C(\gamma)}{1-c(\gamma)} T^{1-c(\gamma)}$$

Optimal selection of $\gamma$ via PDS reduces $C$ and/or increases $c$, leading directly to improved sample efficiency and lower loss trajectories.

5. Connections with Related Data Selection Paradigms

PDS builds upon and extends a range of data selection methodologies:

• Bi-Level Optimization (BLO): Prior DPS approaches use outer-loop optimization for instance weights, but face scalability and convergence issues, especially with minibatches.
• Bayesian Approaches: Recent alternatives infer instance weights and network parameters jointly via posterior distributions and SGLD, avoiding the computational intensity of BLO. These have similar aims and can be interpreted as probabilistically grounded analogues of the PMP paradigm (2411.03768).
• Representation and Preference-Based Methods: Techniques such as RDS+ leverage LM hidden states for efficient scaling, while preference-oriented frameworks (e.g., ProDS) integrate human feedback for data ranking (2503.01807, 2505.12754). PDS aligns naturally with these, as PMP-optimal scores can incorporate representation- or preference-derived signals by appropriate gradient formulations.

6. Integration and Deployment in Large-Scale LM Pipelines

The PDS procedure is fully offline and designed for seamless integration into contemporary LM pre-training workflows. Key implementation characteristics include:

Stage	Key Component	Resource Considerations
Proxy PMP Solver	Small proxy LM + dataset	Feasible on a single GPU, moderate memory/compute
Data Scorer Training	Supervised fine-tuning	Standard fine-tuning; efficient at <1B params
Full-Corpus Scoring	Batch inference	Scalable; can be distributed across many workers
Sample Selection	Top-K/Gumbel sampling	Negligible overhead; parallelizable

The process imposes negligible runtime overhead on final LM training and, once trained, the data scorer can be reused for continual data curation as new text corpora become available.

7. Practical Implications and Theoretical Significance

PMP-based Data Selection provides a principled, empirically validated, and operationally tractable solution to the central challenge of data curation in LLM development. By leveraging optimal control theory, PMP-derived gradients, and scalable approximation techniques, PDS maximizes model utility from fixed or constrained data, accelerates convergence, and ensures sustained downstream performance. Its framework accommodates future advances in data scoring—such as integrating preference signals or Bayesian uncertainty—cementing its role as a foundational tool in data-centric, scalable artificial intelligence.