Scaling-Aware Data Selection for End-to-End Autonomous Driving Systems

AI summary

Overview

Research area: Data selection and data mixture optimization for large-scale machine learning, applied to End-to-End (E2E) autonomous driving (AD).

Technical level: Intermediate. The paper assumes familiarity with neural scaling laws, clustering, and data mixture baselines, but the framework itself is described in accessible terms.

Scope: The paper proposes MOSAIC (Mixture Optimization via Scaling-Aware Iterative Collection), a framework that clusters a data pool, fits per-cluster scaling laws against a multi-metric utility, and iteratively mines the cluster with the highest estimated marginal gain, evaluated on E2E autonomous driving with the Hydra-MDP planner.

What This Paper Is About

Training large physical-AI models on all collected data is computationally inefficient, and in autonomous driving the pool can scale to hundreds of millions of hours of clips. Unlike prior data-mixture methods that assume well-defined, homogeneous domains, this paper targets settings where models are judged on multiple competing metrics and where different data samples improve different metrics at different rates. The goal is to select a subset of a data pool, under a fixed budget, that maximizes an aggregate utility (here, Extended Predictive Driving Model Score, or EPDMS).

Key Contributions

1. MOSAIC framework. A generic data mixture optimization pipeline that (i) clusters and ranks data, (ii) models domain-specific data scaling, and (iii) mines samples to maximize expected gain over aggregate metrics.
2. Application to E2E AD. MOSAIC is applied on the NAVSIM and OpenScene benchmarks using the NAVSIM Challenge 2024-winning Hydra-MDP model, achieving higher driving performance than existing data selection and mixture baselines and improving data efficiency by up to 82%. The paper also reports that MOSAIC reaches full training performance with 42% fewer data samples.
3. Empirical validation of clustering plus scaling. MOSAIC outperforms baselines regardless of the clustering approach (semantic captions, geolocation), and embedding scaling laws on top of clustering significantly outperforms clustering-only strategies.
4. Ablation of components. The paper separately disables clustering and ranking to show the contribution of each, finding that ranking largely drives gains in the low-data regime.

Main Findings

• OpenScene results. Across budgets B ∈ {250, 1000, 4000}, MOSAIC achieves the highest EPDMS, roughly one point higher than the next best method. At B = 250: MOSAIC 77.38 ±1.58 (BRMR 0.15) versus Coreset 76.26 ±0.48 (0.20), Chameleon 72.97 ±1.72 (0.86), Random 72.84 ±1.14 (1.00), Uncertainty 70.78 ±0.59 (14.58). At B = 1000: MOSAIC 81.68 ±0.52 (0.18) versus Coreset 80.46 ±0.02 (0.22). At B = 4000: MOSAIC 84.25 ±0.14 (0.18) versus Coreset 83.63 ±0.36 (0.25), Chameleon 82.92 ±0.13 (0.39), Random 80.38 ±0.55 (1.00), Uncertainty 73.46 ±0.19 (2.00).
• OpenScene sample efficiency. MOSAIC requires over 80% fewer samples to match the performance achieved by random selection (BRMR < 0.2).
• Navtrain results. MOSAIC consistently delivers the strongest performance, achieving up to 1.1 points higher EPDMS than the next best method across all budgets. At B = 100: MOSAIC 86.29 ±0.43 (BRMR 0.30) versus Coreset 85.29 ±0.47 (0.53), Chameleon 84.57 ±0.18 (1.07), Random 84.66 ±0.60 (1.00), Uncertainty 84.50 ±0.48 (1.47). At B = 1600: MOSAIC 90.18 ±0.25 (0.37) versus Chameleon 89.50 ±0.20 (0.62), Coreset 89.30 ±0.19 (0.58), Random 88.62 ±0.22 (1.00), Uncertainty 87.75 ±0.37 (1.36).
• Navtrain sample efficiency. MOSAIC attains the lowest BRMR values (< 0.4), corresponding to a 60–70% reduction in the number of samples needed to match Random.
• Metric breakdown (OpenScene). The base model is particularly limited in DAC (83.9) and EC (81.88), which impact EPDMS. MOSAIC achieves the largest gain in DAC, nearly 10 points higher than the base (93.59 ±0.11), while maintaining balanced gains across other metrics and Top-2 performance across all rule-compliance metrics. Other methods improve TTC and EC, which have less effect on the final EPDMS.
• Metric breakdown (Navtrain). MOSAIC achieves consistent improvements across all metrics, with the largest gains in DAC (98.61 ±0.32), NC (98.04 ±0.24), and LK (97.07 ±0.06), shifting collection effort from saturated, less impactful metrics toward those needing more improvement.
• Heterogeneous cluster scaling. With geolocation clustering into four domains (Las Vegas, Boston, Singapore, Pittsburgh) and 2 pilot runs, Boston and Singapore yield the largest initial gains in the low-data regime (< 500 clips), Pittsburgh maintains steadier improvements and eventually supersedes all other domains at high budgets, and Las Vegas provides the smallest gains and saturates early.
• Selection dynamics. Early stages mine only Boston and Singapore; Pittsburgh becomes increasingly favorable between indices 500 to 3700; at 4000 clips the selected set is dominated by Pittsburgh samples. After around 3700 collection rounds the Pittsburgh domain is exhausted, and beyond approximately 2500 sample selections the other clusters approach saturation, leading MOSAIC to mine from Las Vegas.
• Robustness to clustering mechanism. With Navtrain clusters formed from captions generated by Qwen-2.5-VL-32B-Instruct and TF-IDF vectors into six clusters, MOSAIC consistently beats all baselines, requires 61% and 52% fewer clips than Random to match its performance at the 1,600- and 2,400-clip budgets, and reaches full-training performance with only 2,400 clips (42% fewer samples). Chameleon degrades under caption-based clustering despite being the strongest baseline under geolocation clustering.
• Ablation. Disabling clustering ("w/o Clustering") or ranking ("w/o Ranking") shows that improvements in the low-data regime (up to a budget of 800 clips) can largely be attributed to ranking; in the higher-data regime, adding clips with low EPDMS scores becomes less effective and "w/o Clustering" lags behind MOSAIC. Both variants still outperform random collection by a large margin.
• Reported abstract-vs-body figures. The abstract states MOSAIC outperforms baselines on EPDMS with "up to 80% less data," while the introduction states "up to 82% less additional data."

