💻 computer science

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

This paper proposes MOSAIC, a scaling-aware data selection framework that optimizes training data mixtures by partitioning datasets into domains and fitting neural scaling laws to evaluation metrics, demonstrating that it achieves superior performance in end-to-end autonomous driving with up to 80% less data compared to existing baselines.

Original authors: Tolga Dimlioglu, Nadine Chang, Maying Shen, Rafid Mahmood, Jose M. Alvarez

Published 2026-04-10

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Tolga Dimlioglu, Nadine Chang, Maying Shen, Rafid Mahmood, Jose M. Alvarez

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are training a robot to drive a car. You have a massive library of video clips showing every kind of driving situation imaginable: sunny highways, rainy city streets, crowded intersections, and quiet country roads.

The problem? You can't teach the robot with all of them. It would take too long, cost too much money, and the robot might get confused by too much information. You need to pick the best clips to teach it.

But here's the tricky part: different clips teach different lessons. A clip of a busy city intersection is great for teaching the robot how to avoid pedestrians, but it's terrible for teaching it how to merge onto a highway. If you just pick clips randomly, you might end up with a robot that is amazing at city driving but crashes on the highway.

This paper introduces a smart new method called MOSAIC (Mixture Optimization via Scaling-Aware Iterative Collection) to solve this problem. Think of MOSAIC as a super-smart curriculum designer for your robot driver.

Here is how it works, broken down into three simple steps:

1. Sorting the Library (Clustering)

Imagine your video library is a giant, messy attic. MOSAIC starts by organizing this attic into neat, labeled boxes.

Box A: "Curvy Mountain Roads"
Box B: "Dense City Traffic"
Box C: "Rainy Highways"

Instead of looking at every single video individually, MOSAIC groups them by where they were filmed or what is happening in them. This helps the system understand that "City Traffic" videos are a different type of lesson than "Mountain Road" videos.

2. The "Growth Chart" (Scaling Laws)

This is the magic sauce. In the past, people just guessed which box had the most useful videos. MOSAIC is smarter; it looks at a growth chart.

Imagine you are feeding a plant. You know that if you add a little water, it grows a bit. If you add a lot, it grows a lot. But eventually, adding more water doesn't help; the plant is already full, and you might even drown it.

MOSAIC tests a few small samples from each "Box" (City, Mountain, Rain) to see how much the robot's driving skills improve.

It might find that City videos give a huge boost at first, but after 100 clips, the robot learns everything it needs from them, and adding more doesn't help.
It might find that Mountain videos start slow, but keep improving the robot's skills even after 1,000 clips.

This creates a "scaling law" for each box: a curve showing exactly how much benefit you get from adding more data from that specific group.

3. The Smart Shopping List (Iterative Collection)

Now, MOSAIC acts like a chef with a limited budget for ingredients. The goal is to make the best possible meal (the best driving model) without wasting money.

Instead of buying 500 tomatoes and 500 potatoes (random selection), MOSAIC looks at the growth charts:

"Right now, adding one more City clip will give us a huge jump in skill." -> Add a City clip.
"Okay, the robot is getting good at City driving. Adding another City clip won't help much anymore. But adding a Mountain clip right now will give a big jump." -> Switch to a Mountain clip.

It does this one clip at a time, constantly checking: "Which type of video will give us the biggest improvement right this second?" It keeps switching between the boxes until the budget is full.

The Result: A Smarter Driver with Less Data

The paper tested this on real autonomous driving data. The results were impressive:

Efficiency: MOSAIC could build a driver that was just as good as the one trained on all the data, but it only needed 42% of the data.
Speed: In some cases, it achieved the same performance with 80% less data than random selection.
Balance: It didn't just make the robot good at one thing; it balanced the robot's skills so it was safe in the city and on the highway.

The Analogy in a Nutshell

Old Way: Throwing a bunch of random flashcards at a student and hoping they learn everything.
MOSAIC: A tutor who looks at the student's progress report, realizes they are great at math but bad at history, and says, "Okay, let's stop doing math problems for a moment and focus entirely on history until you catch up, then we'll switch back."

Why does this matter?
Training self-driving cars is expensive and slow. MOSAIC proves that you don't need more data to get better results; you just need smarter data selection. It's the difference between eating a buffet where you grab everything randomly, and a chef who carefully selects the perfect ingredients to create a masterpiece meal.

1. Problem Statement

Large-scale physical AI systems, particularly End-to-End (E2E) Autonomous Driving (AD), rely on massive, diverse datasets. However, training on all collected data is computationally prohibitive. The core challenge lies in data mixture selection:

Competing Metrics: AD models must satisfy multiple, often conflicting, evaluation metrics (e.g., lane keeping vs. comfort vs. collision avoidance).
Heterogeneous Influence: Different data samples (e.g., dense city traffic vs. curvy suburban roads) influence these metrics at different rates.
Ambiguity in Domains: Existing data selection methods often assume homogeneous domains or fail to account for how specific data subsets scale performance across different metrics.
Goal: Select a subset of data from a large pool that maximizes a global utility function (aggregating multiple metrics) under a strict data budget, without needing to train on the full dataset.

