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

Abstract: Dynamic data pruning techniques aim to reduce computational cost while minimizing information loss by periodically selecting representative subsets of input data during model training. However, existing methods often struggle to maintain strong worst-group accuracy, particularly at high pruning rates, across balanced and imbalanced datasets. To address this challenge, we propose RCAP, a Robust, Class-Aware, Probabilistic dynamic dataset pruning algorithm for classification tasks. RCAP applies a closed-form solution to estimate the fraction of samples to be included in the training subset for each individual class. This fraction is adaptively adjusted in every epoch using class-wise aggregated loss. Thereafter, it employs an adaptive sampling strategy that prioritizes samples having high loss for populating the class-wise subsets. We evaluate RCAP on six diverse datasets ranging from class-balanced to highly imbalanced using five distinct models across three training paradigms: training from scratch, transfer learning, and fine-tuning. Our approach consistently outperforms state-of-the-art dataset pruning methods, achieving superior worst-group accuracy at all pruning rates. Remarkably, with only $10\%$ data, RCAP delivers $>1\%$ improvement in performance on class-imbalanced datasets compared to full data training while providing an average $8.69\times$ speedup. The code can be accessed at https://github.com/atif-hassan/RCAP-dynamic-dataset-pruning

• The paper proposes a dynamic data pruning method that computes per-class selection fractions using a closed-form, loss-based formula.
• It employs probabilistic, high-loss sample prioritization with dynamic clipping to mitigate outlier effects and ensure robust subgroup performance.
• Empirical results demonstrate up to 8.69x training speedup and significant accuracy gains on worst-group metrics, notably in imbalanced datasets.

RCAP: A Robust, Class-Aware, Probabilistic Dynamic Pruning Algorithm

Overview

RCAP addresses the critical challenge in contemporary deep learning of reducing training costs via dataset pruning without compromising model robustness, especially under severe class imbalance. Existing dynamic pruning methods typically focus on overall performance metrics and often neglect class or group-level robustness, leading to suboptimal worst-group accuracy at high pruning rates. RCAP introduces a dynamic, closed-form, class-aware approach to subsetting, combined with probabilistic, loss-driven sample selection. This design is shown to yield significant improvements in both worst-group and average-group accuracies while simultaneously achieving substantial training speedups.

Methodological Contributions

RCAP frames dynamic data pruning as a per-epoch process with two tightly coupled optimization modules:

1. Adaptive Per-Class Subset Size Determination: RCAP computes, at every epoch, the desired fraction of samples to select from each class using a closed-form expression derived from class-wise aggregated losses. Classes with elevated cumulative training error are allocated higher representation, theoretically minimizing the overall empirical error. This formulation leverages recent advances in loss-proportional allocation, establishing adaptivity to evolving model uncertainties and intrinsic class-level hardness.
2. Probabilistic, High-Loss Sample Prioritization: Within each class-selected subset, RCAP deploys a Softmax-transformed probability distribution over per-sample losses, with sampling temperatures controlling selection sharpness. A crucial innovation is dynamic clipping of per-sample losses, capping them to the maximum initial epoch loss, which mitigates the effect of persistent noisy or corrupted samples and enhances robustness against outliers.

Crucially, RCAP is computationally efficient: the estimation of selection fractions and construction of the per-class sampling distribution exploit loss values computed during the forward pass, yielding $O(1)$ time per sample.

Empirical Evaluation

RCAP is benchmarked against seven state-of-the-art static and dynamic pruning baselines, including CCS, MetriQ, TDDS, UCB, InfoBatch, and RS2 (with and without replacement). The empirical analysis covers a spectrum of dataset types and scales—CIFAR10, CIFAR100, ImageNet, Waterbirds, CelebA, and iNaturalist—capturing both class-balanced and highly imbalanced scenarios. Multiple model architectures (ResNet18/50, EfficientNetV2, Dinov2, EfficientFormer V2) and training paradigms (from scratch, transfer, fine-tune) were used for comprehensive coverage.

Key Numerical Findings

• Worst-Group Accuracy: RCAP outperforms all baselines on worst-group accuracy across all datasets and pruning rates, with especially large margins at very high pruning ($\geq 90\%$). For class-imbalanced datasets (CelebA, iNaturalist), RCAP consistently achieves worst-group accuracies exceeding full data training, owing to its loss-adaptive sample reallocation which counters representational imbalance.
• Efficiency: RCAP demonstrates an average 8.69x speedup over full-data training. On ImageNet, Waterbirds, CelebA, and iNaturalist, RCAP achieves training performance with <$1\%$ accuracy drop—even with only $10\%$ of the data—while improving worst-group metrics. This indicates efficient data utilization and effective mitigation of training redundancy.
• Average-Group Accuracy: On class-balanced datasets, RCAP matches or slightly outperforms top baselines in average-group top-1 accuracy, indicating minimal trade-off between robustness and global performance.

Sensitivity and Ablation

The algorithm's only sensitive hyperparameter is the Softmax temperature controlling the sampling distribution. Ablation studies indicate that optimal temperature settings depend on both dataset characteristics and pruning rate. At high pruning, sharper distributions (lower temperature) yield improved worst-group performance but require careful tuning; at moderate pruning, performance is more robust to the temperature value.

Theoretical and Practical Implications

RCAP offers a principled solution to dynamic pruning by unifying adaptive sample allocation and class-aware loss-driven selection, which is both theoretically sound and empirically effective. Theoretical analysis justifies the closed-form per-class fraction, while practical implementations confirm robust operation across a wide variety of data regimes.

Practically, RCAP's approach facilitates democratization of deep learning: resource-constrained environments can train models competitively at a fraction of data and computation costs, without sacrificing subgroup fairness or robustness. Additionally, the method's architecture-agnostic nature broadens its applicability to diverse real-world pipelines and domains.

On a theoretical level, RCAP's adaptive class allocation paradigm addresses longstanding open questions on balancing model learning dynamics and equity across class groups, particularly in settings where conventional importance sampling and static coresets underperform or introduce brittleness.

Limitations and Future Directions

A notable limitation is RCAP's need for a few epochs to stabilize per-class allocation fractions, which may hinder rapid adaptation in ultra-short training regimes such as few-shot settings or LLM pretraining phases. The manual selection of the Softmax temperature also restricts full automation; adaptive scheduling or learnable parameterization of this hyperparameter is highlighted as promising future work.

Extensions of RCAP to self-supervised, multi-modal, or reinforcement learning domains, as well as integration with active learning or data condensation techniques, are natural progression areas. Furthermore, reducing approximation error of class-wise allocations in the early training stages could amplify RCAP's effectiveness in scenarios with extremely scarce computational budgets.

Conclusion

RCAP advances the state of the art in robust, efficient dynamic dataset pruning by formalizing and operationalizing class-aware, probabilistic sample selection. It delivers strong, consistent improvements in worst-group accuracy and substantial reductions in training cost across a range of architectures, paradigms, and data distributions. The method's robustness to imbalance, coupled with its efficiency and minimal overhead, positions it as a compelling standard for large-scale, fair, and sustainable deep learning workflows (2606.11761).