🤖 machine learning

RCAP: Robust, Class-Aware, Probabilistic Dynamic Dataset Pruning

The paper introduces RCAP, a robust, class-aware, probabilistic dynamic dataset pruning algorithm that adaptively selects high-loss samples per class to significantly improve worst-group accuracy and training efficiency, even outperforming full-data training on imbalanced datasets with only 10% of the data.

Original authors: Atif Hassan, Swanand Khare, Jiaul H. Paik

Published 2026-06-11

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Atif Hassan, Swanand Khare, Jiaul H. Paik

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 trying to teach a student for a massive final exam. You have a library containing millions of books (the full dataset). Traditionally, to get the best results, you make the student read every single book. This takes forever, costs a fortune in electricity and time, and often leads to burnout.

Data pruning is the idea of saying, "Let's just pick the most important books and throw the rest away." The goal is to teach the student just as well, but much faster and cheaper.

However, there's a catch. If you just pick the "hardest" books, you might accidentally ignore the rare, difficult topics that the student struggles with the most. If the student fails on those rare topics, they fail the whole exam, even if they aced the easy stuff. This is what researchers call a lack of robustness.

Enter RCAP (Robust, Class-Aware, Probabilistic Dynamic Dataset Pruning). Think of RCAP as a super-smart, adaptive tutor who doesn't just pick a static list of books once and stick to it. Instead, RCAP changes the reading list every single day (every "epoch" of training) based on how the student is doing.

Here is how RCAP works, broken down into simple steps:

1. The "Class-Aware" Strategy (The Balanced Diet)

Imagine your student is studying for a test with 100 different subjects. Some subjects are easy (like "Apples"), and some are very hard (like "Quantum Physics").

The Problem: If you just grab the "hardest" questions from the whole pile, you might accidentally grab 100 questions about "Apples" (because there are so many of them) and zero questions about "Quantum Physics." The student gets great at Apples but fails the exam because they never saw a Quantum Physics question.
RCAP's Solution: RCAP looks at the "score" (loss) for each subject separately. If the student is struggling with "Quantum Physics," RCAP says, "Okay, for today's study session, we need to make sure we pick a specific percentage of questions just for Quantum Physics." It calculates exactly how many questions to keep for every single subject to ensure no topic is left behind.

2. The "Dynamic" Strategy (The Daily Schedule)

Most other methods pick their study list once at the beginning and never change it.

RCAP's Solution: RCAP is dynamic. Every day, it checks the student's homework.
- If the student did poorly on "Quantum Physics" yesterday, RCAP increases the number of Quantum Physics questions for today.
- If the student is breezing through "Apples," RCAP might reduce the number of Apple questions.
- It recalculates this balance every single day, ensuring the study plan evolves as the student learns.

3. The "Probabilistic" Strategy (The Lottery Ticket)

Once RCAP decides, "We need 50 questions about Quantum Physics," how does it pick which 50?

The Old Way: Sort all 1,000 Quantum Physics questions by difficulty and pick the top 50. This is slow and computationally expensive.
RCAP's Way: It uses a "weighted lottery." It gives every question a ticket. The harder the question is for the student (the higher the "loss"), the bigger the ticket they get.
- A very easy question has a tiny ticket (low chance of being picked).
- A very hard question has a giant ticket (high chance of being picked).
- RCAP then draws 50 tickets. This is much faster than sorting, but it still guarantees that the hardest questions are picked most often.

4. The "Robust" Result (The Safety Net)

The paper claims that by doing this, RCAP solves a major problem: Worst-Group Accuracy.

In the real world, "groups" can be different types of data (e.g., photos of cats vs. dogs, or images of people with different hair colors).
Other methods might get a high average score but fail miserably on the hardest group.
RCAP ensures that even the "worst-performing" group gets enough attention. The paper shows that on difficult, unbalanced datasets, RCAP actually performs better than training on the full dataset, while using only 10% of the data.

The Bottom Line

RCAP is like a smart, adaptive coach who:

Monitors every specific topic the student struggles with.
Adjusts the daily study plan to focus more on the weak spots.
Picks the hardest practice problems using a fast, fair lottery system.

The Results:

Speed: It makes training 8.69 times faster on average.
Quality: It doesn't just save time; it often improves the model's ability to handle difficult, rare cases (robustness).
Efficiency: It achieves this without needing extra computing power; it just uses the "score" the model already calculates during normal training.

In short, RCAP proves that you don't need to read the whole library to get an A. You just need the right tutor to pick the right books at the right time.

Technical Summary: RCAP – Robust, Class-Aware, Probabilistic Dynamic Dataset Pruning

1. Problem Statement

Deep learning models require massive datasets, leading to exorbitant computational costs and unsustainable training times. While data pruning techniques aim to reduce these costs by removing less informative examples, existing methods face significant limitations:

Static vs. Dynamic: Static pruning selects a fixed subset before training, often relying on expensive pre-training to determine sample importance. Dynamic pruning adapts subsets during training but often fails to maintain robustness.
Worst-Group Accuracy: Current state-of-the-art dynamic pruning methods (e.g., RS2, InfoBatch) prioritize average performance but frequently neglect worst-group accuracy. This is particularly detrimental in class-imbalanced datasets or when specific classes (e.g., "cats" vs. "dogs" in CIFAR10) exhibit higher loss, leading to non-robust models that fail on minority or difficult groups.
Class-Awareness: Existing methods often lack a principled mechanism to dynamically adjust the representation of individual classes based on their difficulty (hardness) during the training process.