2. Methodology: MOSAIC

The authors propose MOSAIC (Mixture Optimization via Scaling-Aware Iterative Collection), a three-stage framework designed to optimize data mixtures by modeling how performance scales with data volume from specific domains.

Stage 1: Clustering and Ranking

Clustering: The data pool ( $D_{pool}$ ) is partitioned into $M$ discrete domains (clusters) based on feature representations (e.g., geolocation or semantic captions). The goal is to group samples that have similar influences on the evaluation metrics.
Ranking: Within each cluster, samples are ranked by an importance score $I(x)$ . This score is defined as the utility gain of adding that specific sample to the current training set. The framework prioritizes high-influence samples within each domain first.

Stage 2: Estimating Scaling Laws

Pilot Runs: The system performs small-scale "pilot" training runs to estimate how the global utility function $U$ improves as data is added from each specific cluster.
Scaling Model: The improvement in utility for cluster $i$ $i$ is modeled using a saturating exponential scaling law:
$\Delta U_i(n) \approx a_i (1 - e^{-n/\tau_i})$
Where:
- $n$ is the number of samples added from cluster $i$ .
- $a_i$ represents the asymptotic maximum improvement (potential gain) of that domain.
- $\tau_i$ governs the saturation rate (how quickly the marginal benefit diminishes).
Approximation: The total utility gain is approximated as the sum of individual cluster gains (linear separability assumption), allowing the complex mixture problem to be solved efficiently.

Stage 3: Iterative Data Collection

Greedy Optimization: The algorithm iteratively selects data points one by one. At each step, it calculates the marginal gain ( $\delta_i$ ) of adding one more sample from each cluster based on the fitted scaling laws.
Selection: It selects the sample from the cluster $j$ that offers the highest marginal gain:
$j = \arg\max_i \delta_i(b_i)$
Where $b_i$ is the current count of samples selected from cluster $i$ .
Dynamic Allocation: As a cluster's scaling curve saturates (diminishing returns), the algorithm naturally shifts focus to other clusters with higher potential gains, optimizing the mixture dynamically.

3. Key Contributions

MOSAIC Framework: A general, principled pipeline for data mixture optimization that jointly leverages clustering, sample ranking, and scaling-law modeling.
Handling Competing Metrics: Unlike methods that optimize for a single metric, MOSAIC optimizes an aggregate utility function (EPDMS) that balances multiple rule-compliance metrics, effectively managing trade-offs.
Robustness to Clustering: The method is shown to be robust across different clustering strategies (geolocation vs. semantic captions), provided the clusters are semantically coherent.
Efficiency: It significantly reduces the data required to reach target performance levels compared to baselines.

4. Experimental Results

The framework was evaluated on End-to-End Autonomous Driving using the Hydra-MDP model (2024 NAVSIM Challenge winner) on two benchmarks: OpenScene and Navtrain. The primary metric was EPDMS (Extended Predictive Driver Model Score), an aggregate of 9 rule-compliance metrics.

Performance Gains:
- MOSAIC consistently outperformed baselines (Random, Uncertainty, Coreset, Chameleon) across all data budgets.
- On OpenScene, MOSAIC achieved the highest EPDMS, roughly 1 point higher than the next best method.
- On Navtrain, MOSAIC achieved up to 1.1 points higher EPDMS than the best baseline.
Data Efficiency:
- MOSAIC required up to 82% less data to match the performance of a Random selection baseline.
- Specifically, it achieved full training performance (using all available data) with only 42% of the data samples.
- The Budget Ratio to Match Random (BRMR) was consistently low (e.g., 0.15–0.38), indicating high sample efficiency.
Metric Breakdown: MOSAIC showed balanced improvements across all 9 sub-metrics, particularly excelling in areas where the base model was weak (e.g., DAC and EC), whereas other methods often improved only specific, less critical metrics.
Ablation Studies:
- Removing the ranking step hurt performance in low-data regimes.
- Removing clustering hurt performance in high-data regimes, confirming that both components are necessary.
- The method remained effective even when using semantic captions for clustering instead of geolocation.

5. Significance and Impact

Scalability for Physical AI: MOSAIC addresses the critical bottleneck of data collection costs in physical AI. By identifying the "most influential" data subsets, it allows developers to achieve state-of-the-art performance with significantly fewer resources.
Beyond Single-Metric Optimization: The paper demonstrates that optimizing for a single metric (like accuracy or loss) is insufficient for complex systems like AD. The scaling-aware approach effectively navigates the trade-offs between safety, comfort, and rule compliance.
Generalizability: While applied to AD, the framework is generic and can be applied to any large-scale deep learning system where data comes from heterogeneous sources and performance is measured by multiple competing objectives.
Compute Efficiency: Although MOSAIC requires an initial overhead for pilot runs to fit scaling laws, the paper shows that this investment pays off in the large-data regime, ultimately saving significant GPU hours compared to training on random or suboptimal data mixtures.

In conclusion, MOSAIC provides a principled, data-efficient strategy for training robust autonomous driving systems by intelligently curating data mixtures based on how different data domains scale model performance.