Methodology in Plain English

The framework has three stages. First, the data pool is partitioned into discrete domains (clusters) using a feature representation, and samples within each cluster are ranked by an importance score—in this application, the utility the current model achieves on that sample. Second, the framework fits neural scaling laws from each domain to the evaluation metrics. Specifically, it assumes the total change in utility is approximately the sum of each domain's independent effect, and models each domain's effect with a saturating exponential law, a_i(1 − e^(−n/τ_i)), where a_i is the asymptotic improvement from that domain and τ_i controls how quickly marginal benefit declines. These parameters are estimated from small-scale pilot runs. Third, the framework allocates data samples one at a time: at each step it computes the marginal improvement of adding one more sample from each cluster, identifies the cluster with the highest marginal improvement, draws the next ranked sample from that cluster, and repeats until the budget is reached. Because each scaling curve is concave, the marginal gain from a cluster decreases as more data is drawn from it, so the policy naturally shifts to other clusters. The method is evaluated with the Hydra-MDP model, a pretrained VoVNetV2-99 backbone, and a trajectory vocabulary of 16,384, on Navtrain and OpenScene, with EPDMS as the metric and Budget Ratio to Match Random (BRMR) as a sample-efficiency measure.

Why This Matters

Impact on research. The work connects two lines of work—data mixture optimization and neural scaling laws—and argues that prior mixture methods overlook heterogeneous, metric-dependent improvement rates. It offers a principled alternative to relying on well-defined homogeneous domains, and its ablation isolates the roles of clustering and ranking.

Real-world applications:

• Autonomous driving fleet data curation, where clips vary by location, weather, and traffic density and annotation/training budgets are limited.
• Robotics and other physical AI systems with multiple competing safety and performance criteria.
• Industrial perception or sensor-fusion pipelines where collected data spans many deployment conditions.
• Guidance for data collection campaigns, indicating which geographic or semantic domains are worth collecting more of.

Industry relevance. The paper reports that MOSAIC reaches full training performance with 42% fewer samples, which is directly relevant to the compute and labeling cost of training large models. The authors' affiliations (NVIDIA, with an internship from New York University, and University of Ottawa) reflect industrial interest in scalable training-set design.

Future Directions

• Mitigate the separable-scaling assumption. The relaxation assuming each cluster's contribution is captured by its own scaling curve ignores cross-cluster interactions; the authors note that violated separation could lead to suboptimal allocation.
• Reduce pilot-run overhead. MOSAIC relies on pilot runs to estimate cluster-specific scaling curves; the paper states (Section 9 of the Appendix) that accurate scaling fits can be obtained efficiently using small pilot subsets or continual training, but removing the initial overhead entirely remains an open question.
• Evaluate under different or closed-loop protocols. The paper notes that conventional open-loop metrics correlate poorly with closed-loop driving quality and that EPDMS has become a standard; extending the framework to additional benchmarks or domains is a natural next step.
• Explain the sensitivity of competing mixture methods. Chameleon degrades under caption-based clustering despite being the strongest baseline under geolocation clustering, which the authors attribute to sensitivity of its kernel ridge weighting to cluster structure—an area requiring further study.

Target Audience

Researchers and practitioners working on data selection, data mixture optimization, and neural scaling laws, particularly those applying these methods to physical AI and autonomous driving. It is also relevant to machine learning engineers in industry who must design training sets under fixed data or compute budgets, and to readers interested in how multi-metric evaluation changes data curation decisions.

Authors’ abstract

Large-scale deep learning models for physical AI applications depend on diverse training data collection efforts. These models and correspondingly, the training data, must address different evaluation criteria necessary for the models to be deployable in real-world environments. Data selection policies can guide the development of the training set, but current frameworks do not account for the ambiguity in how data points affect different metrics. In this work, we propose Mixture Optimization via Scaling-Aware Iterative Collection (MOSAIC), a general data selection framework that operates by: (i) partitioning the dataset into domains; (ii) fitting neural scaling laws from each data domain to the evaluation metrics; and (iii) optimizing a data mixture by iteratively adding data from domains that maximize the change in metrics. We apply MOSAIC to autonomous driving (AD), where an End-to-End (E2E) planner model is evaluated on the Extended Predictive Driver Model Score (EPDMS), an aggregate of driving rule compliance metrics. Here, MOSAIC outperforms a diverse set of baselines on EPDMS with up to 80\% less data.