The core question addressed is: Does incorporating class hardness during data pruning enhance model robustness across both balanced and imbalanced settings?

2. Methodology: RCAP

The authors propose RCAP (Robust, Class-Aware, Probabilistic), a dynamic dataset pruning algorithm designed to optimize worst-group accuracy while maintaining efficiency. The method operates in two primary phases per epoch:

2.1 Adaptive Per-Class Subset Size

RCAP addresses the problem of determining the optimal number of samples to retain for each class.

Theoretical Foundation: The authors derive a closed-form solution (Theorem 3.1) for minimizing the total empirical error under a fixed pruning rate constraint. The solution demonstrates that classes with higher cumulative loss (indicating greater difficulty) should be allocated a proportionally larger share of the training subset.
Implementation: In every epoch $t$ $t$ , the fraction of samples $\alpha_j$ $α_{j}$ to retain for class $j$ $j$ is calculated based on the aggregated loss of that class from the previous epoch ( $t-1$ $t - 1$ ).
- The formula is: $\hat{\alpha}_j^{t+1} \propto \frac{\sqrt{p_j \tilde{E}_j^t}}{n_j}$ , where $p_j$ is the class proportion, $\tilde{E}_j^t$ is the aggregated loss, and $n_j$ is the total samples in class $j$ .
- Clipping Mechanism: To ensure $\alpha_j \leq 1$ , if the calculated fraction exceeds 1, the excess is redistributed among other classes. This prevents the algorithm from attempting to select more samples than exist for a specific class.

2.2 Adaptive Per-Class Sample Selection

Once the subset size for each class is determined, RCAP selects specific samples.

Probabilistic Sampling: Instead of deterministic sorting (which is computationally expensive), RCAP samples from a probability distribution defined over per-sample losses.
Loss-Based Probability: The probability of selecting a sample $X_i$ for class $j$ is proportional to $e^{\phi_j(X_i)/\beta}$ , where $\phi_j$ represents the sample's loss and $\beta$ is a temperature hyper-parameter.
Loss Clipping: To prevent instability caused by outliers (e.g., label noise causing extremely high losses), the loss values are clipped against the maximum loss observed at initialization (epoch 0). This ensures the sampling distribution remains robust and does not overly bias selection toward noisy outliers.
Iterative Update: The loss values $\phi$ are updated at the end of each epoch based on the current model state, and the process repeats.

Complexity: The method incurs no additional computational overhead beyond the standard forward pass, as all metrics are derived from losses computed during training. The per-sample time complexity is $O(1)$ .

3. Key Contributions

Novel Algorithm: Introduction of RCAP, the first dynamic pruning method to explicitly combine a closed-form solution for class-wise subset sizing with an adaptive probabilistic sampling strategy.
Robustness Focus: The method specifically targets the improvement of worst-group accuracy, addressing a critical gap in existing pruning literature where minority or difficult classes are often underrepresented.
Comprehensive Evaluation: The paper provides the first comprehensive evaluation of dynamic pruning across both class-balanced and highly imbalanced datasets, spanning six datasets (CIFAR10, CIFAR100, ImageNet, Waterbirds, CelebA, iNaturalist) and five distinct architectures (ResNet, EfficientNet, Dinov2, etc.).
Efficiency: The approach achieves significant speedups without sacrificing performance, particularly on imbalanced data.

4. Experimental Results

The authors evaluated RCAP against seven state-of-the-art baselines (including CCS, MetriQ, TDDS, UCB, InfoBatch, and RS2) across three training paradigms: training from scratch, transfer learning, and fine-tuning.

Worst-Group Accuracy: RCAP consistently outperformed all baselines across all datasets and pruning rates. Notably, on highly imbalanced datasets (Waterbirds, CelebA, iNaturalist), RCAP achieved superior worst-group accuracy even at 90% pruning rates, often surpassing the performance of models trained on the full dataset.
- Example: On iNaturalist (90% pruning), RCAP achieved ~68.5% worst-group accuracy, while other methods dropped significantly (e.g., InfoBatch at ~5.5%, TDDS at 0%).
Average Accuracy: On class-balanced datasets (CIFAR10, CIFAR100), RCAP maintained competitive average accuracy comparable to or better than existing methods.
Efficiency: RCAP delivered an average 8.69× speedup compared to full-data training. On imbalanced datasets, it achieved >1% performance improvement over full-data training while reducing training time drastically.
Hyper-parameter Sensitivity: An ablation study on the temperature parameter $\beta$ revealed that $\beta < 1$ is crucial for high pruning rates (90%) to prioritize high-loss samples, while $\beta > 1$ offers more stability at moderate rates (50%).

5. Significance and Claims

The paper claims that RCAP establishes a new state-of-the-art in robust dynamic dataset pruning. Its significance lies in:

Solving the Imbalance Problem: By dynamically adjusting class-wise representation based on loss, RCAP naturally mitigates the bias against minority groups that plagues other pruning methods, especially at high pruning rates.
Practical Efficiency: It offers a "best-of-both-worlds" scenario: drastically reduced training time (speedup) coupled with improved model robustness (worst-group accuracy), challenging the notion that pruning must trade off performance for speed.
Generalizability: The method is model-agnostic and effective across diverse architectures and training paradigms.

Limitations Acknowledged by the Authors:

RCAP requires a few initial epochs to accurately approximate optimal class-wise fractions, which may limit its utility in scenarios requiring immediate effectiveness (e.g., few-shot learning with Large Language Models).
The temperature hyper-parameter $\beta$ currently requires manual tuning. The authors suggest future work to make $\beta$ adaptive during training.